Skip to content

Commit e924724

Browse files
authored
Fix openai @function_tool tools crashing under default spawn start method- #451
2 parents 8d8fce3 + dd941af commit e924724

4 files changed

Lines changed: 198 additions & 61 deletions

File tree

‎src/conductor/ai/agents/frameworks/serializer.py‎

Lines changed: 2 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
from dataclasses import fields as dc_fields
1616
from typing import Any, Callable, Dict, List, Optional, Set, Tuple
1717

18+
from conductor.ai.agents.runtime._worker_entries import _find_embedded_function
19+
1820
logger = logging.getLogger("conductor.ai.agents.frameworks")
1921

2022

@@ -340,64 +342,6 @@ def _try_extract_tool_object(obj: Any) -> Optional[WorkerInfo]:
340342
)
341343

342344

343-
def _find_embedded_function(obj: Any, max_depth: int = 2) -> Optional[Any]:
344-
"""Walk an object's attributes to find an embedded plain function.
345-
346-
Searches closures and nested attribute objects for the original
347-
user-defined function. Returns ``None`` if not found.
348-
"""
349-
if max_depth <= 0:
350-
return None
351-
352-
# Check direct attributes for callables that look like user functions
353-
for attr_name in vars(obj) if hasattr(obj, "__dict__") else []:
354-
val = getattr(obj, attr_name, None)
355-
if val is None:
356-
continue
357-
358-
# Plain function with a clean signature
359-
if inspect.isfunction(val):
360-
# Check if it has a closure containing the original function
361-
func = _extract_from_closure(val)
362-
if func is not None:
363-
return func
364-
# The function itself might be usable
365-
return val
366-
367-
# Nested object — recurse one level
368-
if hasattr(val, "__dict__") and not isinstance(val, type):
369-
result = _find_embedded_function(val, max_depth - 1)
370-
if result is not None:
371-
return result
372-
373-
return None
374-
375-
376-
def _extract_from_closure(func: Any) -> Optional[Any]:
377-
"""Extract the original user function from a closure's cell variables."""
378-
closure = getattr(func, "__closure__", None)
379-
if not closure:
380-
return None
381-
382-
for cell in closure:
383-
try:
384-
val = cell.cell_contents
385-
except ValueError:
386-
continue
387-
if inspect.isfunction(val):
388-
# Skip internal wrappers that take (ctx, input) or (context, ...)
389-
try:
390-
sig = inspect.signature(val)
391-
param_names = list(sig.parameters.keys())
392-
# Internal wrappers typically start with ctx/context as first param
393-
if param_names and param_names[0] in ("ctx", "context"):
394-
continue
395-
return val
396-
except (ValueError, TypeError):
397-
continue
398-
return None
399-
400-
401345
def _extract_callable(func: Any) -> WorkerInfo:
402346
"""Extract name, description, and JSON schema from a callable."""
403347
from conductor.ai.agents._internal.schema_utils import schema_from_function

‎src/conductor/ai/agents/runtime/_worker_entries.py‎

Lines changed: 108 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@
3535
import pickle
3636
import sys
3737
from dataclasses import dataclass
38-
from typing import Callable, Dict, FrozenSet, Optional
38+
from typing import Any, Callable, Dict, FrozenSet, Optional
3939

4040
# Mirrors conductor.client.worker.worker._MAX_UNWRAP_DEPTH.
4141
_MAX_UNWRAP_DEPTH = 32
@@ -67,6 +67,79 @@ def _walk_qualname(module_obj, qualname: str):
6767
# .coroutine; our Guardrail / ToolDef → .func.
6868
_CONTAINER_ATTRS = ("func", "coroutine")
6969

