Merge pull request #972 from UiPath/akshaya/use_io_schema

fix(UseOutputSchema): use custom outputschema for mocking

Author: akshaylive

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath"
-version = "2.2.10"
+version = "2.2.11"
 description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools."
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath/_cli/_evals/mocks/llm_mocker.py

@@ -4,7 +4,7 @@
 import logging
 from typing import Any, Callable
 
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 
 from uipath.tracing import traced
 from uipath.tracing._utils import _SpanUtils
@@ -105,15 +105,16 @@ async def response(
             if return_type is None:
                 return_type = Any
 
-            class OutputSchema(BaseModel):
-                response: return_type
+            output_schema = params.get(
+                "output_schema", TypeAdapter(return_type).json_schema()
+            )
 
             response_format = {
                 "type": "json_schema",
                 "json_schema": {
-                    "name": OutputSchema.__name__.lower(),
+                    "name": "OutputSchema",
                     "strict": True,
-                    "schema": _cleanup_schema(OutputSchema),
+                    "schema": _cleanup_schema(output_schema),
                 },
             }
             try:
@@ -199,10 +200,7 @@ class OutputSchema(BaseModel):
                     response_format=response_format,
                     **completion_kwargs,
                 )
-                mocked_response = OutputSchema(
-                    **json.loads(response.choices[0].message.content)
-                )
-                result = mocked_response.model_dump(mode="json")["response"]
+                result = json.loads(response.choices[0].message.content)
 
                 if cache_manager is not None:
                     cache_manager.set(

src/uipath/platform/llm_gateway/_llm_gateway_service.py

@@ -77,15 +77,15 @@ class EmbeddingModels(object):
     text_embedding_ada_002 = "text-embedding-ada-002"
 
 
-def _cleanup_schema(model_class: type[BaseModel]) -> dict[str, Any]:
-    """Clean up a Pydantic model schema for use with LLM Gateway.
+def _cleanup_schema(schema: dict[str, Any]) -> dict[str, Any]:
+    """Clean up a JSON schema for use with LLM Gateway.
 
-    This function converts a Pydantic model's JSON schema to a format that's
+    This function converts a JSON schema to a format that's
     compatible with the LLM Gateway's JSON schema requirements by removing
     titles and other metadata that might cause validation issues.
 
     Args:
-        model_class (type[BaseModel]): A Pydantic BaseModel class to convert to schema.
+        schema (dict[str, Any]): an input JSON schema.
 
     Returns:
         dict: A cleaned JSON schema dictionary suitable for LLM Gateway response_format.
@@ -99,11 +99,10 @@ class Country(BaseModel):
             capital: str
             languages: list[str]
 
-        schema = _cleanup_schema(Country)
+        schema = _cleanup_schema(Country.model_json_schema())
         # Returns a clean schema without titles and unnecessary metadata
         ```
     """
-    schema = model_class.model_json_schema()
 
     def clean_type(type_def):
         """Clean property definitions by removing titles and cleaning nested items. Additionally, `additionalProperties` is ensured on all objects."""
@@ -303,7 +302,7 @@ class Country(BaseModel):
                 response_format, BaseModel
             ):
                 # Convert Pydantic model to JSON schema format
-                cleaned_schema = _cleanup_schema(response_format)
+                cleaned_schema = _cleanup_schema(response_format.model_json_schema())
                 request_body["response_format"] = {
                     "type": "json_schema",
                     "json_schema": {
@@ -511,7 +510,7 @@ class Country(BaseModel):
                 response_format, BaseModel
             ):
                 # Convert Pydantic model to JSON schema format
-                cleaned_schema = _cleanup_schema(response_format)
+                cleaned_schema = _cleanup_schema(response_format.model_json_schema())
                 request_body["response_format"] = {
                     "type": "json_schema",
                     "json_schema": {

tests/cli/eval/mocks/test_mocks.py

@@ -1,8 +1,10 @@
+import json
 from typing import Any
 from unittest.mock import MagicMock
 
 import pytest
 from _pytest.monkeypatch import MonkeyPatch
+from pydantic import BaseModel
 from pytest_httpx import HTTPXMock
 
 from uipath._cli._evals._models._evaluation_set import (
@@ -142,7 +144,7 @@ def test_llm_mockable_sync(httpx_mock: HTTPXMock, monkeypatch: MonkeyPatch):
 
     # Arrange
     @mockable()
-    def foo(*args, **kwargs):
+    def foo(*args, **kwargs) -> str:
         raise NotImplementedError()
 
     @mockable()
@@ -189,7 +191,7 @@ def foofoo(*args, **kwargs):
                     "index": 0,
                     "message": {
                         "role": "ai",
-                        "content": '{"response": "bar1"}',
+                        "content": '"bar1"',
                         "tool_calls": None,
                     },
                     "finish_reason": "EOS",
@@ -206,6 +208,19 @@ def foofoo(*args, **kwargs):
     set_execution_context(evaluation, _mock_span_collector, "test-execution-id")
 
     assert foo() == "bar1"
+
+    mock_request = httpx_mock.get_request()
+    assert mock_request
+    request = json.loads(mock_request.content.decode("utf-8"))
+    assert request["response_format"] == {
+        "type": "json_schema",
+        "json_schema": {
+            "name": "OutputSchema",
+            "strict": True,
+            "schema": {"type": "string"},
+        },
+    }
+
     with pytest.raises(NotImplementedError):
         assert foofoo()
     httpx_mock.add_response(
@@ -228,7 +243,7 @@ async def test_llm_mockable_async(httpx_mock: HTTPXMock, monkeypatch: MonkeyPatc
 
     # Arrange
     @mockable()
-    async def foo(*args, **kwargs):
+    async def foo(*args, **kwargs) -> str:
         raise NotImplementedError()
 
     @mockable()
@@ -277,7 +292,7 @@ async def foofoo(*args, **kwargs):
                     "index": 0,
                     "message": {
                         "role": "ai",
-                        "content": '{"response": "bar1"}',
+                        "content": '"bar1"',
                         "tool_calls": None,
                     },
                     "finish_reason": "EOS",
@@ -294,6 +309,19 @@ async def foofoo(*args, **kwargs):
     set_execution_context(evaluation, _mock_span_collector, "test-execution-id")
 
     assert await foo() == "bar1"
+
+    mock_request = httpx_mock.get_request()
+    assert mock_request
+    request = json.loads(mock_request.content.decode("utf-8"))
+    assert request["response_format"] == {
+        "type": "json_schema",
+        "json_schema": {
+            "name": "OutputSchema",
+            "strict": True,
+            "schema": {"type": "string"},
+        },
+    }
+
     with pytest.raises(NotImplementedError):
         assert await foofoo()
 
@@ -305,3 +333,188 @@ async def foofoo(*args, **kwargs):
     )
     with pytest.raises(UiPathMockResponseGenerationError):
         assert await foo()
+
+
+@pytest.mark.httpx_mock(assert_all_responses_were_requested=False)
+def test_llm_mockable_with_output_schema_sync(
+    httpx_mock: HTTPXMock, monkeypatch: MonkeyPatch
+):
+    monkeypatch.setenv("UIPATH_URL", "https://example.com")
+    monkeypatch.setenv("UIPATH_ACCESS_TOKEN", "1234567890")
+    monkeypatch.setattr(CacheManager, "get", lambda *args, **kwargs: None)
+    monkeypatch.setattr(CacheManager, "set", lambda *args, **kwargs: None)
+
+    class ToolResponseMock(BaseModel):
+        content: str
+
+    # Arrange
+    @mockable(output_schema=ToolResponseMock.model_json_schema())
+    def foo(*args, **kwargs) -> dict[str, Any]:
+        raise NotImplementedError()
+
+    evaluation_item: dict[str, Any] = {
+        "id": "evaluation-id",
+        "name": "Mock foo",
+        "inputs": {},
+        "evaluationCriterias": {
+            "ExactMatchEvaluator": None,
+        },
+        "mockingStrategy": {
+            "type": "llm",
+            "prompt": "response content is 'bar1'",
+            "toolsToSimulate": [{"name": "foo"}],
+        },
+    }
+    evaluation = EvaluationItem(**evaluation_item)
+    assert isinstance(evaluation.mocking_strategy, LLMMockingStrategy)
+    httpx_mock.add_response(
+        url="https://example.com/agenthub_/llm/api/capabilities",
+        status_code=200,
+        json={},
+    )
+    httpx_mock.add_response(
+        url="https://example.com/orchestrator_/llm/api/capabilities",
+        status_code=200,
+        json={},
+    )
+
+    httpx_mock.add_response(
+        url="https://example.com/llm/api/chat/completions"
+        "?api-version=2024-08-01-preview",
+        status_code=200,
+        json={
+            "id": "response-id",
+            "object": "",
+            "created": 0,
+            "model": "model",
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "ai",
+                        "content": '{"content": "bar1"}',
+                        "tool_calls": None,
+                    },
+                    "finish_reason": "EOS",
+                }
+            ],
+            "usage": {
+                "prompt_tokens": 1,
+                "completion_tokens": 1,
+                "total_tokens": 2,
+            },
+        },
+    )
+    # Act & Assert
+    set_execution_context(evaluation, _mock_span_collector, "test-execution-id")
+
+    assert foo() == {"content": "bar1"}
+    mock_request = httpx_mock.get_request()
+    assert mock_request
+    request = json.loads(mock_request.content.decode("utf-8"))
+    assert request["response_format"] == {
+        "type": "json_schema",
+        "json_schema": {
+            "name": "OutputSchema",
+            "strict": True,
+            "schema": {
+                "required": ["content"],
+                "type": "object",
+                "additionalProperties": False,
+                "properties": {"content": {"type": "string"}},
+            },
+        },
+    }
+
+
+@pytest.mark.asyncio
+@pytest.mark.httpx_mock(assert_all_responses_were_requested=False)
+async def test_llm_mockable_with_output_schema_async(
+    httpx_mock: HTTPXMock, monkeypatch: MonkeyPatch
+):
+    monkeypatch.setenv("UIPATH_URL", "https://example.com")
+    monkeypatch.setenv("UIPATH_ACCESS_TOKEN", "1234567890")
+    monkeypatch.setattr(CacheManager, "get", lambda *args, **kwargs: None)
+    monkeypatch.setattr(CacheManager, "set", lambda *args, **kwargs: None)
+
+    class ToolResponseMock(BaseModel):
+        content: str
+
+    # Arrange
+    @mockable(output_schema=ToolResponseMock.model_json_schema())
+    async def foo(*args, **kwargs) -> dict[str, Any]:
+        raise NotImplementedError()
+
+    evaluation_item: dict[str, Any] = {
+        "id": "evaluation-id",
+        "name": "Mock foo",
+        "inputs": {},
+        "evaluationCriterias": {
+            "ExactMatchEvaluator": None,
+        },
+        "mockingStrategy": {
+            "type": "llm",
+            "prompt": "response content is 'bar1'",
+            "toolsToSimulate": [{"name": "foo"}],
+        },
+    }
+    evaluation = EvaluationItem(**evaluation_item)
+    assert isinstance(evaluation.mocking_strategy, LLMMockingStrategy)
+    httpx_mock.add_response(
+        url="https://example.com/agenthub_/llm/api/capabilities",
+        status_code=200,
+        json={},
+    )
+    httpx_mock.add_response(
+        url="https://example.com/orchestrator_/llm/api/capabilities",
+        status_code=200,
+        json={},
+    )
+
+    httpx_mock.add_response(
+        url="https://example.com/llm/api/chat/completions"
+        "?api-version=2024-08-01-preview",
+        status_code=200,
+        json={
+            "id": "response-id",
+            "object": "",
+            "created": 0,
+            "model": "model",
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "ai",
+                        "content": '{"content": "bar1"}',
+                        "tool_calls": None,
+                    },
+                    "finish_reason": "EOS",
+                }
+            ],
+            "usage": {
+                "prompt_tokens": 1,
+                "completion_tokens": 1,
+                "total_tokens": 2,
+            },
+        },
+    )
+    # Act & Assert
+    set_execution_context(evaluation, _mock_span_collector, "test-execution-id")
+
+    assert await foo() == {"content": "bar1"}
+    mock_request = httpx_mock.get_request()
+    assert mock_request
+    request = json.loads(mock_request.content.decode("utf-8"))
+    assert request["response_format"] == {
+        "type": "json_schema",
+        "json_schema": {
+            "name": "OutputSchema",
+            "strict": True,
+            "schema": {
+                "required": ["content"],
+                "type": "object",
+                "additionalProperties": False,
+                "properties": {"content": {"type": "string"}},
+            },
+        },
+    }