flyteplugins-agents-langgraph 2.5.17__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- flyteplugins/agents/langgraph/__init__.py +24 -0
- flyteplugins/agents/langgraph/_memory.py +59 -0
- flyteplugins/agents/langgraph/_nodes.py +165 -0
- flyteplugins/agents/langgraph/_run.py +196 -0
- flyteplugins/agents/langgraph/_tools.py +170 -0
- flyteplugins_agents_langgraph-2.5.17.dist-info/METADATA +9 -0
- flyteplugins_agents_langgraph-2.5.17.dist-info/RECORD +9 -0
- flyteplugins_agents_langgraph-2.5.17.dist-info/WHEEL +5 -0
- flyteplugins_agents_langgraph-2.5.17.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
"""LangGraph adapter for Flyte.
|
|
2
|
+
|
|
3
|
+
Bring your own LangGraph ``StateGraph`` and run it durably on Flyte. You build the
|
|
4
|
+
graph; the adapter provides the durable, observable building blocks:
|
|
5
|
+
|
|
6
|
+
- :func:`tool` — turn a Flyte ``@env.task`` into a LangChain ``StructuredTool``
|
|
7
|
+
(a first-class LangGraph tool) that executes as a durable child action (own
|
|
8
|
+
container/GPU, retries, caching).
|
|
9
|
+
- :func:`ai_node` — the model-calling node: binds the tools to your chat model
|
|
10
|
+
and records each model turn durably (replayed on retry).
|
|
11
|
+
- :func:`tool_node` — the tool-executing node: runs the model's tool calls as
|
|
12
|
+
durable Flyte child actions.
|
|
13
|
+
- :func:`run_agent` — drive a compiled graph (or build a default one from tools)
|
|
14
|
+
inside your task and return the final answer.
|
|
15
|
+
|
|
16
|
+
Each tool call runs as a durable Flyte child action, and the run timeline is
|
|
17
|
+
rendered into the Flyte task report.
|
|
18
|
+
"""
|
|
19
|
+
|
|
20
|
+
from ._nodes import ai_node, tool_node
|
|
21
|
+
from ._run import run_agent, run_agent_sync
|
|
22
|
+
from ._tools import tool
|
|
23
|
+
|
|
24
|
+
__all__ = ["ai_node", "run_agent", "run_agent_sync", "tool", "tool_node"]
|
|
@@ -0,0 +1,59 @@
|
|
|
1
|
+
"""Cross-run LangGraph memory — a thin bridge over Flyte's keyed ``MemoryStore``.
|
|
2
|
+
|
|
3
|
+
LangGraph keeps conversation state in-memory. This module persists the message
|
|
4
|
+
transcript to a durable, keyed :class:`~flyte.ai.agents.memory.MemoryStore` (an
|
|
5
|
+
object-store slot addressed by ``memory_key``) so a later run with the same key
|
|
6
|
+
continues the conversation — across workers and restarts.
|
|
7
|
+
|
|
8
|
+
The transcript is stored (via ``read_json`` / ``write_json``) as the serialized
|
|
9
|
+
LangChain message list, so it round-trips faithfully through
|
|
10
|
+
``messages_from_dict`` / ``messages_to_dict``.
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from __future__ import annotations
|
|
14
|
+
|
|
15
|
+
import typing
|
|
16
|
+
|
|
17
|
+
from flyte._logging import logger
|
|
18
|
+
from flyteplugins.agents.core import resolve_memory as _resolve_memory
|
|
19
|
+
|
|
20
|
+
# Path-addressed slot holding the serialized message transcript inside the MemoryStore.
|
|
21
|
+
_MEMORY_HISTORY_PATH = "langgraph/history.json"
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
async def resolve_memory(memory_key: str | None) -> typing.Any | None:
|
|
25
|
+
"""Resolve a keyed MemoryStore for LangGraph cross-run memory, or ``None``.
|
|
26
|
+
|
|
27
|
+
Best-effort: returns ``None`` when ``memory_key`` is falsy or no durable
|
|
28
|
+
store can be resolved, so memory never breaks a run.
|
|
29
|
+
"""
|
|
30
|
+
if not memory_key:
|
|
31
|
+
return None
|
|
32
|
+
return await _resolve_memory(memory_key)
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
async def load_messages(store: typing.Any) -> list[typing.Any]:
|
|
36
|
+
"""Load the prior conversation as LangChain messages (empty list if none)."""
|
|
37
|
+
if store is None:
|
|
38
|
+
return []
|
|
39
|
+
try:
|
|
40
|
+
from langchain_core.messages import messages_from_dict
|
|
41
|
+
|
|
42
|
+
raw = await store.read_json.aio(_MEMORY_HISTORY_PATH, [])
|
|
43
|
+
return messages_from_dict(raw) if raw else []
|
|
44
|
+
except Exception: # pragma: no cover - memory is best-effort, never fatal
|
|
45
|
+
logger.warning("Could not load LangGraph memory; continuing without prior history.")
|
|
46
|
+
return []
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
async def save_messages(store: typing.Any, messages: typing.Sequence[typing.Any]) -> None:
|
|
50
|
+
"""Persist the full conversation transcript back to the keyed store."""
|
|
51
|
+
if store is None or not messages:
|
|
52
|
+
return
|
|
53
|
+
try:
|
|
54
|
+
from langchain_core.messages import messages_to_dict
|
|
55
|
+
|
|
56
|
+
await store.write_json.aio(_MEMORY_HISTORY_PATH, messages_to_dict(list(messages)))
|
|
57
|
+
await store.save.aio()
|
|
58
|
+
except Exception: # pragma: no cover - memory is best-effort, never fatal
|
|
59
|
+
logger.warning("Could not persist LangGraph memory; continuing.")
|
|
@@ -0,0 +1,165 @@
|
|
|
1
|
+
"""LangGraph node factories that make Flyte the durable runtime under a graph.
|
|
2
|
+
|
|
3
|
+
The intended devex: you build the ``StateGraph`` yourself, and these two factories
|
|
4
|
+
provide the nodes that Flyte makes durable and observable.
|
|
5
|
+
|
|
6
|
+
- :func:`ai_node` — the model-calling node. It binds your ``@tool``-wrapped tasks
|
|
7
|
+
to the chat model and runs one model turn. Each turn is recorded as a durable
|
|
8
|
+
``flyte.trace`` leaf (via :func:`~flyteplugins.agents.core.durable_step`), so a
|
|
9
|
+
crash/retry replays the recorded response instead of re-calling (and re-billing)
|
|
10
|
+
the model.
|
|
11
|
+
- :func:`tool_node` — the tool-executing node. It runs the tool calls the model
|
|
12
|
+
emitted; each ``@tool``-wrapped task runs as a durable Flyte child action (its
|
|
13
|
+
own container/resources, retries, caching).
|
|
14
|
+
|
|
15
|
+
Both render their turns into the Flyte task report. Wire them into a standard
|
|
16
|
+
tool-calling loop::
|
|
17
|
+
|
|
18
|
+
from langgraph.graph import StateGraph, MessagesState, START
|
|
19
|
+
from langgraph.prebuilt import tools_condition
|
|
20
|
+
|
|
21
|
+
builder = StateGraph(MessagesState)
|
|
22
|
+
builder.add_node("ai", ai_node(model, tools))
|
|
23
|
+
builder.add_node("tools", tool_node(tools))
|
|
24
|
+
builder.add_edge(START, "ai")
|
|
25
|
+
builder.add_conditional_edges("ai", tools_condition)
|
|
26
|
+
builder.add_edge("tools", "ai")
|
|
27
|
+
graph = builder.compile()
|
|
28
|
+
"""
|
|
29
|
+
|
|
30
|
+
from __future__ import annotations
|
|
31
|
+
|
|
32
|
+
import json
|
|
33
|
+
import typing
|
|
34
|
+
|
|
35
|
+
from flyteplugins.agents.core import ReportTimeline, abbrev, durable_step, fingerprint
|
|
36
|
+
|
|
37
|
+
if typing.TYPE_CHECKING:
|
|
38
|
+
from langchain_core.language_models.chat_models import BaseChatModel
|
|
39
|
+
|
|
40
|
+
# A node is an (optionally async) callable ``state -> partial_state``.
|
|
41
|
+
Node = typing.Callable[[dict], typing.Awaitable[dict]]
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def _message_text(message: typing.Any) -> str:
|
|
45
|
+
content = getattr(message, "content", message)
|
|
46
|
+
return content if isinstance(content, str) else str(content)
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
def ai_node(
|
|
50
|
+
model: "BaseChatModel",
|
|
51
|
+
tools: typing.Sequence[typing.Any],
|
|
52
|
+
*,
|
|
53
|
+
name: str = "ai",
|
|
54
|
+
durable: bool = True,
|
|
55
|
+
observability: bool = True,
|
|
56
|
+
) -> Node:
|
|
57
|
+
"""Build the model-calling node for a tool-calling graph.
|
|
58
|
+
|
|
59
|
+
The returned node binds ``tools`` to ``model`` and runs a single model turn
|
|
60
|
+
over ``state["messages"]``, appending the model's response. Pass
|
|
61
|
+
``@tool``-wrapped tasks (or any LangChain ``BaseTool``) as ``tools``.
|
|
62
|
+
|
|
63
|
+
Args:
|
|
64
|
+
model: A LangChain chat model (e.g. ``ChatOpenAI(model="gpt-4o")``).
|
|
65
|
+
tools: The tools to expose to the model.
|
|
66
|
+
name: Node label (used for the graph node and the trace/report entry).
|
|
67
|
+
durable: Record each model turn via ``flyte.trace`` so retries replay it.
|
|
68
|
+
observability: Render each model turn into the Flyte task report.
|
|
69
|
+
|
|
70
|
+
Returns:
|
|
71
|
+
An async node ``state -> {"messages": [ai_message]}``.
|
|
72
|
+
"""
|
|
73
|
+
bound = model.bind_tools(list(tools))
|
|
74
|
+
timeline = ReportTimeline() if observability else None
|
|
75
|
+
|
|
76
|
+
async def _ai(state: dict) -> dict:
|
|
77
|
+
from langchain_core.messages import message_to_dict, messages_from_dict, messages_to_dict
|
|
78
|
+
|
|
79
|
+
messages = state["messages"]
|
|
80
|
+
|
|
81
|
+
async def _call() -> typing.Any:
|
|
82
|
+
return await bound.ainvoke(messages)
|
|
83
|
+
|
|
84
|
+
if durable:
|
|
85
|
+
key = fingerprint({"node": name, "messages": messages_to_dict(list(messages))})
|
|
86
|
+
response = await durable_step(
|
|
87
|
+
key,
|
|
88
|
+
_call,
|
|
89
|
+
name=f"{name}:model",
|
|
90
|
+
dumps=lambda m: json.dumps(message_to_dict(m)),
|
|
91
|
+
loads=lambda s: messages_from_dict([json.loads(s)])[0],
|
|
92
|
+
)
|
|
93
|
+
else:
|
|
94
|
+
response = await _call()
|
|
95
|
+
|
|
96
|
+
if timeline is not None:
|
|
97
|
+
tool_calls = getattr(response, "tool_calls", None) or []
|
|
98
|
+
if tool_calls:
|
|
99
|
+
detail = "→ " + ", ".join(tc["name"] for tc in tool_calls)
|
|
100
|
+
else:
|
|
101
|
+
detail = abbrev(_message_text(response), 200)
|
|
102
|
+
timeline.row(icon="🤖", label=name, meta="assistant", detail=detail)
|
|
103
|
+
|
|
104
|
+
return {"messages": [response]}
|
|
105
|
+
|
|
106
|
+
_ai.__name__ = name
|
|
107
|
+
return _ai
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
def tool_node(
|
|
111
|
+
tools: typing.Sequence[typing.Any],
|
|
112
|
+
*,
|
|
113
|
+
name: str = "tools",
|
|
114
|
+
observability: bool = True,
|
|
115
|
+
) -> Node:
|
|
116
|
+
"""Build the tool-executing node for a tool-calling graph.
|
|
117
|
+
|
|
118
|
+
The returned node reads the tool calls from the last message and runs each
|
|
119
|
+
one, appending a ``ToolMessage`` per call. ``@tool``-wrapped tasks run as
|
|
120
|
+
durable Flyte child actions; anything else runs as the tool defines.
|
|
121
|
+
|
|
122
|
+
Args:
|
|
123
|
+
tools: The tools available to execute (``@tool``-wrapped tasks or any
|
|
124
|
+
LangChain ``BaseTool``).
|
|
125
|
+
name: Node label (used for the report entry).
|
|
126
|
+
observability: Render each tool call/result into the Flyte task report.
|
|
127
|
+
|
|
128
|
+
Returns:
|
|
129
|
+
An async node ``state -> {"messages": [tool_message, ...]}``.
|
|
130
|
+
"""
|
|
131
|
+
registry = {getattr(t, "name", getattr(t, "__name__", "")): t for t in tools}
|
|
132
|
+
timeline = ReportTimeline() if observability else None
|
|
133
|
+
|
|
134
|
+
async def _tools(state: dict) -> dict:
|
|
135
|
+
from langchain_core.messages import ToolMessage
|
|
136
|
+
|
|
137
|
+
messages = state["messages"]
|
|
138
|
+
last = messages[-1] if messages else None
|
|
139
|
+
calls = getattr(last, "tool_calls", None) or []
|
|
140
|
+
|
|
141
|
+
results: list[typing.Any] = []
|
|
142
|
+
for call in calls:
|
|
143
|
+
tool_name = call["name"]
|
|
144
|
+
args = call.get("args", {}) or {}
|
|
145
|
+
call_id = call.get("id", "")
|
|
146
|
+
if timeline is not None:
|
|
147
|
+
timeline.row(icon="🛠️", label=tool_name, meta="tool", detail=abbrev(str(args), 160))
|
|
148
|
+
|
|
149
|
+
selected = registry.get(tool_name)
|
|
150
|
+
if selected is None:
|
|
151
|
+
output = f"Error: unknown tool '{tool_name}'"
|
|
152
|
+
else:
|
|
153
|
+
try:
|
|
154
|
+
output = await selected.ainvoke(args)
|
|
155
|
+
except Exception as exc: # surface tool errors back to the model
|
|
156
|
+
output = f"Error: {exc}"
|
|
157
|
+
|
|
158
|
+
if timeline is not None:
|
|
159
|
+
timeline.row(icon="🔧", label=tool_name, meta="tool result", detail=abbrev(str(output), 160))
|
|
160
|
+
results.append(ToolMessage(content=str(output), tool_call_id=call_id, name=tool_name))
|
|
161
|
+
|
|
162
|
+
return {"messages": results}
|
|
163
|
+
|
|
164
|
+
_tools.__name__ = name
|
|
165
|
+
return _tools
|
|
@@ -0,0 +1,196 @@
|
|
|
1
|
+
"""``run_agent`` — drive a LangGraph graph on Flyte.
|
|
2
|
+
|
|
3
|
+
The intended devex is that *you* build the ``StateGraph`` (with
|
|
4
|
+
:func:`~flyteplugins.agents.langgraph.ai_node` /
|
|
5
|
+
:func:`~flyteplugins.agents.langgraph.tool_node`), compile it, and hand the
|
|
6
|
+
compiled graph to ``run_agent(agent=...)``. ``run_agent`` runs that graph inside
|
|
7
|
+
your ``@env.task``: each model turn is durable (replayed on retry) and each tool
|
|
8
|
+
call runs as a durable Flyte child action.
|
|
9
|
+
|
|
10
|
+
As a convenience, passing ``tools`` (instead of ``agent``) builds a default
|
|
11
|
+
tool-calling graph for you from the same ``ai_node`` / ``tool_node`` building
|
|
12
|
+
blocks.
|
|
13
|
+
|
|
14
|
+
Observability: the run timeline — model turns and tool calls — is rendered into
|
|
15
|
+
the Flyte task report.
|
|
16
|
+
"""
|
|
17
|
+
|
|
18
|
+
from __future__ import annotations
|
|
19
|
+
|
|
20
|
+
import typing
|
|
21
|
+
|
|
22
|
+
from flyteplugins.agents.core import ReportTimeline, flush_report, sync_variant
|
|
23
|
+
|
|
24
|
+
from ._nodes import ai_node, tool_node
|
|
25
|
+
from ._tools import _coerce_tool
|
|
26
|
+
|
|
27
|
+
try: # langgraph is a hard dependency; guarded only so imports never hard-crash.
|
|
28
|
+
from langgraph.graph import END, START, MessagesState, StateGraph
|
|
29
|
+
from langgraph.prebuilt import tools_condition
|
|
30
|
+
except Exception: # pragma: no cover - langgraph missing
|
|
31
|
+
END = START = MessagesState = StateGraph = tools_condition = None # type: ignore[assignment]
|
|
32
|
+
|
|
33
|
+
# Module-level aliases so the builder path can be redirected in tests.
|
|
34
|
+
_StateGraph = StateGraph
|
|
35
|
+
_MessagesState = MessagesState
|
|
36
|
+
_START = START
|
|
37
|
+
_tools_condition = tools_condition
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def _resolve_chat_model(model: typing.Any) -> typing.Any:
|
|
41
|
+
"""Return a LangChain chat model for the default-graph builder.
|
|
42
|
+
|
|
43
|
+
A chat-model instance passes through unchanged. A ``provider:model`` string
|
|
44
|
+
resolves via ``langchain.chat_models.init_chat_model`` (requires the
|
|
45
|
+
``langchain`` package). ``None`` is an error — the caller must choose a model.
|
|
46
|
+
"""
|
|
47
|
+
if model is None:
|
|
48
|
+
raise ValueError(
|
|
49
|
+
"Provide `model=` when building the agent (or pass a pre-built `agent=`). "
|
|
50
|
+
'For example: `model=ChatOpenAI(model="gpt-4o")`.'
|
|
51
|
+
)
|
|
52
|
+
if not isinstance(model, str):
|
|
53
|
+
return model
|
|
54
|
+
try:
|
|
55
|
+
from langchain.chat_models import init_chat_model
|
|
56
|
+
except ImportError as e:
|
|
57
|
+
raise ImportError(
|
|
58
|
+
f"Resolving the model string {model!r} requires the `langchain` package. "
|
|
59
|
+
"Pass a chat-model instance instead, or install `langchain` to use "
|
|
60
|
+
"`provider:model` strings."
|
|
61
|
+
) from e
|
|
62
|
+
|
|
63
|
+
return init_chat_model(model)
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def _build_default_graph(
|
|
67
|
+
*,
|
|
68
|
+
model: typing.Any,
|
|
69
|
+
tools: typing.Sequence[typing.Any],
|
|
70
|
+
durable: bool,
|
|
71
|
+
observability: bool,
|
|
72
|
+
) -> typing.Any:
|
|
73
|
+
"""Build the standard tool-calling graph from ``ai_node`` + ``tool_node``."""
|
|
74
|
+
chat_model = _resolve_chat_model(model)
|
|
75
|
+
builder = _StateGraph(_MessagesState)
|
|
76
|
+
builder.add_node("ai", ai_node(chat_model, tools, name="ai", durable=durable, observability=observability))
|
|
77
|
+
builder.add_node("tools", tool_node(tools, name="tools", observability=observability))
|
|
78
|
+
builder.add_edge(_START, "ai")
|
|
79
|
+
builder.add_conditional_edges("ai", _tools_condition)
|
|
80
|
+
builder.add_edge("tools", "ai")
|
|
81
|
+
return builder.compile()
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def _final_text(result: typing.Any) -> str:
|
|
85
|
+
"""Extract the final assistant text from a graph's output state."""
|
|
86
|
+
if isinstance(result, dict):
|
|
87
|
+
messages = result.get("messages", [])
|
|
88
|
+
if messages:
|
|
89
|
+
last = messages[-1]
|
|
90
|
+
content = last.get("content") if isinstance(last, dict) else getattr(last, "content", None)
|
|
91
|
+
if content is not None:
|
|
92
|
+
return content if isinstance(content, str) else str(content)
|
|
93
|
+
return ""
|
|
94
|
+
return str(result) if result is not None else ""
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
async def run_agent(
|
|
98
|
+
input: str | typing.Any,
|
|
99
|
+
*,
|
|
100
|
+
tools: typing.Sequence[typing.Any] = (),
|
|
101
|
+
model: typing.Any = None,
|
|
102
|
+
instructions: str | None = None,
|
|
103
|
+
agent: typing.Any = None,
|
|
104
|
+
name: str = "langgraph-agent",
|
|
105
|
+
durable: bool = True,
|
|
106
|
+
observability: bool = True,
|
|
107
|
+
memory_key: str | None = None,
|
|
108
|
+
**run_kwargs: typing.Any,
|
|
109
|
+
) -> str:
|
|
110
|
+
"""Run a LangGraph graph and return the final text.
|
|
111
|
+
|
|
112
|
+
Await this from an async task as ``await run_agent(...)``; from a sync task
|
|
113
|
+
use :func:`run_agent_sync` instead.
|
|
114
|
+
|
|
115
|
+
Call this from inside an ``@env.task`` — that task is the durable parent.
|
|
116
|
+
Within it, each model turn is recorded via ``flyte.trace`` (replayed on
|
|
117
|
+
retry) and each tool call runs as a durable Flyte child action. Give the
|
|
118
|
+
enclosing task ``retries=...`` for self-healing and ``report=True`` to see
|
|
119
|
+
the agent timeline.
|
|
120
|
+
|
|
121
|
+
Provide either a pre-built ``agent`` (a compiled ``StateGraph`` you built
|
|
122
|
+
with :func:`ai_node` / :func:`tool_node`) or ``tools`` to have a default
|
|
123
|
+
tool-calling graph built for you. The two are mutually exclusive.
|
|
124
|
+
|
|
125
|
+
Args:
|
|
126
|
+
input: The user prompt (a ``str``) or a full graph input state (a dict).
|
|
127
|
+
tools: ``@tool``-wrapped tools or bare ``@env.task`` templates (used only
|
|
128
|
+
when ``agent`` is not given).
|
|
129
|
+
model: A LangChain chat-model instance (e.g. ``ChatOpenAI(model="gpt-4o")``)
|
|
130
|
+
or a ``provider:model`` string (resolved via ``init_chat_model``, which
|
|
131
|
+
requires the ``langchain`` package). Required when building the graph
|
|
132
|
+
(i.e. when ``agent`` is not given).
|
|
133
|
+
instructions: System prompt prepended to a built graph's messages.
|
|
134
|
+
agent: A pre-built compiled LangGraph graph. Mutually exclusive with ``tools``.
|
|
135
|
+
name: Graph name (for debugging/observability).
|
|
136
|
+
durable: Record each model turn via ``flyte.trace`` (built graphs only).
|
|
137
|
+
observability: Render the run timeline into the Flyte task report.
|
|
138
|
+
memory_key: Stable id (e.g. a user/thread id) for cross-run memory. When
|
|
139
|
+
set, the conversation transcript is persisted to a keyed ``MemoryStore``
|
|
140
|
+
and resumed on a later run with the same key.
|
|
141
|
+
**run_kwargs: Additional kwargs forwarded to the graph's ``ainvoke``.
|
|
142
|
+
|
|
143
|
+
Returns:
|
|
144
|
+
The graph's final assistant message as a string.
|
|
145
|
+
"""
|
|
146
|
+
from langchain_core.messages import HumanMessage, SystemMessage
|
|
147
|
+
|
|
148
|
+
from ._memory import load_messages, resolve_memory, save_messages
|
|
149
|
+
|
|
150
|
+
if agent is not None and tools:
|
|
151
|
+
raise ValueError("Pass either `agent` (with its own tools) or `tools`, not both.")
|
|
152
|
+
|
|
153
|
+
timeline = ReportTimeline() if observability else None
|
|
154
|
+
if timeline is not None:
|
|
155
|
+
timeline.heading("LangGraph agent")
|
|
156
|
+
|
|
157
|
+
# Cross-run memory: the prior transcript (if any) is prepended to the run's
|
|
158
|
+
# messages, and the full transcript is persisted back afterwards.
|
|
159
|
+
store = await resolve_memory(memory_key)
|
|
160
|
+
prior = await load_messages(store)
|
|
161
|
+
|
|
162
|
+
if agent is None:
|
|
163
|
+
agent = _build_default_graph(
|
|
164
|
+
model=model,
|
|
165
|
+
tools=[_coerce_tool(t) for t in tools],
|
|
166
|
+
durable=durable,
|
|
167
|
+
observability=observability,
|
|
168
|
+
)
|
|
169
|
+
seed: list[typing.Any] = []
|
|
170
|
+
# The system prompt is only needed once; on resumed runs it already lives
|
|
171
|
+
# in the prior transcript.
|
|
172
|
+
if instructions and not prior:
|
|
173
|
+
seed.append(SystemMessage(content=instructions))
|
|
174
|
+
seed.extend(prior)
|
|
175
|
+
seed.append(HumanMessage(content=input))
|
|
176
|
+
input_state: typing.Any = {"messages": seed}
|
|
177
|
+
elif isinstance(input, str):
|
|
178
|
+
input_state = {"messages": [*prior, HumanMessage(content=input)]}
|
|
179
|
+
else:
|
|
180
|
+
input_state = input or {}
|
|
181
|
+
if prior:
|
|
182
|
+
input_state = {**input_state, "messages": [*prior, *input_state.get("messages", [])]}
|
|
183
|
+
|
|
184
|
+
try:
|
|
185
|
+
result = await agent.ainvoke(input_state, **run_kwargs)
|
|
186
|
+
finally:
|
|
187
|
+
if observability:
|
|
188
|
+
await flush_report()
|
|
189
|
+
|
|
190
|
+
if store is not None and isinstance(result, dict) and result.get("messages"):
|
|
191
|
+
await save_messages(store, result["messages"])
|
|
192
|
+
|
|
193
|
+
return _final_text(result)
|
|
194
|
+
|
|
195
|
+
|
|
196
|
+
run_agent_sync = sync_variant(run_agent)
|
|
@@ -0,0 +1,170 @@
|
|
|
1
|
+
"""Turn Flyte tasks into LangGraph-compatible tools that run as durable actions.
|
|
2
|
+
|
|
3
|
+
LangGraph (via LangChain) drives tools that are ``BaseTool`` instances: it binds
|
|
4
|
+
them to the model (``model.bind_tools([...])``) so the LLM can call them, and it
|
|
5
|
+
executes them from a tool node. :func:`tool` wraps a Flyte ``@env.task`` as a
|
|
6
|
+
LangChain ``StructuredTool`` whose async body dispatches to the task via
|
|
7
|
+
``task.aio()`` — so when the graph executes the tool, it runs as a durable Flyte
|
|
8
|
+
child action (its own container/resources, with retries and caching) rather than
|
|
9
|
+
inline in the graph's process.
|
|
10
|
+
|
|
11
|
+
The returned tool is a first-class ``StructuredTool``: pass it straight to
|
|
12
|
+
``model.bind_tools(...)``, to :func:`~flyteplugins.agents.langgraph.tool_node`,
|
|
13
|
+
or to LangGraph's ``ToolNode``. It additionally exposes ``__wrapped_task__`` /
|
|
14
|
+
``task`` (so the backing task resolves to itself on the worker, no recursion) and
|
|
15
|
+
``__name__`` for convenience.
|
|
16
|
+
"""
|
|
17
|
+
|
|
18
|
+
from __future__ import annotations
|
|
19
|
+
|
|
20
|
+
import functools
|
|
21
|
+
import inspect
|
|
22
|
+
import json
|
|
23
|
+
import typing
|
|
24
|
+
from functools import partial
|
|
25
|
+
|
|
26
|
+
from flyte._task import AsyncFunctionTaskTemplate
|
|
27
|
+
from flyteplugins.agents.core import attach_tool_resolver, coerce_tool_args
|
|
28
|
+
|
|
29
|
+
try: # pragma: no cover - import shape only
|
|
30
|
+
from langchain_core.tools import StructuredTool
|
|
31
|
+
except Exception: # pragma: no cover
|
|
32
|
+
StructuredTool = None # type: ignore[assignment,misc]
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
if StructuredTool is not None:
|
|
36
|
+
|
|
37
|
+
class FlyteStructuredTool(StructuredTool):
|
|
38
|
+
"""A LangChain ``StructuredTool`` backed by a Flyte task.
|
|
39
|
+
|
|
40
|
+
Behaves exactly like a ``StructuredTool`` (so LangGraph's ``bind_tools`` /
|
|
41
|
+
``ToolNode`` accept it), while carrying the backing task so it resolves to
|
|
42
|
+
itself on the worker.
|
|
43
|
+
"""
|
|
44
|
+
|
|
45
|
+
flyte_task: typing.Any = None
|
|
46
|
+
|
|
47
|
+
@property
|
|
48
|
+
def task(self) -> typing.Any:
|
|
49
|
+
return self.flyte_task
|
|
50
|
+
|
|
51
|
+
@property
|
|
52
|
+
def __wrapped_task__(self) -> typing.Any:
|
|
53
|
+
return self.flyte_task
|
|
54
|
+
else: # pragma: no cover - langchain-core missing
|
|
55
|
+
FlyteStructuredTool = None # type: ignore[assignment,misc]
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def tool(
|
|
59
|
+
func: AsyncFunctionTaskTemplate | typing.Callable | None = None,
|
|
60
|
+
*,
|
|
61
|
+
name: str | None = None,
|
|
62
|
+
description: str | None = None,
|
|
63
|
+
) -> typing.Any:
|
|
64
|
+
"""Convert a Flyte task (or plain callable) into a LangChain ``StructuredTool``.
|
|
65
|
+
|
|
66
|
+
- For an ``@env.task``: returns a ``StructuredTool`` whose async body runs the
|
|
67
|
+
task as a durable Flyte child action when the graph invokes it. The input
|
|
68
|
+
schema is inferred from the task's typed signature. The backing task is
|
|
69
|
+
wired to :class:`~flyteplugins.agents.core.ToolTaskResolver` and exposed via
|
|
70
|
+
``__wrapped_task__`` so it resolves to itself on the worker (no recursion).
|
|
71
|
+
- For a plain (async) callable: returns a ``StructuredTool`` that runs it inline.
|
|
72
|
+
|
|
73
|
+
The result is a first-class LangGraph tool — bind it to a model or hand it to
|
|
74
|
+
:func:`~flyteplugins.agents.langgraph.tool_node` /
|
|
75
|
+
:func:`~flyteplugins.agents.langgraph.ai_node`.
|
|
76
|
+
|
|
77
|
+
Usable bare, parametrized, or as a direct call::
|
|
78
|
+
|
|
79
|
+
@tool
|
|
80
|
+
@env.task
|
|
81
|
+
async def get_weather(city: str) -> str: ...
|
|
82
|
+
"""
|
|
83
|
+
if func is None:
|
|
84
|
+
return partial(tool, name=name, description=description)
|
|
85
|
+
if isinstance(func, AsyncFunctionTaskTemplate):
|
|
86
|
+
return _task_to_tool(func, name=name, description=description)
|
|
87
|
+
return _callable_to_tool(func, name=name, description=description)
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
def _task_to_tool(
|
|
91
|
+
task: AsyncFunctionTaskTemplate,
|
|
92
|
+
*,
|
|
93
|
+
name: str | None = None,
|
|
94
|
+
description: str | None = None,
|
|
95
|
+
) -> typing.Any:
|
|
96
|
+
"""Build a LangChain ``StructuredTool`` from a Flyte task."""
|
|
97
|
+
tool_name = name or task.func.__name__
|
|
98
|
+
desc = (description or task.func.__doc__ or f"Run {tool_name}").strip()
|
|
99
|
+
|
|
100
|
+
# ``functools.wraps`` copies the task function's signature (via ``__wrapped__``)
|
|
101
|
+
# so ``StructuredTool.from_function`` infers the correct args schema, while the
|
|
102
|
+
# body dispatches to ``task.aio`` for durable execution. ``coerce_tool_args``
|
|
103
|
+
# relaxes LLM int->float args so Flyte's type engine doesn't reject e.g.
|
|
104
|
+
# ``amount_usd=42`` for a ``float`` param.
|
|
105
|
+
@functools.wraps(task.func)
|
|
106
|
+
async def _arun(**kwargs: typing.Any) -> str:
|
|
107
|
+
result = await task.aio(**coerce_tool_args(task, kwargs or {}))
|
|
108
|
+
return _as_content(result)
|
|
109
|
+
|
|
110
|
+
_arun.__name__ = tool_name
|
|
111
|
+
|
|
112
|
+
# Wire the shared resolver so the task resolves to itself on the worker.
|
|
113
|
+
attach_tool_resolver(task)
|
|
114
|
+
|
|
115
|
+
structured = FlyteStructuredTool.from_function(
|
|
116
|
+
coroutine=_arun,
|
|
117
|
+
name=tool_name,
|
|
118
|
+
description=desc,
|
|
119
|
+
flyte_task=task,
|
|
120
|
+
)
|
|
121
|
+
structured.__name__ = tool_name # type: ignore[attr-defined]
|
|
122
|
+
return structured
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
def _callable_to_tool(
|
|
126
|
+
func: typing.Callable,
|
|
127
|
+
*,
|
|
128
|
+
name: str | None = None,
|
|
129
|
+
description: str | None = None,
|
|
130
|
+
) -> typing.Any:
|
|
131
|
+
"""Build a LangChain ``StructuredTool`` from a plain callable."""
|
|
132
|
+
from langchain_core.tools import StructuredTool as _StructuredTool
|
|
133
|
+
|
|
134
|
+
tool_name = name or getattr(func, "__name__", "tool")
|
|
135
|
+
desc = (description or func.__doc__ or f"Run {tool_name}").strip()
|
|
136
|
+
|
|
137
|
+
@functools.wraps(func)
|
|
138
|
+
async def _arun(**kwargs: typing.Any) -> str:
|
|
139
|
+
out = func(**(kwargs or {}))
|
|
140
|
+
if inspect.isawaitable(out):
|
|
141
|
+
out = await out
|
|
142
|
+
return _as_content(out)
|
|
143
|
+
|
|
144
|
+
_arun.__name__ = tool_name
|
|
145
|
+
|
|
146
|
+
structured = _StructuredTool.from_function(coroutine=_arun, name=tool_name, description=desc)
|
|
147
|
+
structured.__name__ = tool_name # type: ignore[attr-defined]
|
|
148
|
+
return structured
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
def _as_content(result: typing.Any) -> str:
|
|
152
|
+
"""Convert a tool result to a string for LangChain's ``ToolMessage``."""
|
|
153
|
+
if isinstance(result, str):
|
|
154
|
+
return result
|
|
155
|
+
try:
|
|
156
|
+
return json.dumps(result, default=str)
|
|
157
|
+
except (TypeError, ValueError):
|
|
158
|
+
return str(result)
|
|
159
|
+
|
|
160
|
+
|
|
161
|
+
def _coerce_tool(t: typing.Any) -> typing.Any:
|
|
162
|
+
"""Coerce a bare ``@env.task`` (or plain callable) into a LangChain tool.
|
|
163
|
+
|
|
164
|
+
Already-wrapped tools (anything exposing ``ainvoke``) pass through unchanged.
|
|
165
|
+
"""
|
|
166
|
+
if isinstance(t, AsyncFunctionTaskTemplate):
|
|
167
|
+
return tool(t)
|
|
168
|
+
if callable(t) and not hasattr(t, "ainvoke"):
|
|
169
|
+
return tool(t)
|
|
170
|
+
return t
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: flyteplugins-agents-langgraph
|
|
3
|
+
Version: 2.5.17
|
|
4
|
+
Summary: Run LangGraph agents on Flyte.
|
|
5
|
+
Author-email: Niels Bantilan <niels@union.ai>
|
|
6
|
+
Requires-Python: >=3.10
|
|
7
|
+
Description-Content-Type: text/markdown
|
|
8
|
+
Requires-Dist: flyteplugins-agents-core
|
|
9
|
+
Requires-Dist: langgraph
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
flyteplugins/agents/langgraph/__init__.py,sha256=cMXMElr3SPLkZvKoq1HMwebjOk9mqPDqeOR3512mTWk,1067
|
|
2
|
+
flyteplugins/agents/langgraph/_memory.py,sha256=tgQ8Qp468m_xI2EI5RRWSVkvxjtBjr0dIiDP3AN8Vvk,2394
|
|
3
|
+
flyteplugins/agents/langgraph/_nodes.py,sha256=RBv3O3ZCrz0R45oXf1JNcPk5bYP28m7X3NBIxggxQJg,6212
|
|
4
|
+
flyteplugins/agents/langgraph/_run.py,sha256=DFCrgh6MQzH4RQqlp13v-i54gPoVp6hwWTpeOFPbibM,7958
|
|
5
|
+
flyteplugins/agents/langgraph/_tools.py,sha256=88SckNo4v-sOoVTSFLNHXQP9hs2Qh1RjzqywgOD40PY,6262
|
|
6
|
+
flyteplugins_agents_langgraph-2.5.17.dist-info/METADATA,sha256=01UTDZxL2vtuonG0bk-31mv6aUVeN2BjSopm_Oyxm3s,289
|
|
7
|
+
flyteplugins_agents_langgraph-2.5.17.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
|
|
8
|
+
flyteplugins_agents_langgraph-2.5.17.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
|
|
9
|
+
flyteplugins_agents_langgraph-2.5.17.dist-info/RECORD,,
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
flyteplugins
|