70+
def _extract_from_closure(func: Callable) -> Optional[Callable]:
71+
"""Extract the original user function from a closure's cell variables.
72+
73+
- Shared by :func:`_find_embedded_function` below.
74+
- In turn shared by :mod:`conductor.ai.agents.frameworks.serializer`'s
75+
discovery and :class:`FunctionRef`'s parent-verification /
76+
child-reconstruction — one implementation, can't drift apart.
77+
"""
78+
closure = getattr(func, "__closure__", None)
79+
if not closure:
80+
return None
81+
82+
for cell in closure:
83+
try:
84+
val = cell.cell_contents
85+
except ValueError:
86+
continue
87+
if inspect.isfunction(val):
88+
# Skip internal wrappers that take (ctx, input) or (context, ...)
89+
try:
90+
sig = inspect.signature(val)
91+
param_names = list(sig.parameters.keys())
92+
# Internal wrappers typically start with ctx/context as first param
93+
if param_names and param_names[0] in ("ctx", "context"):
94+
continue
95+
return val
96+
except (ValueError, TypeError):
97+
continue
98+
return None
99+
100+
101+
# How many attribute-nesting levels _find_embedded_function will descend.
102+
# openai-agents' FunctionTool -> on_invoke_tool (invoker instance) ->
103+
# _invoke_tool_impl (closure fn) is 2 levels deep; matches serializer.py's
104+
# own default so parent discovery and child reconstruction stay in lockstep.
105+
_DEEP_EXTRACT_MAX_DEPTH = 2
106+
107+
108+
def _find_embedded_function(obj: Any, max_depth: int = _DEEP_EXTRACT_MAX_DEPTH) -> Optional[Callable]:
109+
"""Walk an object's attributes to find an embedded plain function.
110+
111+
- Generic, framework-agnostic: no hardcoded attribute names.
112+
- For containers that expose the original function only behind further
113+
nested attributes and/or closures, not a single named attribute — e.g.
114+
openai-agents' ``FunctionTool.on_invoke_tool`` is itself a callable
115+
*instance*, not a function; the real closure lives one level deeper,
116+
on that instance's own ``_invoke_tool_impl`` attribute.
117+
- Deliberately avoids hardcoding that private attribute name: a
118+
shape-based walk degrades to "not found" (the pre-existing, actionable
119+
``SpawnSafetyError``) instead of breaking outright if openai-agents
120+
restructures its internals.
121+
"""
122+
if max_depth <= 0:
123+
return None
124+
125+
for attr_name in vars(obj) if hasattr(obj, "__dict__") else []:
126+
val = getattr(obj, attr_name, None)
127+
if val is None:
128+
continue
129+
130+
if inspect.isfunction(val):
131+
func = _extract_from_closure(val)
132+
if func is not None:
133+
return func
134+
return val
135+
136+
if hasattr(val, "__dict__") and not isinstance(val, type):
137+
result = _find_embedded_function(val, max_depth - 1)
138+
if result is not None:
139+
return result
140+
141+
return None
142+
70143

