mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
langgraph: remove FewShotExamples managed value (#1195)
This commit is contained in:
Vadym Barda
File diff suppressed because one or more lines are too long
@@ -1,105 +0,0 @@
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
AsyncGenerator,
|
||||
AsyncIterator,
|
||||
Callable,
|
||||
Dict,
|
||||
Generator,
|
||||
Generic,
|
||||
Iterator,
|
||||
Optional,
|
||||
Sequence,
|
||||
Union,
|
||||
)
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from typing_extensions import Self
|
||||
|
||||
from langgraph.channels.manager import AsyncChannelsManager, ChannelsManager
|
||||
from langgraph.managed.base import ConfiguredManagedValue, ManagedValue, V
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.pregel.io import read_channels
|
||||
from langgraph.pregel.types import PregelTaskDescription
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langgraph.pregel import Pregel
|
||||
|
||||
# Metadata filter can be a dict (static) or a function (dynamic) that takes a
|
||||
# RunnableConfig and returns a dict. Functions are used for filtering on
|
||||
# metadata values that are only available at runtime.
|
||||
MetadataFilter = Union[Dict[str, Any], Callable[[RunnableConfig], Dict[str, Any]]]
|
||||
|
||||
|
||||
class FewShotExamples(ManagedValue[Sequence[V]], Generic[V]):
|
||||
examples: list[V]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
graph: Pregel,
|
||||
k: int = 5,
|
||||
metadata_filter: Optional[MetadataFilter] = None,
|
||||
) -> None:
|
||||
super().__init__(config, graph)
|
||||
self.k = k
|
||||
self.metadata_filter = metadata_filter or {}
|
||||
|
||||
@classmethod
|
||||
def configure(
|
||||
cls, k: int = 5, metadata_filter: Optional[MetadataFilter] = None
|
||||
) -> ConfiguredManagedValue:
|
||||
return ConfiguredManagedValue(
|
||||
cls,
|
||||
{
|
||||
"k": k,
|
||||
"metadata_filter": metadata_filter,
|
||||
},
|
||||
)
|
||||
|
||||
@property
|
||||
def metadata_filter_dict(self) -> Dict[str, Any]:
|
||||
if isinstance(self.metadata_filter, Callable):
|
||||
return self.metadata_filter(self.config)
|
||||
else:
|
||||
return self.metadata_filter
|
||||
|
||||
def iter(self, score: int = 1) -> Iterator[V]:
|
||||
for example in self.graph.checkpointer.list(
|
||||
None, filter={"score": score, **self.metadata_filter_dict}, limit=self.k
|
||||
):
|
||||
with ChannelsManager(
|
||||
self.graph.channels, example.checkpoint, self.config
|
||||
) as channels:
|
||||
yield read_channels(channels, self.graph.output_channels)
|
||||
|
||||
async def aiter(self, score: int = 1) -> AsyncIterator[V]:
|
||||
async for example in self.graph.checkpointer.alist(
|
||||
None, filter={"score": score, **self.metadata_filter_dict}, limit=self.k
|
||||
):
|
||||
async with AsyncChannelsManager(
|
||||
self.graph.channels, example.checkpoint, self.config
|
||||
) as channels:
|
||||
yield read_channels(channels, self.graph.output_channels)
|
||||
|
||||
@classmethod
|
||||
@contextmanager
|
||||
def enter(
|
||||
cls, config: RunnableConfig, graph: "Pregel", **kwargs: Any
|
||||
) -> Generator[Self, None, None]:
|
||||
with super().enter(config, graph, **kwargs) as value:
|
||||
value.examples = list(value.iter())
|
||||
yield value
|
||||
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
async def aenter(
|
||||
cls, config: RunnableConfig, graph: "Pregel", **kwargs: Any
|
||||
) -> AsyncGenerator[Self, None]:
|
||||
async with super().aenter(config, graph, **kwargs) as value:
|
||||
value.examples = [e async for e in value.aiter()]
|
||||
yield value
|
||||
|
||||
def __call__(self, step: int, task: PregelTaskDescription) -> Sequence[V]:
|
||||
return self.examples
|
||||
@@ -54,7 +54,6 @@ from langgraph.graph import END, Graph
|
||||
from langgraph.graph.graph import START
|
||||
from langgraph.graph.message import MessageGraph, add_messages
|
||||
from langgraph.graph.state import StateGraph
|
||||
from langgraph.managed.few_shot import FewShotExamples
|
||||
from langgraph.prebuilt.chat_agent_executor import (
|
||||
create_function_calling_executor,
|
||||
create_tool_calling_executor,
|
||||
@@ -3394,233 +3393,6 @@ def test_state_graph_w_config(snapshot: SnapshotAssertion) -> None:
|
||||
assert app.config_schema().schema_json() == snapshot
|
||||
|
||||
|
||||
def test_state_graph_few_shot() -> None:
|
||||
from langchain_core.language_models.fake_chat_models import (
|
||||
FakeMessagesListChatModel,
|
||||
)
|
||||
from langchain_core.messages import AIMessage, AnyMessage, ToolMessage
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.tools import tool
|
||||
|
||||
def filter_by_source(config: RunnableConfig) -> Dict[str, Any]:
|
||||
"""This function is a trivial example that demonstrates that passing
|
||||
a Callable to metadata_filter works as expected.
|
||||
"""
|
||||
return {"source": "loop"}
|
||||
|
||||
class BaseState(TypedDict):
|
||||
messages: Annotated[list[AnyMessage], add_messages]
|
||||
|
||||
class AgentState(BaseState):
|
||||
examples: Annotated[
|
||||
Sequence[BaseState],
|
||||
FewShotExamples[BaseState].configure(k=1, metadata_filter=filter_by_source),
|
||||
]
|
||||
|
||||
# Assemble the tools
|
||||
@tool()
|
||||
def search_api(query: str) -> str:
|
||||
"""Searches the API for the query."""
|
||||
return f"result for {query}"
|
||||
|
||||
tools = [search_api]
|
||||
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
(
|
||||
"system",
|
||||
"""You are a nice assistant.
|
||||
Some examples of past conversations:
|
||||
{examples}""",
|
||||
),
|
||||
("placeholder", "{messages}"),
|
||||
]
|
||||
)
|
||||
|
||||
model = FakeMessagesListChatModel(
|
||||
responses=[
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call123",
|
||||
"name": "search_api",
|
||||
"args": {"query": "query"},
|
||||
},
|
||||
],
|
||||
),
|
||||
AIMessage(content="answer"),
|
||||
]
|
||||
)
|
||||
|
||||
def agent(state: AgentState, config: RunnableConfig) -> AgentState:
|
||||
# begin: testing code
|
||||
assert state["examples"] == config["configurable"]["expected_examples"]
|
||||
# end: testing code
|
||||
formatted = prompt.invoke(state)
|
||||
response = model.invoke(formatted)
|
||||
return {"messages": response}
|
||||
|
||||
# Define decision-making logic
|
||||
def should_continue(data: AgentState) -> str:
|
||||
# Logic to decide whether to continue in the loop or exit
|
||||
if not data["messages"][-1].tool_calls:
|
||||
return "exit"
|
||||
else:
|
||||
return "continue"
|
||||
|
||||
# Define a new graph
|
||||
workflow = StateGraph(AgentState)
|
||||
|
||||
workflow.add_node("agent", agent)
|
||||
workflow.add_node("tools", ToolNode(tools))
|
||||
workflow.set_entry_point("agent")
|
||||
workflow.add_conditional_edges(
|
||||
"agent", should_continue, {"continue": "tools", "exit": END}
|
||||
)
|
||||
workflow.add_edge("tools", "agent")
|
||||
|
||||
with SqliteSaver.from_conn_string(":memory:") as saver:
|
||||
app = workflow.compile(checkpointer=saver)
|
||||
|
||||
first_messages = [
|
||||
_AnyIdHumanMessage(content="what is weather in sf"),
|
||||
AIMessage(
|
||||
content="",
|
||||
id=AnyStr(),
|
||||
tool_calls=[
|
||||
{
|
||||
"name": "search_api",
|
||||
"args": {"query": "query"},
|
||||
"id": "tool_call123",
|
||||
"type": "tool_call",
|
||||
}
|
||||
],
|
||||
),
|
||||
ToolMessage(
|
||||
content="result for query",
|
||||
name="search_api",
|
||||
id=AnyStr(),
|
||||
tool_call_id="tool_call123",
|
||||
),
|
||||
_AnyIdAIMessage(content="answer"),
|
||||
]
|
||||
actual = app.invoke(
|
||||
{"messages": "what is weather in sf"},
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"expected_examples": [],
|
||||
},
|
||||
},
|
||||
)
|
||||
expected = {"messages": first_messages}
|
||||
assert actual == expected
|
||||
|
||||
# get first checkpoint
|
||||
chkpnt_tuple_1 = saver.get_tuple({"configurable": {"thread_id": "1"}})
|
||||
config = chkpnt_tuple_1.config
|
||||
checkpoint = chkpnt_tuple_1.checkpoint
|
||||
metadata = chkpnt_tuple_1.metadata
|
||||
|
||||
# not needed in application code, only for testing
|
||||
hiscored = list(saver.list(None, filter={"score": 1}))
|
||||
assert hiscored == []
|
||||
|
||||
# mark as "good"
|
||||
metadata["score"] = 1
|
||||
saver.put(config, checkpoint, metadata)
|
||||
|
||||
# not needed in application code, only for testing
|
||||
hiscored = list(saver.list(None, filter={"score": 1}))
|
||||
assert len(hiscored) == 1
|
||||
assert hiscored[0].checkpoint["channel_values"]["messages"] == first_messages
|
||||
|
||||
second_messages = [
|
||||
_AnyIdHumanMessage(content="what is weather in la"),
|
||||
AIMessage(
|
||||
content="",
|
||||
id=AnyStr(),
|
||||
tool_calls=[
|
||||
{
|
||||
"name": "search_api",
|
||||
"args": {"query": "query"},
|
||||
"id": "tool_call123",
|
||||
"type": "tool_call",
|
||||
}
|
||||
],
|
||||
),
|
||||
ToolMessage(
|
||||
content="result for query",
|
||||
name="search_api",
|
||||
id=AnyStr(),
|
||||
tool_call_id="tool_call123",
|
||||
),
|
||||
_AnyIdAIMessage(content="answer"),
|
||||
]
|
||||
assert app.invoke(
|
||||
{"messages": "what is weather in la"},
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": "2",
|
||||
# below is only for testing purposes, not part of few shot api
|
||||
"expected_examples": [{"messages": first_messages}],
|
||||
}
|
||||
},
|
||||
) == {"messages": second_messages}
|
||||
|
||||
# get first checkpoint
|
||||
chkpnt_tuple_2 = saver.get_tuple({"configurable": {"thread_id": "2"}})
|
||||
config = chkpnt_tuple_2.config
|
||||
checkpoint = chkpnt_tuple_2.checkpoint
|
||||
metadata = chkpnt_tuple_2.metadata
|
||||
|
||||
# not needed in application code, only for testing
|
||||
hiscored = list(saver.list(None, filter={"score": 1}))
|
||||
assert len(hiscored) == 1
|
||||
|
||||
# mark as "good"
|
||||
metadata["score"] = 1
|
||||
saver.put(config, checkpoint, metadata)
|
||||
|
||||
hiscored = list(saver.list(None, filter={"score": 1}))
|
||||
assert len(hiscored) == 2
|
||||
|
||||
assert app.invoke(
|
||||
{"messages": "what is weather in ny"},
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": "3",
|
||||
# below is only for testing purposes, not part of few shot api
|
||||
"expected_examples": [{"messages": second_messages}],
|
||||
}
|
||||
},
|
||||
) == {
|
||||
"messages": [
|
||||
_AnyIdHumanMessage(content="what is weather in ny"),
|
||||
AIMessage(
|
||||
content="",
|
||||
id=AnyStr(),
|
||||
tool_calls=[
|
||||
{
|
||||
"name": "search_api",
|
||||
"args": {"query": "query"},
|
||||
"id": "tool_call123",
|
||||
"type": "tool_call",
|
||||
}
|
||||
],
|
||||
),
|
||||
ToolMessage(
|
||||
content="result for query",
|
||||
name="search_api",
|
||||
id=AnyStr(),
|
||||
tool_call_id="tool_call123",
|
||||
),
|
||||
_AnyIdAIMessage(content="answer"),
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None:
|
||||
class AgentState(TypedDict, total=False):
|
||||
input: str
|
||||
|
||||
@@ -14,7 +14,6 @@ from typing import (
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Sequence,
|
||||
Tuple,
|
||||
TypedDict,
|
||||
Union,
|
||||
@@ -52,7 +51,6 @@ from langgraph.errors import InvalidUpdateError
|
||||
from langgraph.graph import END, Graph, StateGraph
|
||||
from langgraph.graph.graph import START
|
||||
from langgraph.graph.message import MessageGraph, add_messages
|
||||
from langgraph.managed.few_shot import FewShotExamples
|
||||
from langgraph.prebuilt.chat_agent_executor import (
|
||||
create_function_calling_executor,
|
||||
create_tool_calling_executor,
|
||||
@@ -3242,187 +3240,6 @@ async def test_conditional_graph_state() -> None:
|
||||
)
|
||||
|
||||
|
||||
async def test_state_graph_few_shot() -> None:
|
||||
from langchain_core.language_models.fake_chat_models import (
|
||||
FakeMessagesListChatModel,
|
||||
)
|
||||
from langchain_core.messages import (
|
||||
AIMessage,
|
||||
AnyMessage,
|
||||
ToolCall,
|
||||
ToolMessage,
|
||||
)
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.tools import tool
|
||||
|
||||
def filter_by_source(config: RunnableConfig) -> Dict[str, Any]:
|
||||
"""This function is a trivial example that demonstrates that passing
|
||||
a Callable to metadata_filter works as expected.
|
||||
"""
|
||||
return {"source": "loop"}
|
||||
|
||||
class BaseState(TypedDict):
|
||||
messages: Annotated[list[AnyMessage], add_messages]
|
||||
|
||||
class AgentState(BaseState):
|
||||
examples: Annotated[
|
||||
Sequence[BaseState],
|
||||
FewShotExamples[BaseState].configure(k=1, metadata_filter=filter_by_source),
|
||||
]
|
||||
|
||||
# Assemble the tools
|
||||
@tool()
|
||||
def search_api(query: str) -> str:
|
||||
"""Searches the API for the query."""
|
||||
return f"result for {query}"
|
||||
|
||||
tools = [search_api]
|
||||
tools_by_name = {t.name: t for t in tools}
|
||||
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
(
|
||||
"system",
|
||||
"""You are a nice assistant.
|
||||
Some examples of past conversations:
|
||||
{examples}""",
|
||||
),
|
||||
("placeholder", "{messages}"),
|
||||
]
|
||||
)
|
||||
|
||||
model = FakeMessagesListChatModel(
|
||||
responses=[
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "tool_call123",
|
||||
"name": "search_api",
|
||||
"args": {"query": "query"},
|
||||
},
|
||||
],
|
||||
),
|
||||
AIMessage(content="answer"),
|
||||
]
|
||||
)
|
||||
|
||||
async def agent(state: AgentState, config: RunnableConfig) -> AgentState:
|
||||
# begin: testing code
|
||||
assert state["examples"] == config["configurable"]["expected_examples"]
|
||||
# end: testing code
|
||||
formatted = await prompt.ainvoke(state)
|
||||
response = await model.ainvoke(formatted)
|
||||
return {"messages": response}
|
||||
|
||||
# Define decision-making logic
|
||||
def should_continue(data: AgentState) -> str:
|
||||
# Logic to decide whether to continue in the loop or exit
|
||||
if tool_calls := data["messages"][-1].tool_calls:
|
||||
return [Send("tools", tool_call) for tool_call in tool_calls]
|
||||
else:
|
||||
return "exit"
|
||||
|
||||
def tools_node(tool_call: ToolCall, config: RunnableConfig) -> AgentState:
|
||||
output = tools_by_name[tool_call["name"]].invoke(tool_call["args"], config)
|
||||
return {
|
||||
"messages": ToolMessage(
|
||||
content=output, name=tool_call["name"], tool_call_id=tool_call["id"]
|
||||
)
|
||||
}
|
||||
|
||||
# Define a new graph
|
||||
workflow = StateGraph(AgentState)
|
||||
|
||||
workflow.add_node("agent", agent)
|
||||
workflow.add_node("tools", tools_node)
|
||||
workflow.set_entry_point("agent")
|
||||
workflow.add_conditional_edges(
|
||||
"agent", should_continue, {"continue": "tools", "exit": END}
|
||||
)
|
||||
workflow.add_edge("tools", "agent")
|
||||
|
||||
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
|
||||
app = workflow.compile(checkpointer=saver)
|
||||
|
||||
first_messages = [
|
||||
_AnyIdHumanMessage(content="what is weather in sf"),
|
||||
AIMessage(
|
||||
content="",
|
||||
id=AnyStr(),
|
||||
tool_calls=[
|
||||
{
|
||||
"name": "search_api",
|
||||
"args": {"query": "query"},
|
||||
"id": "tool_call123",
|
||||
}
|
||||
],
|
||||
),
|
||||
ToolMessage(
|
||||
content="result for query",
|
||||
name="search_api",
|
||||
id=AnyStr(),
|
||||
tool_call_id="tool_call123",
|
||||
),
|
||||
_AnyIdAIMessage(content="answer"),
|
||||
]
|
||||
assert await app.ainvoke(
|
||||
{"messages": "what is weather in sf"},
|
||||
{"configurable": {"thread_id": "1", "expected_examples": []}},
|
||||
) == {"messages": first_messages}
|
||||
|
||||
# get first checkpoint
|
||||
chkpnt_tuple_1 = await saver.aget_tuple({"configurable": {"thread_id": "1"}})
|
||||
config = chkpnt_tuple_1.config
|
||||
checkpoint = chkpnt_tuple_1.checkpoint
|
||||
metadata = chkpnt_tuple_1.metadata
|
||||
|
||||
# not needed in application code, only for testing
|
||||
assert [c async for c in saver.alist(None, filter={"score": 1})] == []
|
||||
|
||||
# mark as "good"
|
||||
metadata["score"] = 1
|
||||
await saver.aput(config, checkpoint, metadata)
|
||||
|
||||
# not needed in application code, only for testing
|
||||
hiscored = [c async for c in saver.alist(None, filter={"score": 1})]
|
||||
assert len(hiscored) == 1
|
||||
assert hiscored[0].checkpoint["channel_values"]["messages"] == first_messages
|
||||
|
||||
assert await app.ainvoke(
|
||||
{"messages": "what is weather in la"},
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": "2",
|
||||
# below is only for testing purposes, not part of few shot api
|
||||
"expected_examples": [{"messages": first_messages}],
|
||||
}
|
||||
},
|
||||
) == {
|
||||
"messages": [
|
||||
_AnyIdHumanMessage(content="what is weather in la"),
|
||||
AIMessage(
|
||||
content="",
|
||||
id=AnyStr(),
|
||||
tool_calls=[
|
||||
{
|
||||
"name": "search_api",
|
||||
"args": {"query": "query"},
|
||||
"id": "tool_call123",
|
||||
}
|
||||
],
|
||||
),
|
||||
ToolMessage(
|
||||
content="result for query",
|
||||
name="search_api",
|
||||
id=AnyStr(),
|
||||
tool_call_id="tool_call123",
|
||||
),
|
||||
_AnyIdAIMessage(content="answer"),
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
async def test_conditional_entrypoint_graph() -> None:
|
||||
async def left(data: str) -> str:
|
||||
return data + "->left"
|
||||
|
||||
Reference in New Issue
Block a user