|
1 | 1 | from typing import Any, Optional, Union |
| 2 | +from unittest.mock import Mock |
2 | 3 |
|
| 4 | +import orjson |
3 | 5 | import pytest |
4 | 6 |
|
5 | 7 | from chatlas import ChatOpenAI |
@@ -403,3 +405,93 @@ def add(x: int, y: int) -> int: |
403 | 405 | assert parsed.tool is not None |
404 | 406 | assert parsed.tool.name == "add" |
405 | 407 | assert parsed.tool.description == "Add two numbers" |
| 408 | + |
| 409 | + |
| 410 | +def test_content_tool_result_pandas_dataframe(): |
| 411 | + """Test ContentToolResult with pandas DataFrame using orient='records'""" |
| 412 | + pandas = pytest.importorskip("pandas") |
| 413 | + |
| 414 | + # Create a simple pandas DataFrame |
| 415 | + df = pandas.DataFrame( |
| 416 | + {"name": ["Alice", "Bob"], "age": [25, 30], "city": ["New York", "London"]} |
| 417 | + ) |
| 418 | + |
| 419 | + # Create ContentToolResult with DataFrame value |
| 420 | + result = ContentToolResult(value=df).get_model_value() |
| 421 | + expected = df.to_json(orient="records") |
| 422 | + assert result == expected |
| 423 | + |
| 424 | + parsed = orjson.loads(str(result)) |
| 425 | + assert isinstance(parsed, list) |
| 426 | + assert len(parsed) == 2 |
| 427 | + assert parsed[0] == {"name": "Alice", "age": 25, "city": "New York"} |
| 428 | + assert parsed[1] == {"name": "Bob", "age": 30, "city": "London"} |
| 429 | + |
| 430 | + |
| 431 | +def test_content_tool_result_object_with_to_pandas(): |
| 432 | + """Test ContentToolResult with objects that have .to_pandas() method""" |
| 433 | + pandas = pytest.importorskip("pandas") |
| 434 | + |
| 435 | + # Create mock object with to_pandas method (like Polars, PyArrow) |
| 436 | + mock_df_lib = Mock() |
| 437 | + pandas_df = pandas.DataFrame({"x": [1, 2, 3], "y": ["a", "b", "c"]}) |
| 438 | + mock_df_lib.to_pandas.return_value = pandas_df |
| 439 | + |
| 440 | + result = ContentToolResult(value=mock_df_lib).get_model_value() |
| 441 | + mock_df_lib.to_pandas.assert_called_once() |
| 442 | + expected = pandas_df.to_json(orient="records") |
| 443 | + assert result == expected |
| 444 | + |
| 445 | + |
| 446 | +def test_content_tool_result_narwhals_dataframe(): |
| 447 | + """Test ContentToolResult with narwhals DataFrame""" |
| 448 | + narwhals = pytest.importorskip("narwhals") |
| 449 | + pandas = pytest.importorskip("pandas") |
| 450 | + |
| 451 | + pandas_df = pandas.DataFrame({"a": [1, 2], "b": ["x", "y"]}) |
| 452 | + nw_df = narwhals.from_native(pandas_df) |
| 453 | + result = ContentToolResult(value=nw_df).get_model_value() |
| 454 | + expected = pandas_df.to_json(orient="records") |
| 455 | + assert result == expected |
| 456 | + |
| 457 | + |
| 458 | +def test_content_tool_result_object_with_to_dict(): |
| 459 | + """Test ContentToolResult with objects that have to_dict method""" |
| 460 | + # Mock object with to_dict method but no to_pandas or to_json |
| 461 | + mock_obj = Mock(spec=["to_dict"]) |
| 462 | + mock_obj.to_dict.return_value = {"key": "value"} |
| 463 | + result = ContentToolResult(value=mock_obj).get_model_value() |
| 464 | + mock_obj.to_dict.assert_called_once() |
| 465 | + # Result should be JSON string representation (orjson format) |
| 466 | + assert result == '{"key":"value"}' |
| 467 | + |
| 468 | + |
| 469 | +def test_content_tool_result_string_passthrough(): |
| 470 | + """Test ContentToolResult with string values (special case - passed through as-is)""" |
| 471 | + result = ContentToolResult(value="plain string").get_model_value() |
| 472 | + assert result == "plain string" |
| 473 | + |
| 474 | + |
| 475 | +def test_content_tool_result_fallback_serialization(): |
| 476 | + """Test ContentToolResult fallback for objects without special methods""" |
| 477 | + # Regular object without to_json, to_pandas, or to_dict (non-string to avoid the string special case) |
| 478 | + result = ContentToolResult(value={"key": "value"}).get_model_value() |
| 479 | + assert result == '{"key":"value"}' |
| 480 | + |
| 481 | + |
| 482 | +def test_content_tool_result_explicit_json_mode(): |
| 483 | + """Test ContentToolResult with explicit JSON mode forces _to_json for non-strings""" |
| 484 | + # Test with non-string object and explicit JSON mode |
| 485 | + result = ContentToolResult( |
| 486 | + value={"key": "value"}, |
| 487 | + model_format="json", |
| 488 | + ).get_model_value() |
| 489 | + # With explicit JSON mode, objects get JSON-encoded |
| 490 | + assert result == '{"key":"value"}' |
| 491 | + # Test that strings still get special treatment even in JSON mode |
| 492 | + string_result = ContentToolResult( |
| 493 | + value="plain string", |
| 494 | + model_format="json", |
| 495 | + ).get_model_value() |
| 496 | + # Strings are still returned as-is even in JSON mode (current behavior) |
| 497 | + assert string_result == "plain string" |
0 commit comments