Jonathan.chavez/llm experiments 2

ddtrace/llmobs/__init__.py

@@ -5,7 +5,6 @@
     from ddtrace.llmobs import LLMObs
     LLMObs.enable()
 """
-
 from ._llmobs import LLMObs
 
 

ddtrace/llmobs/_constants.py

@@ -15,7 +15,9 @@
 INPUT_VALUE = "_ml_obs.meta.input.value"
 INPUT_PARAMETERS = "_ml_obs.meta.input.parameters"
 INPUT_PROMPT = "_ml_obs.meta.input.prompt"
-
+EXPECTED_OUTPUT = "_ml_obs.meta.input.expected_output"
+EXPERIMENT_INPUT = "_ml_obs.meta.input"
+EXPERIMENT_OUTPUT = "_ml_obs.meta.output"
 OUTPUT_DOCUMENTS = "_ml_obs.meta.output.documents"
 OUTPUT_MESSAGES = "_ml_obs.meta.output.messages"
 OUTPUT_VALUE = "_ml_obs.meta.output.value"

ddtrace/llmobs/_integrations/base.py

@@ -210,7 +210,8 @@ def llmobs_set_tags(
             return
         try:
             self._llmobs_set_tags(span, args, kwargs, response, operation)
-        except Exception:
+        except Exception as e:
+            print(e)
             log.error("Error extracting LLMObs fields for span %s, likely due to malformed data", span, exc_info=True)
 
     @abc.abstractmethod

ddtrace/llmobs/_llmobs.py

@@ -52,6 +52,9 @@
 from ddtrace.llmobs._constants import SPAN_KIND
 from ddtrace.llmobs._constants import SPAN_START_WHILE_DISABLED_WARNING
 from ddtrace.llmobs._constants import TAGS
+from ddtrace.llmobs._constants import EXPECTED_OUTPUT
+from ddtrace.llmobs._constants import EXPERIMENT_INPUT
+from ddtrace.llmobs._constants import EXPERIMENT_OUTPUT
 from ddtrace.llmobs._evaluators.runner import EvaluatorRunner
 from ddtrace.llmobs._utils import AnnotationContext
 from ddtrace.llmobs._utils import _get_llmobs_parent_id
@@ -193,6 +196,14 @@ def _llmobs_span_event(cls, span: Span) -> Tuple[Dict[str, Any], bool]:
         span._set_ctx_item(ML_APP, ml_app)
         parent_id = str(_get_llmobs_parent_id(span) or "undefined")
 
+        # Experiments related
+        if span._get_ctx_item(EXPECTED_OUTPUT) is not None:
+            meta["expected_output"] = span._get_ctx_item(EXPECTED_OUTPUT)
+        if span._get_ctx_item(EXPERIMENT_INPUT) is not None:
+            meta["input"] = span._get_ctx_item(EXPERIMENT_INPUT)
+        if span._get_ctx_item(EXPERIMENT_OUTPUT) is not None:
+            meta["output"] = span._get_ctx_item(EXPERIMENT_OUTPUT)
+
         llmobs_span_event = {
             "trace_id": "{:x}".format(span.trace_id),
             "span_id": str(span.span_id),
@@ -212,6 +223,7 @@ def _llmobs_span_event(cls, span: Span) -> Tuple[Dict[str, Any], bool]:
         llmobs_span_event["tags"] = cls._llmobs_tags(
             span, ml_app, session_id, is_ragas_integration_span=is_ragas_integration_span
         )
+
         return llmobs_span_event, is_ragas_integration_span
 
     @staticmethod
@@ -626,6 +638,22 @@ def agent(cls, name: Optional[str] = None, session_id: Optional[str] = None, ml_
         if cls.enabled is False:
             log.warning(SPAN_START_WHILE_DISABLED_WARNING)
         return cls._instance._start_span("agent", name=name, session_id=session_id, ml_app=ml_app)
+
+    @classmethod
+    def _experiment(cls, name: Optional[str] = None, session_id: Optional[str] = None, ml_app: Optional[str] = None) -> Span:
+        """
+        Trace a dynamic workflow in which an embedded language model (agent) decides what sequence of actions to take.
+
+        :param str name: The name of the traced operation. If not provided, a default value of "agent" will be set.
+        :param str session_id: The ID of the underlying user session. Required for tracking sessions.
+        :param str ml_app: The name of the ML application that the agent is orchestrating. If not provided, the default
+                           value will be set to the value of `DD_LLMOBS_ML_APP`.
+
+        :returns: The Span object representing the traced operation.
+        """
+        if cls.enabled is False:
+            log.warning(SPAN_START_WHILE_DISABLED_WARNING)
+        return cls._instance._start_span("experiment", name=name, session_id=session_id, ml_app=ml_app)
 
     @classmethod
     def workflow(
@@ -785,9 +813,20 @@ def annotate(
                 cls._tag_embedding_io(span, input_documents=input_data, output_text=output_data)
             elif span_kind == "retrieval":
                 cls._tag_retrieval_io(span, input_text=input_data, output_documents=output_data)
+            elif span_kind == "experiment":
+                cls._tag_experiment_io(span, input_data=input_data, output_data=output_data)
             else:
                 cls._tag_text_io(span, input_value=input_data, output_value=output_data)
 
+    @staticmethod
+    def _tag_expected_output(span, expected_output: dict) -> None:
+        """Tags a given LLMObs span with a prompt"""
+        try:
+            span._set_ctx_item(EXPECTED_OUTPUT, expected_output)
+        except TypeError:
+            log.warning("Failed to validate expected output with error: ", exc_info=True)
+            return
+
     @staticmethod
     def _tag_prompt(span, prompt: dict) -> None:
         """Tags a given LLMObs span with a prompt"""
@@ -877,6 +916,16 @@ def _tag_text_io(cls, span, input_value=None, output_value=None):
         if output_value is not None:
             span._set_ctx_item(OUTPUT_VALUE, str(output_value))
 
+    @classmethod
+    def _tag_experiment_io(cls, span, input_data=None, output_data=None):
+        """Tags input/output values for experiment kind spans.
+        Will be mapped to span's `meta.{input,output}.values` fields.
+        """
+        if input_data is not None:
+            span._set_ctx_item(EXPERIMENT_INPUT, input_data)
+        if output_data is not None:
+            span._set_ctx_item(EXPERIMENT_OUTPUT, output_data)
+
     @staticmethod
     def _tag_span_tags(span: Span, span_tags: Dict[str, Any]) -> None:
         """Tags a given LLMObs span with a dictionary of key-value tag pairs.

