fix: remove loaded object via main

Author: cosmyo

packages/uipath-openai-agents/src/uipath_openai_agents/runtime/agent.py

@@ -28,9 +28,6 @@ def __init__(self, name: str, file_path: str, variable_name: str):
         self.file_path = file_path
         self.variable_name = variable_name
         self._context_manager: Any = None
-        self._loaded_object: Any = (
-            None  # Store original loaded object for type inference
-        )
 
     @classmethod
     def from_path_string(cls, name: str, file_path: str) -> Self:
@@ -102,9 +99,6 @@ async def load(self) -> Agent:
                 category=UiPathErrorCategory.USER,
             )
 
-        # Store the original loaded object for type inference
-        self._loaded_object = agent_object
-
         agent = await self._resolve_agent(agent_object)
         if not isinstance(agent, Agent):
             raise UiPathOpenAIAgentsRuntimeError(
@@ -179,17 +173,6 @@ async def _resolve_agent(self, agent_object: Any) -> Agent:
 
         return agent_instance
 
-    def get_loaded_object(self) -> Any:
-        """
-        Get the original loaded object before agent resolution.
-
-        This is useful for extracting type annotations from wrapper functions.
-
-        Returns:
-            The original loaded object (could be an Agent, function, or callable)
-        """
-        return self._loaded_object
-
     async def cleanup(self) -> None:
         """Clean up resources (e.g., exit async context managers)."""
         if self._context_manager:

packages/uipath-openai-agents/src/uipath_openai_agents/runtime/factory.py

@@ -303,17 +303,11 @@ async def _create_runtime_instance(
         storage = await self._get_or_create_storage()
         storage_path = storage.storage_path if storage else None
 
-        # Get the loaded object from the agent loader for schema inference
-        loaded_object = None
-        if entrypoint in self._agent_loaders:
-            loaded_object = self._agent_loaders[entrypoint].get_loaded_object()
-
         return UiPathOpenAIAgentRuntime(
             agent=agent,
             runtime_id=runtime_id,
             entrypoint=entrypoint,
             storage_path=storage_path,
-            loaded_object=loaded_object,
             storage=storage,
         )
 

packages/uipath-openai-agents/src/uipath_openai_agents/runtime/runtime.py

@@ -41,7 +41,6 @@ def __init__(
         runtime_id: str | None = None,
         entrypoint: str | None = None,
         storage_path: str | None = None,
-        loaded_object: Any | None = None,
         storage: SqliteAgentStorage | None = None,
     ):
         """
@@ -52,14 +51,12 @@ def __init__(
             runtime_id: Unique identifier for this runtime instance
             entrypoint: Optional entrypoint name (for schema generation)
             storage_path: Path to SQLite database for session persistence
-            loaded_object: Original loaded object (for schema inference)
             storage: Optional storage instance for state persistence
         """
         self.agent: Agent = agent
         self.runtime_id: str = runtime_id or "default"
         self.entrypoint: str | None = entrypoint
         self.storage_path: str | None = storage_path
-        self.loaded_object: Any | None = loaded_object
         self.storage: SqliteAgentStorage | None = storage
 
         # Configure OpenAI Agents SDK to use Responses API
@@ -477,7 +474,7 @@ async def get_schema(self) -> UiPathRuntimeSchema:
         Returns:
             UiPathRuntimeSchema with input/output schemas and graph structure
         """
-        entrypoints_schema = get_entrypoints_schema(self.agent, self.loaded_object)
+        entrypoints_schema = get_entrypoints_schema(self.agent)
 
         return UiPathRuntimeSchema(
             filePath=self.entrypoint,

packages/uipath-openai-agents/src/uipath_openai_agents/runtime/schema.py

@@ -1,7 +1,7 @@
 """Schema extraction utilities for OpenAI Agents."""
 
 import inspect
-from typing import Any, get_args, get_origin, get_type_hints
+from typing import Any, get_args, get_origin
 
 from agents import Agent
 from pydantic import BaseModel, TypeAdapter
@@ -41,106 +41,14 @@ def _is_pydantic_model(type_hint: Any) -> bool:
     return False
 
 
-def _extract_schema_from_callable(callable_obj: Any) -> dict[str, Any] | None:
-    """
-    Extract input/output schemas from a callable's type annotations.
-
-    Args:
-        callable_obj: A callable object (function, async function, etc.)
-
-    Returns:
-        Dictionary with input and output schemas if type hints are found,
-        None otherwise
-    """
-    if not callable(callable_obj):
-        return None
-
-    try:
-        # Get type hints from the callable
-        type_hints = get_type_hints(callable_obj)
-
-        if not type_hints:
-            return None
-
-        # Get function signature to identify parameters
-        sig = inspect.signature(callable_obj)
-        params = list(sig.parameters.values())
-
-        # Find the first parameter (usually the input)
-        input_type = None
-
-        for param in params:
-            if param.name in ("self", "cls"):
-                continue
-            if param.name in type_hints:
-                input_type = type_hints[param.name]
-                break
-
-        # Get return type
-        return_type = type_hints.get("return")
-
-        schema: dict[str, Any] = {
-            "input": {"type": "object", "properties": {}, "required": []},
-            "output": {"type": "object", "properties": {}, "required": []},
-        }
-
-        # Extract input schema from Pydantic model
-        if input_type and _is_pydantic_model(input_type):
-            adapter = TypeAdapter(input_type)
-            input_schema = adapter.json_schema()
-            unpacked_input = _resolve_refs(input_schema)
-
-            schema["input"]["properties"] = _process_nullable_types(
-                unpacked_input.get("properties", {})
-            )
-            schema["input"]["required"] = unpacked_input.get("required", [])
-
-            # Add title and description if available
-            if "title" in unpacked_input:
-                schema["input"]["title"] = unpacked_input["title"]
-            if "description" in unpacked_input:
-                schema["input"]["description"] = unpacked_input["description"]
-
-        # Extract output schema from Pydantic model
-        if return_type and _is_pydantic_model(return_type):
-            adapter = TypeAdapter(return_type)
-            output_schema = adapter.json_schema()
-            unpacked_output = _resolve_refs(output_schema)
-
-            schema["output"]["properties"] = _process_nullable_types(
-                unpacked_output.get("properties", {})
-            )
-            schema["output"]["required"] = unpacked_output.get("required", [])
-
-            # Add title and description if available
-            if "title" in unpacked_output:
-                schema["output"]["title"] = unpacked_output["title"]
-            if "description" in unpacked_output:
-                schema["output"]["description"] = unpacked_output["description"]
-
-        # Only return schema if we found at least one Pydantic model
-        if schema["input"]["properties"] or schema["output"]["properties"]:
-            return schema
-
-    except Exception:
-        # If schema extraction fails, return None to fall back to default
-        pass
-
-    return None
-
-
-def get_entrypoints_schema(
-    agent: Agent, loaded_object: Any | None = None
-) -> dict[str, Any]:
+def get_entrypoints_schema(agent: Agent) -> dict[str, Any]:
     """
     Extract input/output schema from an OpenAI Agent.
 
-    Prioritizes the agent's native output_type attribute (OpenAI Agents pattern),
-    with optional fallback to wrapper function type hints (UiPath pattern).
+    Uses the agent's native output_type attribute for schema extraction.
 
     Args:
         agent: An OpenAI Agent instance
-        loaded_object: Optional original loaded object (function/callable) with type annotations
 
     Returns:
         Dictionary with input and output schemas
@@ -170,19 +78,19 @@ def get_entrypoints_schema(
         "required": ["message"],
     }
 
-    # Extract output schema - PRIORITY 1: Agent's output_type (native OpenAI Agents pattern)
+    # Extract output schema - Agent's output_type (native OpenAI Agents pattern)
     output_type = getattr(agent, "output_type", None)
     output_extracted = False
 
     # Unwrap AgentOutputSchema if present (OpenAI Agents SDK wrapper)
-    # Check for AgentOutputSchema by looking for 'schema' attribute on non-type instances
+    # AgentOutputSchema wraps the actual Pydantic model in an 'output_type' attribute
     if (
         output_type is not None
-        and hasattr(output_type, "schema")
+        and hasattr(output_type, "output_type")
         and not isinstance(output_type, type)
     ):
         # This is an AgentOutputSchema wrapper instance, extract the actual model
-        output_type = output_type.schema
+        output_type = output_type.output_type
 
     if output_type is not None and _is_pydantic_model(output_type):
         try:
@@ -207,15 +115,6 @@ def get_entrypoints_schema(
             # Continue to fallback if extraction fails
             pass
 
-    # Extract output schema - PRIORITY 2: Wrapper function type hints (UiPath pattern)
-    # This allows UiPath-specific patterns where agents are wrapped in typed functions
-    if not output_extracted and loaded_object is not None:
-        wrapper_schema = _extract_schema_from_callable(loaded_object)
-        if wrapper_schema is not None:
-            # Use the wrapper's output schema, but keep the default input (messages)
-            schema["output"] = wrapper_schema["output"]
-            output_extracted = True
-
     # Fallback: Default output schema for agents without explicit output_type
     if not output_extracted:
         schema["output"] = {

packages/uipath-openai-agents/tests/test_agent_as_tools_schema.py

@@ -10,7 +10,6 @@
 from main import (  # type: ignore  # noqa: E402
     TranslationInput,
     TranslationOutput,
-    main,
     orchestrator_agent,
 )
 
@@ -19,7 +18,7 @@
 
 def test_agent_as_tools_input_schema():
     """Test that input schema uses default messages format (OpenAI Agents pattern)."""
-    schema = get_entrypoints_schema(orchestrator_agent, main)
+    schema = get_entrypoints_schema(orchestrator_agent)
 
     # Verify input schema structure - should use default messages
     assert "input" in schema
@@ -42,7 +41,7 @@ def test_agent_as_tools_input_schema():
 
 def test_agent_as_tools_output_schema():
     """Test that output schema is extracted from agent's output_type."""
-    schema = get_entrypoints_schema(orchestrator_agent, main)
+    schema = get_entrypoints_schema(orchestrator_agent)
 
     # Verify output schema structure
     assert "output" in schema
@@ -73,7 +72,7 @@ def test_agent_as_tools_output_schema():
 
 def test_agent_as_tools_schema_metadata():
     """Test that schema includes model metadata from agent's output_type."""
-    schema = get_entrypoints_schema(orchestrator_agent, main)
+    schema = get_entrypoints_schema(orchestrator_agent)
 
     # Input uses default messages format (no custom title/description)
     assert "input" in schema

packages/uipath-openai-agents/tests/test_integration.py

@@ -12,7 +12,6 @@
 from main import (  # type: ignore  # noqa: E402
     TranslationInput,
     TranslationOutput,
-    main,
     orchestrator_agent,
 )
 
@@ -52,7 +51,7 @@ def test_error_handling():
 
 def test_schema_extraction_with_new_serialization():
     """Test that schema extraction works with the serialization improvements."""
-    schema = get_entrypoints_schema(orchestrator_agent, main)
+    schema = get_entrypoints_schema(orchestrator_agent)
 
     # Verify input schema (messages format)
     assert "input" in schema
@@ -86,13 +85,11 @@ async def test_runtime_initialization_with_storage():
             entrypoint="test",
             storage_path=storage_path,
             storage=storage,
-            loaded_object=main,
         )
 
         # Verify runtime initialized correctly
         assert runtime.storage is not None
         assert runtime.runtime_id == "test_runtime"
-        assert runtime.loaded_object == main
 
         # Test schema generation
         schema = await runtime.get_schema()

packages/uipath-openai-agents/tests/test_schema_inference.py

@@ -29,8 +29,8 @@ class OutputModel(BaseModel):
 
 
 def test_schema_inference_from_agent_output_type():
-    """Test that output schema is correctly inferred from agent's output_type (PRIMARY)."""
-    schema = get_entrypoints_schema(agent_with_output_type, None)
+    """Test that output schema is correctly inferred from agent's output_type."""
+    schema = get_entrypoints_schema(agent_with_output_type)
 
     # Check input schema - should be default messages format
     assert "input" in schema
@@ -58,7 +58,7 @@ def test_schema_inference_from_agent_output_type():
 
 def test_schema_fallback_without_types():
     """Test that schemas fall back to defaults when no types are provided."""
-    schema = get_entrypoints_schema(test_agent, None)
+    schema = get_entrypoints_schema(test_agent)
 
     # Should use default message-based input schema
     assert "input" in schema
@@ -70,8 +70,8 @@ def test_schema_fallback_without_types():
 
 
 def test_schema_with_plain_agent():
-    """Test schema extraction with a plain agent (no wrapper function)."""
-    schema = get_entrypoints_schema(test_agent, test_agent)
+    """Test schema extraction with a plain agent."""
+    schema = get_entrypoints_schema(test_agent)
 
     # Should use default message input
     assert "input" in schema
@@ -80,33 +80,3 @@ def test_schema_with_plain_agent():
     # Should use default result output
     assert "output" in schema
     assert "result" in schema["output"]["properties"]
-
-
-class WrapperOutputModel(BaseModel):
-    """Output model for wrapper function test."""
-
-    status: str
-    data: dict[str, str]
-
-
-async def typed_wrapper_function(message: str) -> WrapperOutputModel:
-    """Wrapper function with type annotations (UiPath pattern - SECONDARY)."""
-    return WrapperOutputModel(status="success", data={})
-
-
-def test_schema_with_wrapper_function():
-    """Test that wrapper function output schema is used as fallback (SECONDARY)."""
-    # Agent without output_type should fallback to wrapper function
-    schema = get_entrypoints_schema(test_agent, typed_wrapper_function)
-
-    # Input should still be default messages (not extracted from wrapper)
-    assert "input" in schema
-    assert "message" in schema["input"]["properties"]
-    assert "required" in schema["input"]
-    assert "message" in schema["input"]["required"]
-
-    # Output should come from wrapper function (secondary pattern)
-    assert "output" in schema
-    assert "properties" in schema["output"]
-    assert "status" in schema["output"]["properties"]
-    assert "data" in schema["output"]["properties"]