71144
def _wrapped_depth_to(obj, fn) -> Optional[int]:
72145
"""Number of ``__wrapped__`` hops from *obj* down to *fn*, or ``None``."""
@@ -97,12 +170,25 @@ class FunctionRef:
97170
object rather than a wraps-wrapper — e.g. langchain's ``@tool`` rebinds
98171
it to a ``StructuredTool`` holding the original in ``.func`` (sync) or
99172
``.coroutine`` (async). The hop is taken before the ``__wrapped__`` walk.
173+
174+
``deep_extract`` handles containers that don't expose the original
175+
function via any single named attribute:
176+
177+
- e.g. openai-agents' ``FunctionTool``, whose ``on_invoke_tool`` is
178+
itself a callable instance, not a function — the real function is
179+
nested further behind that instance's own attributes and a closure.
180+
- Uses the generic :func:`_find_embedded_function` walk instead of a
181+
fixed ``attr_hop``, so it isn't tied to a specific (private,
182+
renameable) attribute path.
183+
- Tried last, after ``attr_hop``/``_CONTAINER_ATTRS``, before the
184+
``__wrapped__`` walk.
100185
"""
101186

102187
module: str
103188
qualname: str
104189
unwrap_depth: int = 0
105190
attr_hop: str = ""
191+
deep_extract: bool = False
106192

107193
@classmethod
108194
def of(cls, fn: Callable) -> "FunctionRef":
@@ -158,10 +244,22 @@ def of(cls, fn: Callable) -> "FunctionRef":
158244
depth = _wrapped_depth_to(inner, fn)
159245
if depth is not None:
160246
return cls(module, qualname, depth, attr)
247+
# Still no match — container may hold the original behind further
248+
# nested attributes and/or a closure, not a single named attribute.
249+
# e.g. openai-agents' FunctionTool.on_invoke_tool: callable instance,
250+
# not a function; real closure is one level deeper.
251+
# Generic walk, no hardcoded attribute names (_find_embedded_function).
252+
# No __wrapped__ fallback needed here (unlike _CONTAINER_ATTRS): fn
253+
# always arrives as _find_embedded_function's own prior return value
254+
# (how serializer.py derives WorkerInfo.func) — identical walk always
255+
# matches by identity.
256+
if _find_embedded_function(obj) is fn:
257+
return cls(module, qualname, 0, "", deep_extract=True)
161258
raise SpawnSafetyError(
162259
f"'{module}.{qualname}' does not resolve back to {fn!r} (rebound "
163-
f"without a __wrapped__ chain or a func/coroutine container "
164-
f"attribute). {_REMEDIES}"
260+
f"without a __wrapped__ chain, a func/coroutine container "
261+
f"attribute, or a discoverable nested/closure-held function). "
262+
f"{_REMEDIES}"
165263
)
166264

167265
def resolve(self) -> Callable:
@@ -173,6 +271,13 @@ def resolve(self) -> Callable:
173271
obj = _walk_qualname(module_obj, self.qualname)
174272
if self.attr_hop:
175273
obj = getattr(obj, self.attr_hop)
274+
if self.deep_extract:
275+
obj = _find_embedded_function(obj)
276+
if obj is None:
277+
raise SpawnSafetyError(
278+
f"'{self.module}.{self.qualname}' no longer yields a "
279+
f"discoverable nested function (definition changed?)."
280+
)
176281
for _ in range(self.unwrap_depth):
177282
obj = obj.__wrapped__
178283
_RESOLVE_CACHE[self] = obj

‎tests/unit/ai/test_worker_entries.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
FunctionRef,
1616
SpawnSafetyError,
1717
ToolWorkerEntry,
18+
_find_embedded_function,
1819
probe_spawn_safety,
1920
)
2021
from conductor.ai.agents.tool import get_tool_def
@@ -160,6 +161,72 @@ def test_cross_process_spawn_roundtrip_with_hop(self):
160161
assert p.exitcode == 0
161162

162163

164+
# ── Deep-extract (real openai-agents @function_tool → FunctionTool) ────────
165+
166+
167+
class TestFunctionRefDeepExtract:
168+
"""openai-agents' real ``@function_tool`` → ``FunctionTool`` with no
169+
``.func``/``.coroutine``/``__wrapped__``; original fn only reachable via
170+
nested attrs + closure."""
171+
172+
def test_sync_function_tool_deep_extract(self):
173+
pytest.importorskip("agents")
174+
from tests.unit.resources import openai_agents_entry_helpers as oa
175+
176+
raw = _find_embedded_function(oa.oa_get_weather)
177+
assert raw is not None
178+
ref = FunctionRef.of(raw)
179+
assert ref == FunctionRef(oa.__name__, "oa_get_weather", 0, "", deep_extract=True)
180+
assert ref.resolve() is raw
181+
182+
def test_async_function_tool_deep_extract(self):
183+
pytest.importorskip("agents")
184+
from tests.unit.resources import openai_agents_entry_helpers as oa
185+
186+
raw = _find_embedded_function(oa.oa_get_weather_async)
187+
assert raw is not None
188+
ref = FunctionRef.of(raw)
189+
assert ref.deep_extract is True
190+
assert ref.attr_hop == ""
191+
assert ref.resolve() is raw
192+
193+
def test_ref_pickles(self):
194+
pytest.importorskip("agents")
195+
from tests.unit.resources import openai_agents_entry_helpers as oa
196+
197+
raw = _find_embedded_function(oa.oa_get_weather)
198+
ref = pickle.loads(pickle.dumps(FunctionRef.of(raw)))
199+
assert ref.resolve() is raw
200+
201+
def test_entry_transports_function_tool_fn_by_ref(self):
202+
# Pre-fix this fell to fn_direct, whose reference pickling then found
203+
# the FunctionTool at the global name: "it's not the same object as …".
204+
pytest.importorskip("agents")
205+
from tests.unit.resources import openai_agents_entry_helpers as oa
206+
207+
raw = _find_embedded_function(oa.oa_get_weather)
208+
entry = ToolWorkerEntry.for_callable(raw, "oa_get_weather")
209+
assert entry.fn_ref is not None
210+
clone = pickle.loads(pickle.dumps(entry))
211+
assert clone._target() is raw
212+
213+
def test_cross_process_spawn_roundtrip(self):
214+
pytest.importorskip("agents")
215+
from tests.unit.resources import openai_agents_entry_helpers as oa
216+
217+
raw = _find_embedded_function(oa.oa_get_weather)
218+
ctx = multiprocessing.get_context("spawn")
219+
q = ctx.Queue()
220+
ref_bytes = pickle.dumps(FunctionRef.of(raw))
221+
p = ctx.Process(target=helpers.resolve_and_call_child, args=(ref_bytes, "Boston", q))
222+
p.start()
223+
try:
224+
assert q.get(timeout=30) == "sunny in Boston"
225+
finally:
226+
p.join(timeout=30)
227+
assert p.exitcode == 0
228+
229+
163230
# ── Guardrail spawn transport ─────────────────────────────────────────────
164231

165232

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
"""
2+
Module-level real openai-agents ``@function_tool`` subjects for FunctionRef
3+
deep-extract tests.
4+
5+
Separate from worker_entry_helpers.py so environments without the
6+
openai-agents extra can still import that module; tests importing THIS
7+
module must be gated with ``pytest.importorskip("agents")``.
8+
"""
9+
from agents import function_tool
10+
11+
12+
@function_tool
13+
def oa_get_weather(city: str) -> str:
14+
"""Return a canned weather string for a city."""
15+
return f"sunny in {city}"
16+
17+
18+
@function_tool
19+
async def oa_get_weather_async(city: str) -> str:
20+
"""Return a canned weather string for a city, asynchronously."""
21+
return f"async sunny in {city}"

0 commit comments

Comments
 (0)