ddtrace/llmobs/_utils.py

@@ -1,8 +1,11 @@
+import http.client
 import json
 from typing import Dict
 from typing import List
 from typing import Optional
 from typing import Union
+import urllib.request
+from urllib.error import HTTPError
 
 import ddtrace
 from ddtrace import Span
@@ -164,6 +167,7 @@ def _get_session_id(span: Span) -> Optional[str]:
 def _inject_llmobs_parent_id(span_context):
     """Inject the LLMObs parent ID into the span context for reconnecting distributed LLMObs traces."""
     span = ddtrace.tracer.current_span()
+
     if span is None:
         log.warning("No active span to inject LLMObs parent ID info.")
         return
@@ -175,6 +179,7 @@ def _inject_llmobs_parent_id(span_context):
         llmobs_parent_id = str(span.span_id)
     else:
         llmobs_parent_id = _get_llmobs_parent_id(span)
+
     span_context._meta[PROPAGATED_PARENT_ID_KEY] = llmobs_parent_id or "undefined"
 
 
@@ -191,3 +196,49 @@ def safe_json(obj):
         return json.dumps(obj, ensure_ascii=False, skipkeys=True, default=_unserializable_default_repr)
     except Exception:
         log.error("Failed to serialize object to JSON.", exc_info=True)
+
+
+class HTTPResponse:
+    def __init__(self, resp) -> None:
+        if resp is None:
+            raise ValueError("Response object cannot be None")
+        self._resp = resp
+        self._content = None  # Cache the content
+
+    @property
+    def status_code(self) -> int:
+        if hasattr(self._resp, 'status'):
+            return self._resp.status
+        elif hasattr(self._resp, 'code'):
+            return self._resp.code
+        elif hasattr(self._resp, 'getcode'):
+            return self._resp.getcode()
+        else:
+            raise AttributeError(f"Could not find status code in response object of type {type(self._resp)}")
+
+    def read(self) -> bytes:
+        if self._content is None:
+            self._content = self._resp.read()
+        return self._content
+
+    def text(self) -> str:
+        return self.read().decode('utf-8')
+
+    def json(self) -> dict:
+        return json.loads(self.text())
+
+
+def http_request(
+    method: str, url: str, headers: Optional[Dict[str, str]] = None, body: Optional[bytes] = None
+) -> HTTPResponse:
+    """Make an HTTP request and return an HTTPResponse object."""
+    # Create the request object
+    req = urllib.request.Request(url, data=body, method=method)
+    if headers:
+        req.headers.update(headers)
+    try:
+        response = urllib.request.urlopen(req)
+        return HTTPResponse(response)
+    except HTTPError as e:
+        # Create an HTTPResponse object from the error response
+        return HTTPResponse(e)

ddtrace/llmobs/_writer.py

@@ -32,6 +32,8 @@
 from ddtrace.llmobs._constants import EVP_SUBDOMAIN_HEADER_NAME
 from ddtrace.llmobs._constants import EVP_SUBDOMAIN_HEADER_VALUE
 from ddtrace.llmobs._utils import safe_json
+from ddtrace.internal.utils.formats import asbool
+import os
 
 
 logger = get_logger(__name__)
@@ -208,12 +210,18 @@ def encode(self):
             events = self._buffer
             self._init_buffer()
         data = {"_dd.stage": "raw", "_dd.tracer_version": ddtrace.__version__, "event_type": "span", "spans": events}
+        if asbool(os.getenv("DD_EXPERIMENTS_RUNNER_ENABLED")):
+            data["_dd.scope"] = "experiments"
         try:
             enc_llm_events = safe_json(data)
+            if isinstance(enc_llm_events, str):
+                enc_llm_events = enc_llm_events.encode('utf-8')
             logger.debug("encode %d LLMObs span events to be sent", len(events))
+
         except TypeError:
             logger.error("failed to encode %d LLMObs span events", len(events), exc_info=True)
             return None, 0
+        # print(enc_llm_events)
         return enc_llm_events, len(events)
 
 

ddtrace/llmobs/experimentation/__init__.py

@@ -0,0 +1,11 @@
+"""
+LLM Datasets and Experiments.
+"""
+from ._experiments import Dataset
+from ._experiments import Experiment
+from ._experiments import task
+from ._experiments import evaluator
+from ._experiments import init
+
+
+__all__ = ["Dataset", "Experiment", "task", "evaluator", "init"]