Bygge fleragentpipelines med LangGraph

# pip install langgraph
from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated
import operator

# Define your shared state schema
class AgentState(TypedDict):
    messages: Annotated[list, operator.add]  # append-only list
    research: str
    draft: str
    status: str

from typing import TypedDict

class ResearchState(TypedDict):
    query: str           # input from user
    research_notes: str  # filled by researcher node
    written_draft: str   # filled by writer node
    review_feedback: str # filled by reviewer node
    final_output: str    # filled by synthesizer node
    iteration_count: int # tracks how many revision loops occurred

# Each node returns a PARTIAL update - only the keys it modifies
def researcher_node(state: ResearchState) -> dict:
    notes = do_research(state['query'])
    return {'research_notes': notes}  # only update this key

from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage

llm = ChatOpenAI(model='gpt-4o')

def researcher_node(state: ResearchState) -> dict:
    messages = [
        SystemMessage(content='You are a research specialist. Find accurate information.'),
        HumanMessage(content=f'Research this topic: {state["query"]}')
    ]
    response = llm.invoke(messages)
    return {'research_notes': response.content}

def writer_node(state: ResearchState) -> dict:
    messages = [
        SystemMessage(content='You are a technical writer. Write clear, engaging prose.'),
        HumanMessage(content=f'Write a draft using these notes:\n{state["research_notes"]}')
    ]
    response = llm.invoke(messages)
    return {'written_draft': response.content}

from langgraph.graph import StateGraph, END

# Build the graph
workflow = StateGraph(ResearchState)

# Add nodes
workflow.add_node('researcher', researcher_node)
workflow.add_node('writer', writer_node)
workflow.add_node('reviewer', reviewer_node)

# Add edges (sequential pipeline)
workflow.set_entry_point('researcher')
workflow.add_edge('researcher', 'writer')
workflow.add_edge('writer', 'reviewer')
workflow.add_edge('reviewer', END)

# Compile into an executable
app = workflow.compile()

# Run it
result = app.invoke({'query': 'What is RAG?', 'iteration_count': 0})
print(result['written_draft'])

def should_revise(state: ResearchState) -> str:
    '''Router function: returns the name of the next node.'''
    if state['iteration_count'] >= 3:
        return 'finalize'  # Too many revisions - accept as is
    if 'insufficient' in state.get('review_feedback', '').lower():
        return 'researcher'  # Need more research
    if 'rewrite' in state.get('review_feedback', '').lower():
        return 'writer'  # Needs rewriting
    return 'finalize'  # Looks good

# Add conditional edge from reviewer
workflow.add_conditional_edges(
    'reviewer',                    # from node
    should_revise,                 # routing function
    {
        'researcher': 'researcher', # route name -> node name
        'writer': 'writer',
        'finalize': 'finalizer'
    }
)

def researcher_node(state: ResearchState) -> dict:
    notes = do_research(state['query'])
    return {
        'research_notes': notes,
        'iteration_count': state['iteration_count'] + 1  # always increment
    }

def should_continue_research(state: ResearchState) -> str:
    # Terminate loop after 3 iterations regardless of quality
    if state['iteration_count'] >= 3:
        return END
    # Continue if research is incomplete
    if len(state.get('research_notes', '')) < 500:
        return 'researcher'  # loop back
    return 'writer'  # proceed to next stage

# Fan-out: one node triggers multiple parallel ones
workflow.add_edge('planner', 'researcher_a')
workflow.add_edge('planner', 'researcher_b')
workflow.add_edge('planner', 'researcher_c')

# Fan-in: aggregator waits for all three
workflow.add_edge('researcher_a', 'aggregator')
workflow.add_edge('researcher_b', 'aggregator')
workflow.add_edge('researcher_c', 'aggregator')
workflow.add_edge('aggregator', 'writer')

# Aggregator merges parallel results
def aggregator_node(state: ResearchState) -> dict:
    combined = state.get('notes_a', '') + '\n' + state.get('notes_b', '') + '\n' + state.get('notes_c', '')
    return {'research_notes': combined}

from langgraph.checkpoint.sqlite import SqliteSaver

# Use SQLite for persistent checkpoints
checkpointer = SqliteSaver.from_conn_string(':memory:')
app = workflow.compile(checkpointer=checkpointer, interrupt_before=['human_review'])

# First run - pauses at human_review node
thread = {'configurable': {'thread_id': 'my-workflow-1'}}
result = app.invoke({'query': 'Analyze competitors'}, config=thread)
# result.next == 'human_review' -- waiting for human input

# Human provides feedback and resumes
app.update_state(thread, {'review_feedback': 'Good research, proceed with writing'})
final = app.invoke(None, config=thread)  # resume from checkpoint

# Stream intermediate results as each node completes
for event in app.stream({'query': 'What is RAG?'}):
    for node_name, node_output in event.items():
        print(f'Node completed: {node_name}')
        if 'research_notes' in node_output:
            print('Research done:', node_output['research_notes'][:100])
        if 'written_draft' in node_output:
            print('Draft done:', node_output['written_draft'][:100])

# Visualize the workflow graph
graph_image = app.get_graph().draw_mermaid_png()
with open('workflow.png', 'wb') as f:
    f.write(graph_image)

# Or print Mermaid syntax
print(app.get_graph().draw_mermaid())
# Outputs:
# graph TD
#    __start__ --> researcher
#    researcher --> writer
#    writer --> reviewer
#    reviewer -->|Good| finalize
#    reviewer -->|Needs work| researcher