jstoppa/langgraph_basic_example_api

1from fastapi import FastAPI2from langgraph.graph import StateGraph3from typing import TypedDict, Annotated, List4from langgraph.graph.message import add_messages5from pydantic import BaseModel6 7# Initialize FastAPI app8app = FastAPI(title="LangGraph Agent API")9 10class State(TypedDict):11    messages: Annotated[list[str], add_messages]12    current_step: str13 14class AgentInput(BaseModel):15    messages: List[str]16 17def collect_info(state: State) -> dict:18    print("\n--> In collect_info")19    print(f"Messages before: {state['messages']}")20    21    messages = state["messages"] + ["Information collected"]22    print(f"Messages after: {messages}")23    24    return {25        "messages": messages,26        "current_step": "process"27    }28 29def process_info(state: State) -> dict:30    print("\n--> In process_info")31    print(f"Messages before: {state['messages']}")32    33    messages = state["messages"] + ["Information processed"]34    print(f"Messages after: {messages}")35    36    return {37        "messages": messages,38        "current_step": "end"39    }40 41# Create and setup graph42workflow = StateGraph(State)43 44# Add nodes45workflow.add_node("collect", collect_info)46workflow.add_node("process", process_info)47 48# Add edges49workflow.add_edge("collect", "process")50 51# Set entry and finish points52workflow.set_entry_point("collect")53workflow.set_finish_point("process")54 55# Compile the workflow56agent = workflow.compile()57 58 59@app.post("/run-agent")60async def run_agent(input_data: AgentInput):61    """62    Run the agent with the provided input messages.63    """64    initial_state = State(messages=input_data.messages, current_step="collect")65    final_state = agent.invoke(initial_state)66    return {"messages": final_state["messages"]}67 68@app.get("/")69async def root():70    """71    Root endpoint that returns basic API information.72    """73    return {"message": "LangGraph Agent API is running", "endpoints": ["Navigate to https://jstoppa-langgraph-basic-example-api.hf.space/docs#/default/run_agent_run_agent_post to run the example"]}74 75if __name__ == "__main__":76    import uvicorn77    uvicorn.run(app, host="0.0.0.0", port=7860)