flyteplugins-agents-langchain 2.5.15__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/langchain/__init__.py +19 -0
- flyteplugins/agents/langchain/_durable.py +132 -0
- flyteplugins/agents/langchain/_memory.py +72 -0
- flyteplugins/agents/langchain/_run.py +175 -0
- flyteplugins/agents/langchain/_tools.py +174 -0
- flyteplugins_agents_langchain-2.5.15.dist-info/METADATA +9 -0
- flyteplugins_agents_langchain-2.5.15.dist-info/RECORD +9 -0
- flyteplugins_agents_langchain-2.5.15.dist-info/WHEEL +5 -0
- flyteplugins_agents_langchain-2.5.15.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
"""LangChain adapter for Flyte.
|
|
2
|
+
|
|
3
|
+
Bring your own LangChain agent and run it durably on Flyte. The adapter provides:
|
|
4
|
+
|
|
5
|
+
- :func:`tool` — turn a Flyte ``@env.task`` into a LangChain ``StructuredTool``
|
|
6
|
+
(a ``BaseTool``) that executes as a durable child action (own container/GPU,
|
|
7
|
+
retries, caching).
|
|
8
|
+
- :func:`run_agent` — run the LangChain agent (a compiled ``create_agent`` graph)
|
|
9
|
+
inside your task and return the final answer. Either pass a pre-built ``agent``
|
|
10
|
+
or let it build one from ``tools`` + ``model`` + ``instructions``.
|
|
11
|
+
|
|
12
|
+
Each tool call runs as a durable Flyte child action, and the run timeline is
|
|
13
|
+
rendered into the Flyte task report.
|
|
14
|
+
"""
|
|
15
|
+
|
|
16
|
+
from ._run import run_agent, run_agent_sync
|
|
17
|
+
from ._tools import tool
|
|
18
|
+
|
|
19
|
+
__all__ = ["run_agent", "run_agent_sync", "tool"]
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
"""Durable, replayable model turns for LangChain chat models.
|
|
2
|
+
|
|
3
|
+
LangChain's ``create_agent`` graph owns the agent loop; each iteration calls the
|
|
4
|
+
chat model once (a "model turn"). :class:`DurableChatModel` wraps any
|
|
5
|
+
``BaseChatModel`` so every turn is recorded through the shared
|
|
6
|
+
:func:`~flyteplugins.agents.core.durable_step` (a ``flyte.trace`` leaf). Inside a
|
|
7
|
+
Flyte task this means a crashed/retried run replays completed turns from their
|
|
8
|
+
recorded outputs instead of re-calling (and re-billing) the model. Tool calls run
|
|
9
|
+
as durable child actions (see :func:`flyteplugins.agents.langchain.tool`), so the
|
|
10
|
+
whole agent run becomes crash-resilient when the enclosing task carries
|
|
11
|
+
``retries=...``.
|
|
12
|
+
|
|
13
|
+
The turn is recorded as JSON: the generated messages of the model's
|
|
14
|
+
``ChatResult`` are serialized with ``message_to_dict`` and rebuilt with
|
|
15
|
+
``messages_from_dict`` (the same round-trip the langgraph adapter uses), which
|
|
16
|
+
keeps the recorded turn human-readable in the Flyte UI.
|
|
17
|
+
|
|
18
|
+
Tool-calling still works because :meth:`DurableChatModel.bind_tools` delegates to
|
|
19
|
+
the inner model to format the tools, then re-binds the resulting kwargs to *this*
|
|
20
|
+
wrapper — so ``create_agent``'s bound runnable still routes generation through the
|
|
21
|
+
durable override.
|
|
22
|
+
"""
|
|
23
|
+
|
|
24
|
+
from __future__ import annotations
|
|
25
|
+
|
|
26
|
+
import json
|
|
27
|
+
import typing
|
|
28
|
+
|
|
29
|
+
from flyteplugins.agents.core import durable_step, fingerprint
|
|
30
|
+
from langchain_core.language_models.chat_models import BaseChatModel
|
|
31
|
+
|
|
32
|
+
if typing.TYPE_CHECKING:
|
|
33
|
+
from langchain_core.messages import BaseMessage
|
|
34
|
+
from langchain_core.outputs import ChatResult
|
|
35
|
+
from langchain_core.runnables import Runnable
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def _dumps_result(result: "ChatResult") -> str:
|
|
39
|
+
"""Serialize a ``ChatResult``'s generated messages to JSON (readable in the UI)."""
|
|
40
|
+
from langchain_core.messages import message_to_dict
|
|
41
|
+
|
|
42
|
+
return json.dumps([message_to_dict(gen.message) for gen in result.generations])
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def _loads_result(payload: str) -> "ChatResult":
|
|
46
|
+
"""Rebuild a ``ChatResult`` from the JSON written by :func:`_dumps_result`."""
|
|
47
|
+
from langchain_core.messages import messages_from_dict
|
|
48
|
+
from langchain_core.outputs import ChatGeneration, ChatResult
|
|
49
|
+
|
|
50
|
+
messages = messages_from_dict(json.loads(payload))
|
|
51
|
+
return ChatResult(generations=[ChatGeneration(message=m) for m in messages])
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
class DurableChatModel(BaseChatModel):
|
|
55
|
+
"""Wrap a ``BaseChatModel`` so each model turn is durable and replayable.
|
|
56
|
+
|
|
57
|
+
``_agenerate`` (async) and ``_generate`` (sync) delegate to the inner model
|
|
58
|
+
and record the turn via ``durable_step``. Pass an instance to
|
|
59
|
+
``create_agent(DurableChatModel(inner=model), tools, ...)``; ``bind_tools``
|
|
60
|
+
and other capabilities are delegated to the inner model so tool-calling
|
|
61
|
+
behaves exactly as the inner model does.
|
|
62
|
+
|
|
63
|
+
Durability is best-effort: if anything in the durable path raises, the turn
|
|
64
|
+
falls back to a direct inner call so a run is never broken by it.
|
|
65
|
+
"""
|
|
66
|
+
|
|
67
|
+
inner: BaseChatModel
|
|
68
|
+
"""The wrapped chat model that actually generates responses."""
|
|
69
|
+
|
|
70
|
+
@property
|
|
71
|
+
def _llm_type(self) -> str:
|
|
72
|
+
return f"durable-{self.inner._llm_type}"
|
|
73
|
+
|
|
74
|
+
def _turn_key(self, messages: typing.Sequence["BaseMessage"], **kwargs: typing.Any) -> str:
|
|
75
|
+
"""Deterministic memo key for a model turn — serialized messages + bound tool names."""
|
|
76
|
+
from langchain_core.messages import messages_to_dict
|
|
77
|
+
|
|
78
|
+
tools = kwargs.get("tools") or []
|
|
79
|
+
tool_names = sorted(
|
|
80
|
+
(t.get("function", {}).get("name") or t.get("name") or str(t)) if isinstance(t, dict) else str(t)
|
|
81
|
+
for t in tools
|
|
82
|
+
)
|
|
83
|
+
return fingerprint(
|
|
84
|
+
{
|
|
85
|
+
"type": self._llm_type,
|
|
86
|
+
"messages": messages_to_dict(list(messages)),
|
|
87
|
+
"tools": tool_names,
|
|
88
|
+
"stop": kwargs.get("stop"),
|
|
89
|
+
}
|
|
90
|
+
)
|
|
91
|
+
|
|
92
|
+
async def _agenerate(self, messages, stop=None, run_manager=None, **kwargs) -> "ChatResult": # type: ignore[override]
|
|
93
|
+
async def _call() -> "ChatResult":
|
|
94
|
+
return await self.inner._agenerate(messages, stop=stop, run_manager=run_manager, **kwargs)
|
|
95
|
+
|
|
96
|
+
try:
|
|
97
|
+
key = self._turn_key(messages, stop=stop, **kwargs)
|
|
98
|
+
return await durable_step(
|
|
99
|
+
key,
|
|
100
|
+
_call,
|
|
101
|
+
name="model_turn",
|
|
102
|
+
dumps=_dumps_result,
|
|
103
|
+
loads=_loads_result,
|
|
104
|
+
)
|
|
105
|
+
except Exception: # pragma: no cover - durability must never break a run
|
|
106
|
+
return await _call()
|
|
107
|
+
|
|
108
|
+
def _generate(self, messages, stop=None, run_manager=None, **kwargs) -> "ChatResult": # type: ignore[override]
|
|
109
|
+
# ``durable_step`` is async; the sync path delegates straight through so
|
|
110
|
+
# ``create_agent``'s sync callers keep working (durability applies to the
|
|
111
|
+
# async agent loop, which is the one Flyte tasks drive).
|
|
112
|
+
return self.inner._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
|
|
113
|
+
|
|
114
|
+
def bind_tools(self, tools: typing.Sequence[typing.Any], **kwargs: typing.Any) -> "Runnable":
|
|
115
|
+
"""Format tools via the inner model, but bind them to *this* wrapper.
|
|
116
|
+
|
|
117
|
+
The inner model knows how to convert tools into its provider format; we
|
|
118
|
+
reuse that, then re-bind the resulting kwargs to ``self`` so the runnable
|
|
119
|
+
``create_agent`` invokes still routes generation through the durable
|
|
120
|
+
override (rather than the inner model directly).
|
|
121
|
+
"""
|
|
122
|
+
bound = self.inner.bind_tools(tools, **kwargs)
|
|
123
|
+
bound_kwargs = dict(getattr(bound, "kwargs", {}) or {})
|
|
124
|
+
return self.bind(**bound_kwargs)
|
|
125
|
+
|
|
126
|
+
def get_num_tokens(self, text: str) -> int:
|
|
127
|
+
return self.inner.get_num_tokens(text)
|
|
128
|
+
|
|
129
|
+
def get_num_tokens_from_messages(self, messages, tools=None) -> int: # type: ignore[override]
|
|
130
|
+
if tools is not None:
|
|
131
|
+
return self.inner.get_num_tokens_from_messages(messages, tools=tools)
|
|
132
|
+
return self.inner.get_num_tokens_from_messages(messages)
|
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
"""Cross-run LangChain memory — a thin bridge over Flyte's keyed ``MemoryStore``.
|
|
2
|
+
|
|
3
|
+
LangChain's ``create_agent`` graph is driven with a messages state
|
|
4
|
+
(``graph.ainvoke({"messages": [...]})``) and, by default, keeps no state across
|
|
5
|
+
runs. This module bridges that: it resolves a keyed :class:`MemoryStore`, loads a
|
|
6
|
+
prior conversation from a path-addressed JSON slot, and writes the full
|
|
7
|
+
transcript back after the run — so a later run with the same ``memory_key``
|
|
8
|
+
continues the conversation.
|
|
9
|
+
|
|
10
|
+
The transcript is stored as ``messages_to_dict(...)`` output (the same
|
|
11
|
+
serialization the langgraph adapter uses) in a single JSON slot, and rebuilt with
|
|
12
|
+
``messages_from_dict``. All operations are best-effort: any failure leaves the run
|
|
13
|
+
untouched (memory never breaks a run).
|
|
14
|
+
"""
|
|
15
|
+
|
|
16
|
+
from __future__ import annotations
|
|
17
|
+
|
|
18
|
+
import typing
|
|
19
|
+
|
|
20
|
+
from flyte._logging import logger
|
|
21
|
+
from flyteplugins.agents.core import resolve_memory as _resolve_memory
|
|
22
|
+
|
|
23
|
+
if typing.TYPE_CHECKING:
|
|
24
|
+
from langchain_core.messages import BaseMessage
|
|
25
|
+
|
|
26
|
+
# Path-addressed memory slot holding the serialized conversation transcript
|
|
27
|
+
# (``messages_to_dict`` output) inside the MemoryStore.
|
|
28
|
+
_HISTORY_PATH = "langchain/history.json"
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
async def resolve_memory(memory_key: str | None) -> typing.Any | None:
|
|
32
|
+
"""Resolve a keyed MemoryStore for LangChain cross-run memory, or ``None``.
|
|
33
|
+
|
|
34
|
+
Best-effort: returns ``None`` when ``memory_key`` is falsy or no durable
|
|
35
|
+
store can be resolved, so memory never breaks a run.
|
|
36
|
+
"""
|
|
37
|
+
if not memory_key:
|
|
38
|
+
return None
|
|
39
|
+
return await _resolve_memory(memory_key)
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
async def load_history(store: typing.Any) -> list["BaseMessage"]:
|
|
43
|
+
"""Load and deserialize the prior conversation from ``store``.
|
|
44
|
+
|
|
45
|
+
Returns an empty list when there is no prior history or on any error.
|
|
46
|
+
"""
|
|
47
|
+
if store is None:
|
|
48
|
+
return []
|
|
49
|
+
try:
|
|
50
|
+
from langchain_core.messages import messages_from_dict
|
|
51
|
+
|
|
52
|
+
raw = await store.read_json.aio(_HISTORY_PATH, [])
|
|
53
|
+
return messages_from_dict(raw) if raw else []
|
|
54
|
+
except Exception: # pragma: no cover - memory is best-effort, never fatal
|
|
55
|
+
logger.warning("Could not load LangChain memory; continuing without prior history.")
|
|
56
|
+
return []
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
async def save_history(store: typing.Any, messages: typing.Sequence["BaseMessage"]) -> None:
|
|
60
|
+
"""Serialize the full conversation ``messages`` and persist them to ``store``.
|
|
61
|
+
|
|
62
|
+
Best-effort: logs and returns on any error so memory never breaks a run.
|
|
63
|
+
"""
|
|
64
|
+
if store is None or not messages:
|
|
65
|
+
return
|
|
66
|
+
try:
|
|
67
|
+
from langchain_core.messages import messages_to_dict
|
|
68
|
+
|
|
69
|
+
await store.write_json.aio(_HISTORY_PATH, messages_to_dict(list(messages)))
|
|
70
|
+
await store.save.aio()
|
|
71
|
+
except Exception: # pragma: no cover - memory is best-effort, never fatal
|
|
72
|
+
logger.warning("Could not save LangChain memory; conversation will not be resumed.")
|
|
@@ -0,0 +1,175 @@
|
|
|
1
|
+
"""``run_agent`` — run a LangChain agent on Flyte.
|
|
2
|
+
|
|
3
|
+
LangChain's agent framework owns the loop. ``run_agent`` runs that loop inside your
|
|
4
|
+
``@env.task``: it builds an agent (a compiled LangGraph graph) with Flyte-task tools,
|
|
5
|
+
drives it, and returns the final answer. Each tool call runs as a durable Flyte child
|
|
6
|
+
action (its own container/resources, with retries and caching).
|
|
7
|
+
|
|
8
|
+
In langchain 1.x the agent is built with ``langchain.agents.create_agent(model, tools,
|
|
9
|
+
system_prompt=...)``, which returns a compiled graph that is driven with a messages
|
|
10
|
+
state: ``await graph.ainvoke({"messages": [{"role": "user", "content": input}]})``, and
|
|
11
|
+
the final text is ``result["messages"][-1].content``.
|
|
12
|
+
|
|
13
|
+
Observability: the run timeline is rendered into the Flyte task report.
|
|
14
|
+
|
|
15
|
+
The adapter minimizes delta between native LangChain code and Flyte integration by
|
|
16
|
+
exposing tools that are drop-in ``BaseTool`` instances.
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
from __future__ import annotations
|
|
20
|
+
|
|
21
|
+
import typing
|
|
22
|
+
|
|
23
|
+
from flyteplugins.agents.core import ReportTimeline, flush_report, sync_variant
|
|
24
|
+
|
|
25
|
+
from ._durable import DurableChatModel
|
|
26
|
+
from ._memory import load_history, resolve_memory, save_history
|
|
27
|
+
from ._tools import _coerce_tool
|
|
28
|
+
|
|
29
|
+
# Module-level alias for test monkeypatching: ``_create_agent`` lets tests
|
|
30
|
+
# substitute the graph builder without constructing a real model or graph.
|
|
31
|
+
_create_agent = None
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _final_text(result: typing.Any) -> str:
|
|
35
|
+
"""Extract the agent's final text from a compiled-graph result.
|
|
36
|
+
|
|
37
|
+
``create_agent`` graphs return a messages state ``{"messages": [...]}``; the
|
|
38
|
+
final answer is the content of the last message. Falls back gracefully for
|
|
39
|
+
other shapes (e.g. a legacy ``{"output": ...}`` dict or a bare string).
|
|
40
|
+
"""
|
|
41
|
+
if isinstance(result, dict):
|
|
42
|
+
messages = result.get("messages")
|
|
43
|
+
if messages:
|
|
44
|
+
content = getattr(messages[-1], "content", messages[-1])
|
|
45
|
+
return content if isinstance(content, str) else str(content)
|
|
46
|
+
if "output" in result:
|
|
47
|
+
return str(result["output"])
|
|
48
|
+
return ""
|
|
49
|
+
return str(result)
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
async def run_agent(
|
|
53
|
+
input: str,
|
|
54
|
+
*,
|
|
55
|
+
tools: typing.Sequence[typing.Any] = (),
|
|
56
|
+
model: typing.Any = None,
|
|
57
|
+
instructions: str | None = None,
|
|
58
|
+
agent: typing.Any = None,
|
|
59
|
+
name: str = "langchain-agent",
|
|
60
|
+
durable: bool = True,
|
|
61
|
+
observability: bool = True,
|
|
62
|
+
memory_key: str | None = None,
|
|
63
|
+
**agent_kwargs: typing.Any,
|
|
64
|
+
) -> str:
|
|
65
|
+
"""Run a LangChain agent with the given tools and prompt; return the final text.
|
|
66
|
+
|
|
67
|
+
Await this from an async task as ``await run_agent(...)``; from a sync task
|
|
68
|
+
use :func:`run_agent_sync` instead.
|
|
69
|
+
|
|
70
|
+
Call this from inside an ``@env.task`` — that task is the durable parent.
|
|
71
|
+
Within it, each tool call runs as a durable Flyte child action. Give the
|
|
72
|
+
enclosing task ``retries=...`` for self-healing and ``report=True`` to see
|
|
73
|
+
the agent timeline.
|
|
74
|
+
|
|
75
|
+
Provide either a pre-built ``agent`` (a compiled graph from ``create_agent``)
|
|
76
|
+
or ``tools`` + ``model`` to have one built for you.
|
|
77
|
+
|
|
78
|
+
Args:
|
|
79
|
+
input: The user prompt.
|
|
80
|
+
tools: ``tool``-wrapped tools or bare ``@env.task`` templates.
|
|
81
|
+
model: A LangChain-compatible chat model (e.g.
|
|
82
|
+
``ChatOpenAI(model="gpt-4o")``). Required when ``agent`` is not
|
|
83
|
+
given — one is built using this model.
|
|
84
|
+
instructions: System prompt for the built agent.
|
|
85
|
+
agent: A pre-built LangChain agent (a compiled ``create_agent`` graph).
|
|
86
|
+
Mutually exclusive with ``tools``.
|
|
87
|
+
name: Agent name (for debugging/observability).
|
|
88
|
+
durable: Record/replay each model turn via ``flyte.trace``. Applies when
|
|
89
|
+
a model is being built (``tools`` + ``model``, or a caller-passed
|
|
90
|
+
``BaseChatModel``) — the model is wrapped in :class:`DurableChatModel`.
|
|
91
|
+
A fully pre-built compiled ``agent`` cannot be rewrapped, so its model
|
|
92
|
+
turns are not made durable (its tool calls remain durable regardless).
|
|
93
|
+
observability: Render the run timeline into the Flyte task report.
|
|
94
|
+
memory_key: Stable id (e.g. a user/thread id) for cross-run memory.
|
|
95
|
+
When set, conversation history is persisted to a keyed ``MemoryStore``
|
|
96
|
+
and resumed on a later run with the same key.
|
|
97
|
+
**agent_kwargs: Additional kwargs forwarded to ``create_agent``.
|
|
98
|
+
|
|
99
|
+
Returns:
|
|
100
|
+
The agent's final output as a string.
|
|
101
|
+
"""
|
|
102
|
+
timeline = ReportTimeline() if observability else None
|
|
103
|
+
if timeline is not None:
|
|
104
|
+
timeline.heading("LangChain agent")
|
|
105
|
+
|
|
106
|
+
if agent is not None and tools:
|
|
107
|
+
raise ValueError("Pass either `agent` (with its own tools) or `tools`, not both.")
|
|
108
|
+
|
|
109
|
+
# Build the agent (a compiled graph) if not provided.
|
|
110
|
+
if agent is None:
|
|
111
|
+
if _create_agent is None:
|
|
112
|
+
from langchain.agents import create_agent as create_agent_fn
|
|
113
|
+
else:
|
|
114
|
+
create_agent_fn = _create_agent
|
|
115
|
+
|
|
116
|
+
if model is None:
|
|
117
|
+
raise ValueError(
|
|
118
|
+
"Provide `model=` when building the agent (or pass a pre-built `agent=`). "
|
|
119
|
+
'For example: `model=ChatOpenAI(model="gpt-4o")`.'
|
|
120
|
+
)
|
|
121
|
+
|
|
122
|
+
# Wrap the inner chat model so each model turn is durable/replayable. We
|
|
123
|
+
# can only do this on the builder path (a fully pre-built compiled ``agent``
|
|
124
|
+
# owns its own model, which we cannot reach to rewrap).
|
|
125
|
+
if durable:
|
|
126
|
+
model = _wrap_durable(model)
|
|
127
|
+
|
|
128
|
+
tool_objs = [_coerce_tool(t) for t in tools]
|
|
129
|
+
system_prompt = instructions or f"You are a helpful assistant named {name}."
|
|
130
|
+
agent = create_agent_fn(model, tool_objs, system_prompt=system_prompt, **agent_kwargs)
|
|
131
|
+
|
|
132
|
+
# Cross-run memory: load prior history and prepend it to the messages passed
|
|
133
|
+
# to the graph, then persist the full transcript back after the run.
|
|
134
|
+
store = await resolve_memory(memory_key)
|
|
135
|
+
prior = await load_history(store)
|
|
136
|
+
messages = [*prior, {"role": "user", "content": input}]
|
|
137
|
+
|
|
138
|
+
result = await agent.ainvoke({"messages": messages})
|
|
139
|
+
|
|
140
|
+
await save_history(store, _result_messages(result))
|
|
141
|
+
|
|
142
|
+
final = _final_text(result)
|
|
143
|
+
|
|
144
|
+
if observability:
|
|
145
|
+
await flush_report()
|
|
146
|
+
|
|
147
|
+
return final or ""
|
|
148
|
+
|
|
149
|
+
|
|
150
|
+
def _wrap_durable(model: typing.Any) -> typing.Any:
|
|
151
|
+
"""Wrap a chat model in :class:`DurableChatModel` when possible.
|
|
152
|
+
|
|
153
|
+
Best-effort: only ``BaseChatModel`` instances are wrappable; anything else
|
|
154
|
+
(or any failure) is returned unchanged so durability never breaks a run.
|
|
155
|
+
"""
|
|
156
|
+
try:
|
|
157
|
+
from langchain_core.language_models.chat_models import BaseChatModel
|
|
158
|
+
|
|
159
|
+
if isinstance(model, BaseChatModel) and not isinstance(model, DurableChatModel):
|
|
160
|
+
return DurableChatModel(inner=model)
|
|
161
|
+
except Exception: # pragma: no cover - durability is best-effort, never fatal
|
|
162
|
+
pass
|
|
163
|
+
return model
|
|
164
|
+
|
|
165
|
+
|
|
166
|
+
def _result_messages(result: typing.Any) -> list[typing.Any]:
|
|
167
|
+
"""Extract the message list from a compiled-graph result (empty on other shapes)."""
|
|
168
|
+
if isinstance(result, dict):
|
|
169
|
+
messages = result.get("messages")
|
|
170
|
+
if messages:
|
|
171
|
+
return list(messages)
|
|
172
|
+
return []
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
run_agent_sync = sync_variant(run_agent)
|
|
@@ -0,0 +1,174 @@
|
|
|
1
|
+
"""Turn Flyte tasks into LangChain tools that execute as durable actions.
|
|
2
|
+
|
|
3
|
+
LangChain accepts ``BaseTool`` instances as tools. :func:`tool` wraps a Flyte
|
|
4
|
+
``@env.task`` as a LangChain ``StructuredTool`` whose async coroutine dispatches
|
|
5
|
+
to the task via ``task.aio()`` — so when the agent calls the tool, it runs as a
|
|
6
|
+
durable Flyte child action (its own container/resources, with retries and
|
|
7
|
+
caching) rather than inline in the agent's process.
|
|
8
|
+
|
|
9
|
+
The returned object is a real ``StructuredTool`` (a ``BaseTool``), so it drops
|
|
10
|
+
straight into ``create_agent(model, tools=[...])``. It additionally exposes
|
|
11
|
+
``__wrapped_task__`` and ``task`` (via direct attribute assignment, which
|
|
12
|
+
``StructuredTool`` permits) and wires the backing task to
|
|
13
|
+
:class:`~flyteplugins.agents.core.ToolTaskResolver` so it resolves to itself on
|
|
14
|
+
the worker (no recursion).
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
from __future__ import annotations
|
|
18
|
+
|
|
19
|
+
import inspect
|
|
20
|
+
import json
|
|
21
|
+
import typing
|
|
22
|
+
from functools import partial
|
|
23
|
+
|
|
24
|
+
from flyte._task import AsyncFunctionTaskTemplate
|
|
25
|
+
from flyteplugins.agents.core import attach_tool_resolver, coerce_tool_args
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def tool(
|
|
29
|
+
func: AsyncFunctionTaskTemplate | typing.Callable | None = None,
|
|
30
|
+
*,
|
|
31
|
+
name: str | None = None,
|
|
32
|
+
description: str | None = None,
|
|
33
|
+
) -> typing.Any:
|
|
34
|
+
"""Convert a Flyte task (or plain callable) into a LangChain ``StructuredTool``.
|
|
35
|
+
|
|
36
|
+
- For an ``@env.task``: returns a ``StructuredTool`` whose async coroutine runs
|
|
37
|
+
the task as a durable Flyte child action when the agent invokes it. The input
|
|
38
|
+
schema is derived from the task's typed signature. The backing task is wired to
|
|
39
|
+
:class:`~flyteplugins.agents.core.ToolTaskResolver` and exposed via
|
|
40
|
+
``__wrapped_task__`` so it resolves to itself on the worker (no recursion).
|
|
41
|
+
- For a plain (async) callable: returns a ``StructuredTool`` that runs it inline.
|
|
42
|
+
|
|
43
|
+
Usable bare, parametrized, or as a direct call::
|
|
44
|
+
|
|
45
|
+
@tool
|
|
46
|
+
@env.task
|
|
47
|
+
async def get_weather(city: str) -> str: ...
|
|
48
|
+
"""
|
|
49
|
+
if func is None:
|
|
50
|
+
return partial(tool, name=name, description=description)
|
|
51
|
+
if isinstance(func, AsyncFunctionTaskTemplate):
|
|
52
|
+
return _task_to_tool(func, name=name, description=description)
|
|
53
|
+
return _callable_to_tool(func, name=name, description=description)
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def _task_to_tool(
|
|
57
|
+
task: AsyncFunctionTaskTemplate,
|
|
58
|
+
*,
|
|
59
|
+
name: str | None = None,
|
|
60
|
+
description: str | None = None,
|
|
61
|
+
) -> typing.Any:
|
|
62
|
+
"""Build a LangChain ``StructuredTool`` from a Flyte task."""
|
|
63
|
+
from langchain_core.tools import StructuredTool
|
|
64
|
+
|
|
65
|
+
tool_name = name or task.func.__name__
|
|
66
|
+
desc = (description or task.func.__doc__ or f"Run {tool_name}").strip()
|
|
67
|
+
|
|
68
|
+
async def _arun(**kwargs: typing.Any) -> str:
|
|
69
|
+
# In a Flyte task context this submits a durable child action; locally it
|
|
70
|
+
# runs inline. ``coerce_tool_args`` relaxes LLM int->float args so Flyte's
|
|
71
|
+
# type engine doesn't reject e.g. ``amount_usd=42`` for a ``float`` param.
|
|
72
|
+
result = await task.aio(**coerce_tool_args(task, kwargs or {}))
|
|
73
|
+
return _as_content(result)
|
|
74
|
+
|
|
75
|
+
# Wire the shared resolver so the task resolves to itself on the worker.
|
|
76
|
+
attach_tool_resolver(task)
|
|
77
|
+
|
|
78
|
+
# Derive an explicit args schema from the task's typed signature. The coroutine
|
|
79
|
+
# above is ``**kwargs``-only, so LangChain's own inference would produce a single
|
|
80
|
+
# ``kwargs`` object param — we build the real pydantic model instead.
|
|
81
|
+
args_schema = _args_schema_from_callable(task.func, tool_name)
|
|
82
|
+
|
|
83
|
+
structured = StructuredTool.from_function(
|
|
84
|
+
func=None,
|
|
85
|
+
coroutine=_arun,
|
|
86
|
+
name=tool_name,
|
|
87
|
+
description=desc,
|
|
88
|
+
args_schema=args_schema,
|
|
89
|
+
)
|
|
90
|
+
|
|
91
|
+
# Expose the real task and a convenient ``task`` alias so callers/tests can reach
|
|
92
|
+
# the backing task. ``StructuredTool`` is a pydantic model whose ``__setattr__``
|
|
93
|
+
# rejects non-field names, so set them through ``object.__setattr__`` (dunders
|
|
94
|
+
# like ``__wrapped_task__`` would bypass it, but keep both paths uniform).
|
|
95
|
+
object.__setattr__(structured, "__wrapped_task__", task)
|
|
96
|
+
object.__setattr__(structured, "task", task)
|
|
97
|
+
return structured
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
def _callable_to_tool(
|
|
101
|
+
func: typing.Callable,
|
|
102
|
+
*,
|
|
103
|
+
name: str | None = None,
|
|
104
|
+
description: str | None = None,
|
|
105
|
+
) -> typing.Any:
|
|
106
|
+
"""Build a LangChain ``StructuredTool`` from a plain callable."""
|
|
107
|
+
from langchain_core.tools import StructuredTool
|
|
108
|
+
|
|
109
|
+
tool_name = name or getattr(func, "__name__", "tool")
|
|
110
|
+
desc = (description or func.__doc__ or f"Run {tool_name}").strip()
|
|
111
|
+
|
|
112
|
+
async def _arun(**kwargs: typing.Any) -> str:
|
|
113
|
+
out = func(**(kwargs or {}))
|
|
114
|
+
if inspect.isawaitable(out):
|
|
115
|
+
out = await out
|
|
116
|
+
return _as_content(out)
|
|
117
|
+
|
|
118
|
+
args_schema = _args_schema_from_callable(func, tool_name)
|
|
119
|
+
|
|
120
|
+
return StructuredTool.from_function(
|
|
121
|
+
func=None,
|
|
122
|
+
coroutine=_arun,
|
|
123
|
+
name=tool_name,
|
|
124
|
+
description=desc,
|
|
125
|
+
args_schema=args_schema,
|
|
126
|
+
)
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
def _args_schema_from_callable(func: typing.Callable, tool_name: str) -> typing.Any | None:
|
|
130
|
+
"""Build a pydantic ``args_schema`` from a callable's typed signature.
|
|
131
|
+
|
|
132
|
+
Returns ``None`` (letting LangChain infer) if the signature can't be resolved.
|
|
133
|
+
The model's fields mirror the callable's parameters, with annotations and
|
|
134
|
+
defaults preserved so the LLM sees a correct tool schema.
|
|
135
|
+
"""
|
|
136
|
+
from pydantic import create_model
|
|
137
|
+
|
|
138
|
+
try:
|
|
139
|
+
hints = typing.get_type_hints(func)
|
|
140
|
+
sig = inspect.signature(func)
|
|
141
|
+
except Exception: # pragma: no cover - unresolved annotations; let LangChain infer
|
|
142
|
+
return None
|
|
143
|
+
|
|
144
|
+
fields: dict[str, typing.Any] = {}
|
|
145
|
+
for pname, param in sig.parameters.items():
|
|
146
|
+
if pname == "self" or param.kind in (
|
|
147
|
+
inspect.Parameter.VAR_POSITIONAL,
|
|
148
|
+
inspect.Parameter.VAR_KEYWORD,
|
|
149
|
+
):
|
|
150
|
+
continue
|
|
151
|
+
annotation = hints.get(pname, typing.Any)
|
|
152
|
+
default = ... if param.default is inspect.Parameter.empty else param.default
|
|
153
|
+
fields[pname] = (annotation, default)
|
|
154
|
+
|
|
155
|
+
if not fields:
|
|
156
|
+
return None
|
|
157
|
+
return create_model(f"{tool_name}Args", **fields)
|
|
158
|
+
|
|
159
|
+
|
|
160
|
+
def _as_content(result: typing.Any) -> str:
|
|
161
|
+
"""Convert a tool result to a string for LangChain's ToolMessage."""
|
|
162
|
+
if isinstance(result, str):
|
|
163
|
+
return result
|
|
164
|
+
try:
|
|
165
|
+
return json.dumps(result, default=str)
|
|
166
|
+
except (TypeError, ValueError):
|
|
167
|
+
return str(result)
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
def _coerce_tool(t: typing.Any) -> typing.Any:
|
|
171
|
+
"""Coerce a tool to a LangChain-compatible tool."""
|
|
172
|
+
if isinstance(t, AsyncFunctionTaskTemplate):
|
|
173
|
+
return tool(t)
|
|
174
|
+
return t
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: flyteplugins-agents-langchain
|
|
3
|
+
Version: 2.5.15
|
|
4
|
+
Summary: Run LangChain 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: langchain
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
flyteplugins/agents/langchain/__init__.py,sha256=OedlCjF8pH0Q36z04eVKpaIDVeezFvee3WNdWFkjCIg,773
|
|
2
|
+
flyteplugins/agents/langchain/_durable.py,sha256=grIsJkTLfTVM-m-AB8cBWQoNV3llDgEpdvtwAcBKYHI,5953
|
|
3
|
+
flyteplugins/agents/langchain/_memory.py,sha256=jaeO5CPUKOh4ii2Q5pBRIPMNFyiDvKI96Em0GViKGtI,2874
|
|
4
|
+
flyteplugins/agents/langchain/_run.py,sha256=BMkupLSV4yYQYl5jiW0MLDqc1b6xJmgUno8RSL3x78o,7152
|
|
5
|
+
flyteplugins/agents/langchain/_tools.py,sha256=HSXsjLxXzDIx6sw81rBy_KjBps6lH_Q6PHRENa6q-Hc,6584
|
|
6
|
+
flyteplugins_agents_langchain-2.5.15.dist-info/METADATA,sha256=Sq0MS7bSkdiJ_8mcd1gEzcLerTnrVo8vvXG8TbEqaJs,289
|
|
7
|
+
flyteplugins_agents_langchain-2.5.15.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
|
|
8
|
+
flyteplugins_agents_langchain-2.5.15.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
|
|
9
|
+
flyteplugins_agents_langchain-2.5.15.dist-info/RECORD,,
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
flyteplugins
|