[Ch 3] Getting Started with LangChain & LangGraph

pip install langchain-core langchain-openai langgraph

# .env.example
OPENAI_API_KEY=your-api-key-here

# messages_demo.py
from langchain_core.messages import (
    SystemMessage,
    HumanMessage,
    AIMessage,
    ToolMessage,
)

# System prompt
system = SystemMessage(content="You are a helpful assistant.")

# User turn
user = HumanMessage(content="What is the capital of France?")

# Assistant reply (no tool call)
assistant = AIMessage(content="The capital of France is Paris.")

# When a tool is called, the AIMessage contains tool_calls
assistant_with_tool = AIMessage(
    content="",
    tool_calls=[
        {
            "id": "call_abc123",
            "name": "search_web",
            "args": {"query": "capital of France"},
        }
    ],
)

# The tool result is always a ToolMessage
tool_result = ToolMessage(
    tool_call_id="call_abc123",
    content="Paris is the capital and most populous city of France.",
)

# Build a conversation list (this is what gets sent to the LLM)
conversation = [system, user, assistant_with_tool, tool_result]
print(f"Conversation has {len(conversation)} messages")

# llm_setup.py
import os
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(
    model="gpt-4o-mini",        # cheaper, fast — good for most agent tasks
    temperature=0.1,             # low randomness for consistent tool-calling
    api_key=os.environ.get("OPENAI_API_KEY"),
)

# Simple invoke
response = llm.invoke([HumanMessage(content="Hello!")])
print(response.content)          # "Hello! How can I help you today?"
print(type(response))            # <class 'langchain_core.messages.ai.AIMessage'>

# tools.py
from langchain_core.tools import tool
from pydantic import BaseModel, Field

# Simple tool — docstring becomes the description
@tool
def get_current_time() -> str:
    """Get the current UTC time in ISO 8601 format."""
    from datetime import datetime, timezone
    return datetime.now(timezone.utc).isoformat()

# Tool with Pydantic input schema — best practice for complex inputs
class SearchInput(BaseModel):
    query: str = Field(description="The search query string")
    max_results: int = Field(default=5, description="Maximum number of results (1-20)")

@tool("search_documents", args_schema=SearchInput)
def search_documents(query: str, max_results: int = 5) -> str:
    """Search the document knowledge base for passages relevant to the query."""
    # In a real system, this calls your vector store
    return f"[Simulated] Found {max_results} results for: '{query}'"

# Inspect what the framework generated
print(get_current_time.name)          # "get_current_time"
print(get_current_time.description)   # "Get the current UTC time in ISO 8601 format."
print(search_documents.args_schema.schema())  # Full JSON schema

# Bind tools to the LLM — this sends schemas in every API call
tools = [get_current_time, search_documents]
llm_with_tools = llm.bind_tools(tools)

# runnable_config_demo.py
from langchain_core.runnables import RunnableConfig

# Pass config through invoke
config: RunnableConfig = {
    "configurable": {
        "user_id": "user-42",
        "session_id": "sess-abc",
        "thread_id": "thread-001",   # LangGraph uses this for checkpointing
    },
    "callbacks": [],  # attach tracing callbacks here
}

response = llm_with_tools.invoke(
    [HumanMessage(content="What time is it?")],
    config=config,
)

@tool
def my_tool(query: str, config: RunnableConfig) -> str:  # config is injected automatically
    """A tool that knows who is calling it."""
    user_id = config.get("configurable", {}).get("user_id", "unknown")
    return f"Hello user {user_id}, you searched for: {query}"

# state.py
from typing import Annotated, TypedDict
from langchain_core.messages import AnyMessage
from langgraph.graph.message import add_messages

class AgentState(TypedDict):
    # add_messages reducer: new messages are appended, not replaced
    messages: Annotated[list[AnyMessage], add_messages]
    # You can add other fields for custom state
    step_count: int
    error: str | None

# nodes.py
from langchain_core.messages import SystemMessage

SYSTEM_PROMPT = "You are a helpful assistant. Use tools when needed."

