feat: Make DSPY instrumentation read from attributes from context (#421)

Arize-ai · May 4, 2024 · 60c9b9c · 60c9b9c
1 parent 54d0931
commit 60c9b9c
Show file tree

Hide file tree

Showing 4 changed files with 720 additions and 124 deletions.
diff --git a/python/instrumentation/openinference-instrumentation-dspy/examples/dspy_basic_qa.py b/python/instrumentation/openinference-instrumentation-dspy/examples/dspy_basic_qa.py
@@ -1,22 +1,22 @@
 import os
 
 import dspy
+from openinference.instrumentation import using_attributes
 from openinference.instrumentation.dspy import DSPyInstrumentor
 from opentelemetry import trace as trace_api
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk import trace as trace_sdk
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor
 
-resource = Resource(attributes={})
-tracer_provider = trace_sdk.TracerProvider(resource=resource)
-span_console_exporter = ConsoleSpanExporter()
-tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter=span_console_exporter))
 # Logs to the Phoenix Collector if running locally
 if os.environ.get("PHOENIX_COLLECTOR_ENDPOINT"):
     endpoint = os.environ["PHOENIX_COLLECTOR_ENDPOINT"] + "/v1/traces"
-    span_otlp_exporter = OTLPSpanExporter(endpoint=endpoint)
-    tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter=span_otlp_exporter))
+
+resource = Resource(attributes={})
+tracer_provider = trace_sdk.TracerProvider(resource=resource)
+tracer_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter(endpoint)))
+tracer_provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))
 
 trace_api.set_tracer_provider(tracer_provider=tracer_provider)
 
@@ -36,11 +36,25 @@ class BasicQA(dspy.Signature):
 
     dspy.settings.configure(lm=turbo)
 
-    # Define the predictor.
-    generate_answer = dspy.Predict(BasicQA)
-
-    # Call the predictor on a particular input.
-    pred = generate_answer(
-        question="What is the capital of the united states?"  # noqa: E501
-    )  # noqa: E501
-    print(f"Predicted Answer: {pred.answer}")
+    with using_attributes(
+        session_id="my-test-session",
+        user_id="my-test-user",
+        metadata={
+            "test-int": 1,
+            "test-str": "string",
+            "test-list": [1, 2, 3],
+            "test-dict": {
+                "key-1": "val-1",
+                "key-2": "val-2",
+            },
+        },
+        tags=["tag-1", "tag-2"],
+    ):
+        # Define the predictor.
+        generate_answer = dspy.Predict(BasicQA)
+
+        # Call the predictor on a particular input.
+        pred = generate_answer(
+            question="What is the capital of the united states?"  # noqa: E501
+        )  # noqa: E501
+        print(f"Predicted Answer: {pred.answer}")
diff --git a/python/instrumentation/openinference-instrumentation-dspy/pyproject.toml b/python/instrumentation/openinference-instrumentation-dspy/pyproject.toml
@@ -27,6 +27,7 @@ dependencies = [
   "opentelemetry-api",
   "opentelemetry-instrumentation",
   "opentelemetry-semantic-conventions",
+  "openinference-instrumentation>=0.1.3",
   "openinference-semantic-conventions",
   "wrapt",
   "typing-extensions",

diff --git a/...ion/openinference-instrumentation-dspy/src/openinference/instrumentation/dspy/__init__.py b/...ion/openinference-instrumentation-dspy/src/openinference/instrumentation/dspy/__init__.py
@@ -17,6 +17,7 @@
 )
 
 import opentelemetry.context as context_api
+from openinference.instrumentation import get_attributes_from_context
 from openinference.instrumentation.dspy.package import _instruments
 from openinference.instrumentation.dspy.version import __version__
 from openinference.semconv.trace import (
@@ -221,6 +222,7 @@ def __call__(
                 )
             ),
         ) as span:
+            span.set_attributes(dict(get_attributes_from_context()))
             try:
                 response = wrapped(*args, **kwargs)
             except Exception as exception:
@@ -300,6 +302,7 @@ def __call__(
                 )
             ),
         ) as span:
+            span.set_attributes(dict(get_attributes_from_context()))
             try:
                 prediction = wrapped(*args, **kwargs)
             except Exception as exception:
@@ -369,6 +372,7 @@ def __call__(
                 )
             ),
         ) as span:
+            span.set_attributes(dict(get_attributes_from_context()))
             try:
                 prediction = wrapped(*args, **kwargs)
             except Exception as exception:
@@ -420,6 +424,7 @@ def __call__(
                 )
             ),
         ) as span:
+            span.set_attributes(dict(get_attributes_from_context()))
             try:
                 prediction = wrapped(*args, **kwargs)
             except Exception as exception:
@@ -472,6 +477,7 @@ def __call__(
                 )
             ),
         ) as span:
+            span.set_attributes(dict(get_attributes_from_context()))
             try:
                 retrieved_documents = wrapped(*args, **kwargs)
             except Exception as exception: