flyteplugins-openai 2.0.0__py3-none-any.whl

flyteplugins/openai/agents/__init__.py

@@ -0,0 +1,3 @@
+from ._function_tools import function_tool
+
+__all__ = ["function_tool"]

flyteplugins/openai/agents/_function_tools.py

@@ -0,0 +1,85 @@
+import inspect
+import json
+import typing
+from dataclasses import asdict, dataclass
+from functools import partial
+
+import agents
+from agents import FunctionTool as OpenAIFunctionTool
+from agents import function_tool as openai_function_tool
+from agents.function_schema import function_schema
+from agents.tool_context import ToolContext
+from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
+from flyte.models import NativeInterface
+from packaging import version
+
+MIN_PACKAGE_VERSION = "0.2.4"
+assert version.parse(agents.__version__) >= version.parse(MIN_PACKAGE_VERSION), (
+    f"The agents package needs to be at least version {MIN_PACKAGE_VERSION}, found version {agents.__version__}"
+)
+
+
+@dataclass
+class FunctionTool(OpenAIFunctionTool):
+    """
+    Flyte-compatible replacement for agents.FunctionTool
+
+    This is a dataclass that includes additional fields that are not present in
+    the OpenAI FunctionTool dataclass so that the tool can be used as a flyte
+    task.
+    """
+
+    task: TaskTemplate | None = None
+    native_interface: NativeInterface | None = None
+    report: bool = False
+
+    async def execute(self, *args, **kwargs):
+        return await self.task.execute(*args, **kwargs)
+
+
+def function_tool(
+    func: AsyncFunctionTaskTemplate | typing.Callable | None = None,
+    **kwargs,
+) -> FunctionTool | OpenAIFunctionTool:
+    """Flyte-compatible replacement for @agents.function_tool
+
+    **kwargs are forwarded to the underlying @agents.function_tool decorator.
+    For @flyte.trace functions, this just forwards all the arguments to the
+    agents.function_tool decorator:
+    https://openai.github.io/openai-agents-python/ref/tool/#agents.tool.function_tool
+
+    For @TaskEnvironment.task functions, this will create a flyte-compatible
+    FunctionTool dataclass that can run tools as flyte tasks.
+    """
+    if func is None:
+        return partial(function_tool, **kwargs)
+
+    if isinstance(func, AsyncFunctionTaskTemplate):
+
+        async def _on_invoke_tool(ctx: ToolContext[typing.Any], input: str) -> typing.Any:
+            json_data: dict[str, typing.Any] = json.loads(input) if input else {}
+            schema = function_schema(func.func)
+            parsed = schema.params_pydantic_model(**json_data) if json_data else schema.params_pydantic_model()
+            args, kwargs_dict = schema.to_call_args(parsed)
+            if inspect.iscoroutinefunction(func.func):
+                out = await func(*args, **kwargs_dict)
+            else:
+                out = func(*args, **kwargs_dict)
+            return out
+
+        _openai_fn_tool = asdict(openai_function_tool(func.func, **kwargs))
+        _openai_fn_tool.pop("on_invoke_tool")
+        _openai_fn_tool.pop("_is_agent_tool")
+        _openai_fn_tool.pop("_is_codex_tool")
+        _openai_fn_tool.pop("_agent_instance")
+        return FunctionTool(
+            native_interface=func.native_interface,
+            task=func,
+            report=func.report,
+            on_invoke_tool=_on_invoke_tool,
+            **_openai_fn_tool,
+        )
+
+    # regular callables or flyte.trace-decorated functions should use the
+    # openai-agents function_tool decorator
+    return openai_function_tool(func, **kwargs)

flyteplugins_openai-2.0.0.dist-info/METADATA

@@ -0,0 +1,21 @@
+Metadata-Version: 2.4
+Name: flyteplugins-openai
+Version: 2.0.0
+Summary: OpenAI plugin for flyte
+Author-email: Niels Bantilan <cosmicbboy@users.noreply.github.com>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: openai-agents>=0.2.4
+Requires-Dist: flyte
+
+# Flyte OpenAI Plugin
+
+This plugin provides a drop-in replacement for OpenAI packages. It provides
+drop-in replacements for functionality in the `openai-agents` package so that
+it works on Flyte.
+
+To install the plugin, run the following command:
+
+```bash
+pip install --pre flyteplugins-openai
+```

flyteplugins_openai-2.0.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+flyteplugins/openai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flyteplugins/openai/agents/__init__.py,sha256=fXfkrE1cbLBPjcytMOHyBexMj4_YMyZumhBC4QTzdcc,72
+flyteplugins/openai/agents/_function_tools.py,sha256=9MNZYta168MwiiPKLhW9JzZLNr2wM8d-rFE8E1Qm_yU,3207
+flyteplugins_openai-2.0.0.dist-info/METADATA,sha256=1J6IdXMkDNra01WfkB0kCJKpLsvqSCbs1GfyIYZiaQE,583
+flyteplugins_openai-2.0.0.dist-info/WHEEL,sha256=YCfwYGOYMi5Jhw2fU4yNgwErybb2IX5PEwBKV4ZbdBo,91
+flyteplugins_openai-2.0.0.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_openai-2.0.0.dist-info/RECORD,,

flyteplugins_openai-2.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flyteplugins_openai-2.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+flyteplugins