def agent_node(state: AgentState) -> dict:
    """The brain: calls the LLM with the current message history."""
    messages = [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}  # add_messages will append this

def increment_step(state: AgentState) -> dict:
    """A simple state-tracking node."""
    return {"step_count": state.get("step_count", 0) + 1}

# routing.py
from langgraph.prebuilt import ToolNode

def should_continue(state: AgentState) -> str:
    """
    Return "tools" if the last AI message has tool calls,
    otherwise return "end" to finish.
    """
    last_message = state["messages"][-1]
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        return "tools"
    return "end"

# graph.py
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode

def build_graph(tools: list) -> StateGraph:
    tool_node = ToolNode(tools)   # handles tool execution + ToolMessage creation

    graph = StateGraph(AgentState)

    # Add nodes
    graph.add_node("agent", agent_node)
    graph.add_node("tools", tool_node)

    # Entry point
    graph.set_entry_point("agent")

    # Conditional edge: after agent runs, decide what happens next
    graph.add_conditional_edges(
        "agent",
        should_continue,
        {
            "tools": "tools",  # if tool call → run tools
            "end": END,        # if no tool call → stop
        },
    )

    # After tools run, always go back to agent
    graph.add_edge("tools", "agent")

    return graph.compile()

# hello_agent.py
import os
from typing import Annotated, TypedDict

from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode

# ── State ──────────────────────────────────────────────────────────────────
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]

# ── LLM ────────────────────────────────────────────────────────────────────
llm = ChatOpenAI(
    model="gpt-4o-mini",
    temperature=0.1,
    api_key=os.environ.get("OPENAI_API_KEY"),
)

# ── Tools ───────────────────────────────────────────────────────────────────
@tool
def add_numbers(a: float, b: float) -> float:
    """Add two numbers together and return the result."""
    return a + b

@tool
def get_word_count(text: str) -> int:
    """Count the number of words in the given text."""
    return len(text.split())

TOOLS = [add_numbers, get_word_count]
llm_with_tools = llm.bind_tools(TOOLS)

# ── Nodes ────────────────────────────────────────────────────────────────────
SYSTEM_PROMPT = "You are a helpful assistant. Use tools when you need to compute things."

def agent_node(state: AgentState) -> dict:
    messages = [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}

def should_continue(state: AgentState) -> str:
    last = state["messages"][-1]
    if hasattr(last, "tool_calls") and last.tool_calls:
        return "tools"
    return "end"

# ── Graph ─────────────────────────────────────────────────────────────────────
graph = StateGraph(AgentState)
graph.add_node("agent", agent_node)
graph.add_node("tools", ToolNode(TOOLS))
graph.set_entry_point("agent")
graph.add_conditional_edges("agent", should_continue, {"tools": "tools", "end": END})
graph.add_edge("tools", "agent")
app = graph.compile()

# ── Run ───────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    initial_state = {
        "messages": [HumanMessage(content="What is 1337 + 42? Also, how many words are in: 'the quick brown fox'?")]
    }

    # Invoke (blocking, returns final state)
    final_state = app.invoke(initial_state)
    print("\n=== Final Answer ===")
    print(final_state["messages"][-1].content)

    # Optional: print all messages to see the full trajectory
    print("\n=== Full Trajectory ===")
    for msg in final_state["messages"]:
        role = msg.__class__.__name__
        content = msg.content if msg.content else f"[tool calls: {msg.tool_calls}]"
        print(f"[{role}] {str(content)[:120]}")

=== Final Answer ===
I calculated both for you:
- 1337 + 42 = **1379**
- The sentence "the quick brown fox" contains **4 words**

=== Full Trajectory ===
[HumanMessage] What is 1337 + 42? Also, how many words are in: 'the quick brown fox'?
[AIMessage] [tool calls: [{'name': 'add_numbers', 'args': {'a': 1337, 'b': 42}...}]]
[ToolMessage] 1379.0
[ToolMessage] 4
[AIMessage] I calculated both for you: ...

# .env.example
OPENAI_API_KEY=your-api-key-here