From 5f6e149d0979a822e07f81af944c22b7530f8fed Mon Sep 17 00:00:00 2001
From: Xander Song <axiomofjoy@gmail.com>
Date: Mon, 11 Nov 2024 21:49:29 -0800
Subject: [PATCH] fix(dspy): instrument `dspy` rather than `dspy-ai` (#1113)

---
 .../examples/requirements.txt                             | 2 +-
 .../openinference-instrumentation-dspy/pyproject.toml     | 4 ++--
 .../src/openinference/instrumentation/dspy/package.py     | 2 +-
 .../instrumentation/dspy/test_instrumentor.py             | 8 ++++----
 4 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/python/instrumentation/openinference-instrumentation-dspy/examples/requirements.txt b/python/instrumentation/openinference-instrumentation-dspy/examples/requirements.txt
index 3c0ec027a..76e21f6e3 100644
--- a/python/instrumentation/openinference-instrumentation-dspy/examples/requirements.txt
+++ b/python/instrumentation/openinference-instrumentation-dspy/examples/requirements.txt
@@ -1,4 +1,4 @@
-dspy-ai >= 2.1.0
+dspy >= 2.5.0
 opentelemetry-sdk
 opentelemetry-exporter-otlp
 openinference-instrumentation-dspy
\ No newline at end of file
diff --git a/python/instrumentation/openinference-instrumentation-dspy/pyproject.toml b/python/instrumentation/openinference-instrumentation-dspy/pyproject.toml
index 5c0c64f4e..b407373df 100644
--- a/python/instrumentation/openinference-instrumentation-dspy/pyproject.toml
+++ b/python/instrumentation/openinference-instrumentation-dspy/pyproject.toml
@@ -35,10 +35,10 @@ dependencies = [
 
 [project.optional-dependencies]
 instruments = [
-  "dspy-ai >= 2.5.0",
+  "dspy >= 2.5.0",
 ]
 test = [
-  "dspy-ai==2.5.0",
+  "dspy>=2.5.27",
   "google-generativeai",
   "opentelemetry-sdk",
   "pytest-recording",
diff --git a/python/instrumentation/openinference-instrumentation-dspy/src/openinference/instrumentation/dspy/package.py b/python/instrumentation/openinference-instrumentation-dspy/src/openinference/instrumentation/dspy/package.py
index 5825ce077..02cd1b833 100644
--- a/python/instrumentation/openinference-instrumentation-dspy/src/openinference/instrumentation/dspy/package.py
+++ b/python/instrumentation/openinference-instrumentation-dspy/src/openinference/instrumentation/dspy/package.py
@@ -1,2 +1,2 @@
-_instruments = ("dspy-ai >= 2.1.0",)
+_instruments = ("dspy >= 2.5.0",)
 _supports_metrics = False
diff --git a/python/instrumentation/openinference-instrumentation-dspy/tests/openinference/instrumentation/dspy/test_instrumentor.py b/python/instrumentation/openinference-instrumentation-dspy/tests/openinference/instrumentation/dspy/test_instrumentor.py
index 538c48b1a..4621d76e0 100644
--- a/python/instrumentation/openinference-instrumentation-dspy/tests/openinference/instrumentation/dspy/test_instrumentor.py
+++ b/python/instrumentation/openinference-instrumentation-dspy/tests/openinference/instrumentation/dspy/test_instrumentor.py
@@ -15,7 +15,6 @@
     backtrack_handler,
 )
 from dspy.teleprompt import BootstrapFewShotWithRandomSearch
-from litellm import AuthenticationError  # type: ignore[attr-defined]
 from opentelemetry import trace as trace_api
 from opentelemetry.sdk import trace as trace_sdk
 from opentelemetry.sdk.resources import Resource
@@ -291,7 +290,7 @@ def test_exception_event_recorded_on_lm_error(
     ) -> None:
         lm = dspy.LM("openai/gpt-4", cache=False)
         prompt = "Who won the World Cup in 2018?"
-        with pytest.raises(AuthenticationError):
+        with pytest.raises(Exception):
             lm(prompt)
         spans = in_memory_span_exporter.get_finished_spans()
         assert len(spans) == 1
@@ -302,9 +301,10 @@ def test_exception_event_recorded_on_lm_error(
         event = span.events[0]
         assert event.name == "exception"
         assert (event_attributes := event.attributes) is not None
-        assert event_attributes["exception.type"] == "litellm.exceptions.AuthenticationError"
+        assert isinstance(exception_type := event_attributes["exception.type"], str)
+        assert exception_type.startswith("litellm.exceptions")
         assert isinstance(exception_message := event_attributes["exception.message"], str)
-        assert "401" in exception_message
+        assert "Connection error" in exception_message or "401" in exception_message
         attributes = dict(span.attributes or {})
         assert attributes.pop(OPENINFERENCE_SPAN_KIND) == LLM
         assert attributes.pop(INPUT_MIME_TYPE) == JSON