Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .github/workflows/pr_build.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
Expand Down
25 changes: 25 additions & 0 deletions scripts/litellm_pydantic_pytest_plugin.py
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,7 @@ def _mark_as_instrumented(library_key, wrapper_instance):
"groq",
"langchain",
"llamaindex",
"mcp",
"mistralai",
"ollama",
"replicate",
Expand Down
11 changes: 11 additions & 0 deletions test/instrumentation/test_instrumentation_definitions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down
Loading