mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-11 20:27:54 +02:00
Add 2 more benchmark graphs (#1707)
This commit is contained in:
Nuno Campos
@@ -1,16 +1,30 @@
|
||||
import random
|
||||
from typing import Optional
|
||||
from uuid import uuid4
|
||||
|
||||
from langchain_core.messages import HumanMessage
|
||||
from pyperf._runner import Runner
|
||||
from uvloop import new_event_loop
|
||||
|
||||
from bench.fanout_to_subgraph import fanout_to_subgraph
|
||||
from bench.react_agent import react_agent
|
||||
from bench.wide_state import wide_state
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
from langgraph.pregel import Pregel
|
||||
|
||||
|
||||
async def run(graph: Pregel, input: dict, config: Optional[dict]):
|
||||
len([c async for c in graph.astream(input, config=config)])
|
||||
async def run(graph: Pregel, input: dict):
|
||||
len(
|
||||
[
|
||||
c
|
||||
async for c in graph.astream(
|
||||
input,
|
||||
{
|
||||
"configurable": {"thread_id": str(uuid4())},
|
||||
"recursion_limit": 1000000000,
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
benchmarks = (
|
||||
@@ -22,7 +36,6 @@ benchmarks = (
|
||||
random.choices("abcdefghijklmnopqrstuvwxyz", k=1000) for _ in range(10)
|
||||
]
|
||||
},
|
||||
None,
|
||||
),
|
||||
(
|
||||
"fanout_to_subgraph_10x_checkpoint",
|
||||
@@ -32,7 +45,6 @@ benchmarks = (
|
||||
random.choices("abcdefghijklmnopqrstuvwxyz", k=1000) for _ in range(10)
|
||||
]
|
||||
},
|
||||
{"configurable": {"thread_id": "1"}},
|
||||
),
|
||||
(
|
||||
"fanout_to_subgraph_100x",
|
||||
@@ -42,7 +54,6 @@ benchmarks = (
|
||||
random.choices("abcdefghijklmnopqrstuvwxyz", k=1000) for _ in range(100)
|
||||
]
|
||||
},
|
||||
None,
|
||||
),
|
||||
(
|
||||
"fanout_to_subgraph_100x_checkpoint",
|
||||
@@ -52,12 +63,121 @@ benchmarks = (
|
||||
random.choices("abcdefghijklmnopqrstuvwxyz", k=1000) for _ in range(100)
|
||||
]
|
||||
},
|
||||
{"configurable": {"thread_id": "1"}},
|
||||
),
|
||||
(
|
||||
"react_agent_10x",
|
||||
react_agent(10, checkpointer=None),
|
||||
{"messages": [HumanMessage("hi?")]},
|
||||
),
|
||||
(
|
||||
"react_agent_10x_checkpoint",
|
||||
react_agent(10, checkpointer=MemorySaver()),
|
||||
{"messages": [HumanMessage("hi?")]},
|
||||
),
|
||||
(
|
||||
"react_agent_100x",
|
||||
react_agent(100, checkpointer=None),
|
||||
{"messages": [HumanMessage("hi?")]},
|
||||
),
|
||||
(
|
||||
"react_agent_100x_checkpoint",
|
||||
react_agent(100, checkpointer=MemorySaver()),
|
||||
{"messages": [HumanMessage("hi?")]},
|
||||
),
|
||||
(
|
||||
"wide_state_25x300",
|
||||
wide_state(300).compile(checkpointer=None),
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
str(i) * 10: {
|
||||
str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
|
||||
for j in range(5)
|
||||
}
|
||||
for i in range(5)
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
"wide_state_25x300_checkpoint",
|
||||
wide_state(300).compile(checkpointer=MemorySaver()),
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
str(i) * 10: {
|
||||
str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
|
||||
for j in range(5)
|
||||
}
|
||||
for i in range(5)
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
"wide_state_15x600",
|
||||
wide_state(600).compile(checkpointer=None),
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
str(i) * 10: {
|
||||
str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
|
||||
for j in range(5)
|
||||
}
|
||||
for i in range(3)
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
"wide_state_15x600_checkpoint",
|
||||
wide_state(600).compile(checkpointer=MemorySaver()),
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
str(i) * 10: {
|
||||
str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
|
||||
for j in range(5)
|
||||
}
|
||||
for i in range(3)
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
"wide_state_9x1200",
|
||||
wide_state(1200).compile(checkpointer=None),
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
str(i) * 10: {
|
||||
str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
|
||||
for j in range(3)
|
||||
}
|
||||
for i in range(3)
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
"wide_state_9x1200_checkpoint",
|
||||
wide_state(1200).compile(checkpointer=MemorySaver()),
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
str(i) * 10: {
|
||||
str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
|
||||
for j in range(3)
|
||||
}
|
||||
for i in range(3)
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
r = Runner()
|
||||
|
||||
for name, graph, input, config in benchmarks:
|
||||
r.bench_async_func(name, run, graph, input, config, loop_factory=new_event_loop)
|
||||
for name, graph, input in benchmarks:
|
||||
r.bench_async_func(name, run, graph, input, loop_factory=new_event_loop)
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
import asyncio
|
||||
import operator
|
||||
from typing import Annotated, TypedDict
|
||||
|
||||
@@ -55,6 +54,7 @@ def fanout_to_subgraph() -> StateGraph:
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import asyncio
|
||||
import random
|
||||
|
||||
import uvloop
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
from typing import Any, Optional
|
||||
from uuid import uuid4
|
||||
|
||||
from langchain_core.callbacks import CallbackManagerForLLMRun
|
||||
from langchain_core.language_models.fake_chat_models import (
|
||||
FakeMessagesListChatModel,
|
||||
)
|
||||
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
|
||||
from langchain_core.outputs import ChatGeneration, ChatResult
|
||||
from langchain_core.tools import StructuredTool
|
||||
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from langgraph.prebuilt.chat_agent_executor import create_react_agent
|
||||
from langgraph.pregel import Pregel
|
||||
|
||||
|
||||
def react_agent(n_tools: int, checkpointer: BaseCheckpointSaver) -> Pregel:
|
||||
class FakeFuntionChatModel(FakeMessagesListChatModel):
|
||||
def bind_tools(self, functions: list):
|
||||
return self
|
||||
|
||||
def _generate(
|
||||
self,
|
||||
messages: list[BaseMessage],
|
||||
stop: Optional[list[str]] = None,
|
||||
run_manager: Optional[CallbackManagerForLLMRun] = None,
|
||||
**kwargs: Any,
|
||||
) -> ChatResult:
|
||||
response = self.responses[self.i].copy()
|
||||
if self.i < len(self.responses) - 1:
|
||||
self.i += 1
|
||||
else:
|
||||
self.i = 0
|
||||
generation = ChatGeneration(message=response)
|
||||
return ChatResult(generations=[generation])
|
||||
|
||||
tool = StructuredTool.from_function(
|
||||
lambda query: f"result for query: {query}" * 10,
|
||||
name=str(uuid4()),
|
||||
description="",
|
||||
)
|
||||
|
||||
model = FakeFuntionChatModel(
|
||||
responses=[
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": str(uuid4()),
|
||||
"name": tool.name,
|
||||
"args": {"query": str(uuid4()) * 100},
|
||||
}
|
||||
],
|
||||
id=str(uuid4()),
|
||||
)
|
||||
for _ in range(n_tools)
|
||||
]
|
||||
+ [
|
||||
AIMessage(content="answer" * 100, id=str(uuid4())),
|
||||
]
|
||||
)
|
||||
|
||||
return create_react_agent(model, [tool], checkpointer=checkpointer)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import asyncio
|
||||
|
||||
import uvloop
|
||||
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
|
||||
graph = react_agent(100, checkpointer=MemorySaver())
|
||||
input = {"messages": [HumanMessage("hi?")]}
|
||||
config = {"configurable": {"thread_id": "1"}, "recursion_limit": 20000000000}
|
||||
|
||||
async def run():
|
||||
len([c async for c in graph.astream(input, config=config)])
|
||||
|
||||
uvloop.install()
|
||||
asyncio.run(run())
|
||||
@@ -0,0 +1,140 @@
|
||||
import operator
|
||||
from dataclasses import dataclass, field
|
||||
from functools import partial
|
||||
from typing import Annotated, Optional, Sequence
|
||||
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph.state import StateGraph
|
||||
|
||||
|
||||
def wide_state(n: int) -> StateGraph:
|
||||
@dataclass(kw_only=True)
|
||||
class State:
|
||||
messages: Annotated[list, operator.add] = field(default_factory=list)
|
||||
trigger_events: Annotated[list, operator.add] = field(default_factory=list)
|
||||
"""The external events that are converted by the graph."""
|
||||
primary_issue_medium: Annotated[str, lambda x, y: y or x] = field(
|
||||
default="email"
|
||||
)
|
||||
autoresponse: Annotated[Optional[dict], lambda _, y: y] = field(
|
||||
default=None
|
||||
) # Always overwrite
|
||||
issue: Annotated[dict | None, lambda x, y: y if y else x] = field(default=None)
|
||||
relevant_rules: Optional[list[dict]] = field(default=None)
|
||||
"""SOPs fetched from the rulebook that are relevant to the current conversation."""
|
||||
memory_docs: Optional[list[dict]] = field(default=None)
|
||||
"""Memory docs fetched from the memory service that are relevant to the current conversation."""
|
||||
categorizations: Annotated[list[dict], operator.add] = field(
|
||||
default_factory=list
|
||||
)
|
||||
"""The issue categorizations auto-generated by the AI."""
|
||||
responses: Annotated[list[dict], operator.add] = field(default_factory=list)
|
||||
"""The draft responses recommended by the AI."""
|
||||
|
||||
user_info: Annotated[Optional[dict], lambda x, y: y if y is not None else x] = (
|
||||
field(default=None)
|
||||
)
|
||||
"""The current user state (by email)."""
|
||||
crm_info: Annotated[Optional[dict], lambda x, y: y if y is not None else x] = (
|
||||
field(default=None)
|
||||
)
|
||||
"""The CRM information for organization the current user is from."""
|
||||
email_thread_id: Annotated[
|
||||
Optional[str], lambda x, y: y if y is not None else x
|
||||
] = field(default=None)
|
||||
"""The current email thread ID."""
|
||||
slack_participants: Annotated[dict, operator.or_] = field(default_factory=dict)
|
||||
"""The growing list of current slack participants."""
|
||||
bot_id: Optional[str] = field(default=None)
|
||||
"""The ID of the bot user in the slack channel."""
|
||||
notified_assignees: Annotated[dict, operator.or_] = field(default_factory=dict)
|
||||
|
||||
def read_write(read: str, write: Sequence[str], input: State) -> dict:
|
||||
val = getattr(input, read)
|
||||
val_single = val[-1] if isinstance(val, list) else val
|
||||
val_list = val if isinstance(val, list) else [val]
|
||||
return {
|
||||
k: val_list if isinstance(getattr(input, k), list) else val_single
|
||||
for k in write
|
||||
}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_edge(START, "one")
|
||||
builder.add_node(
|
||||
"one",
|
||||
partial(read_write, "messages", ["trigger_events", "primary_issue_medium"]),
|
||||
)
|
||||
builder.add_edge("one", "two")
|
||||
builder.add_node(
|
||||
"two",
|
||||
partial(read_write, "trigger_events", ["autoresponse", "issue"]),
|
||||
)
|
||||
builder.add_edge("two", "three")
|
||||
builder.add_edge("two", "four")
|
||||
builder.add_node(
|
||||
"three",
|
||||
partial(read_write, "autoresponse", ["relevant_rules"]),
|
||||
)
|
||||
builder.add_node(
|
||||
"four",
|
||||
partial(
|
||||
read_write,
|
||||
"trigger_events",
|
||||
["categorizations", "responses", "memory_docs"],
|
||||
),
|
||||
)
|
||||
builder.add_node(
|
||||
"five",
|
||||
partial(
|
||||
read_write,
|
||||
"categorizations",
|
||||
[
|
||||
"user_info",
|
||||
"crm_info",
|
||||
"email_thread_id",
|
||||
"slack_participants",
|
||||
"bot_id",
|
||||
"notified_assignees",
|
||||
],
|
||||
),
|
||||
)
|
||||
builder.add_edge(["three", "four"], "five")
|
||||
builder.add_edge("five", "six")
|
||||
builder.add_node(
|
||||
"six",
|
||||
partial(read_write, "responses", ["messages"]),
|
||||
)
|
||||
builder.add_conditional_edges(
|
||||
"six", lambda state: END if len(state.messages) > n else "one"
|
||||
)
|
||||
|
||||
return builder
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import asyncio
|
||||
|
||||
import uvloop
|
||||
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
|
||||
graph = wide_state(1000).compile(checkpointer=MemorySaver())
|
||||
input = {
|
||||
"messages": [
|
||||
{
|
||||
str(i) * 10: {
|
||||
str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
|
||||
for j in range(5)
|
||||
}
|
||||
for i in range(5)
|
||||
}
|
||||
]
|
||||
}
|
||||
config = {"configurable": {"thread_id": "1"}, "recursion_limit": 20000000000}
|
||||
|
||||
async def run():
|
||||
async for c in graph.astream(input, config=config):
|
||||
print(c.keys())
|
||||
|
||||
uvloop.install()
|
||||
asyncio.run(run())
|
||||
Reference in New Issue
Block a user