diff --git a/.github/workflows/pr_build.yaml b/.github/workflows/pr_build.yaml index 2c3bfff..758d541 100644 --- a/.github/workflows/pr_build.yaml +++ b/.github/workflows/pr_build.yaml @@ -75,7 +75,7 @@ jobs: - name: Unit tests run: | export RUN_SDK_INTEGRATION_TESTS=1 - python -m pytest --log-cli-level=INFO --tb=long + python -m pytest -p scripts.litellm_pydantic_pytest_plugin --log-cli-level=INFO --tb=long env: PYTHONUNBUFFERED: 1 HA_ENABLE_CONSOLE_SPAN_EXPORTER: "true" @@ -84,7 +84,7 @@ jobs: if: matrix.python-version == '3.12' run: | export RUN_SDK_INTEGRATION_TESTS=1 - python -m pytest --log-cli-level=INFO --tb=long \ + python -m pytest -p scripts.litellm_pydantic_pytest_plugin --log-cli-level=INFO --tb=long \ --cov=harness_sdk --cov-report=xml:coverage.xml --junitxml=pytest-report.xml env: PYTHONUNBUFFERED: 1 diff --git a/scripts/litellm_pydantic_pytest_plugin.py b/scripts/litellm_pydantic_pytest_plugin.py new file mode 100644 index 0000000..926079e --- /dev/null +++ b/scripts/litellm_pydantic_pytest_plugin.py @@ -0,0 +1,25 @@ +"""Pytest plugin to rebuild litellm pydantic models in the test process.""" + + +def pytest_configure(config): # pylint: disable=unused-argument + try: + from litellm.types.llms.openai import ( # noqa: F401 + ChatCompletionReasoningSummaryTextBlock, + ) + import litellm.types.utils as litellm_utils + from pydantic import BaseModel + except ImportError: + return + + for name in dir(litellm_utils): + obj = getattr(litellm_utils, name) + if ( + isinstance(obj, type) + and issubclass(obj, BaseModel) + and obj is not BaseModel + and hasattr(obj, "model_rebuild") + ): + try: + obj.model_rebuild() + except Exception: + pass diff --git a/src/harness_sdk/instrumentation/instrumentation_definitions.py b/src/harness_sdk/instrumentation/instrumentation_definitions.py index 57b7419..e583418 100644 --- a/src/harness_sdk/instrumentation/instrumentation_definitions.py +++ b/src/harness_sdk/instrumentation/instrumentation_definitions.py @@ -164,6 +164,7 @@ def _mark_as_instrumented(library_key, wrapper_instance): "groq", "langchain", "llamaindex", + "mcp", "mistralai", "ollama", "replicate", diff --git a/test/instrumentation/test_instrumentation_definitions.py b/test/instrumentation/test_instrumentation_definitions.py index 3735f63..4c77535 100644 --- a/test/instrumentation/test_instrumentation_definitions.py +++ b/test/instrumentation/test_instrumentation_definitions.py @@ -160,6 +160,17 @@ def test_instrument_skips_denylisted_libraries(): assert "aiobotocore" not in _GENERIC_INSTRUMENTATION_STATE +def test_instrument_skips_mcp_generic_contrib(): + mcp_ep = _make_entry_point("mcp") + + with patch("harness_sdk.instrumentation.instrumentation_definitions._get_contrib_instrumentation_entry_points", + return_value=[mcp_ep]): + instrument_supported_contrib_without_wrapper() + + mcp_ep.load.assert_not_called() + assert "mcp" not in _GENERIC_INSTRUMENTATION_STATE + + def test_instrument_skips_ai_focused_generic_instrumentors(): langchain_ep = _make_entry_point("langchain")