[BUZZOK-25849] Bump requirements to add opentelemetry for LLMs (#1438)

* Squash and rebase

* Update run_agent.py

* Delint

Author: mjnitz02

public_dropin_environments/python311_genai_agents/env_info.json

@@ -3,7 +3,7 @@
   "name": "[DataRobot] Python 3.11 GenAI Agents",
   "description": "This template environment can be used to create GenAI-powered agents using CrewAI, LangGraph, or Llama-Index. Similar to other drop-in environments, you can either include a .pth artifact or any other code needed to deserialize your model, and optionally a custom.py file. You can also use this environment in codespaces.",
   "programmingLanguage": "python",
-  "environmentVersionId": "68271d84cbce0247227de599",
+  "environmentVersionId": "682c93429604e979f80dc8f6",
   "isPublic": true,
   "useCases": [
     "notebook",

public_dropin_environments/python311_genai_agents/requirements.in

@@ -25,4 +25,14 @@ llama-index-llms-langchain
 llama-index-llms-litellm
 openai
 ragas
+
+# Generative AI Telemetry
 traceloop-sdk
+opentelemetry-api
+opentelemetry-sdk
+opentelemetry-instrumentation-aiohttp-client
+opentelemetry-instrumentation-crewai
+opentelemetry-instrumentation-langchain
+opentelemetry-instrumentation-llamaindex
+opentelemetry-instrumentation-openai
+opentelemetry-instrumentation-requests

public_dropin_environments/python311_genai_agents/requirements.txt

@@ -201,6 +201,7 @@ opentelemetry-exporter-otlp-proto-common==1.33.0
 opentelemetry-exporter-otlp-proto-grpc==1.33.0
 opentelemetry-exporter-otlp-proto-http==1.33.0
 opentelemetry-instrumentation==0.54b0
+opentelemetry-instrumentation-aiohttp-client==0.54b0
 opentelemetry-instrumentation-alephalpha==0.40.3
 opentelemetry-instrumentation-anthropic==0.40.3
 opentelemetry-instrumentation-asgi==0.54b0

public_dropin_environments/python311_genai_agents/run_agent.py

@@ -12,33 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: skip-file
-
 import argparse
 import json
 import logging
 import os
 import sys
-from typing import Any, cast
+from typing import cast
 
 import requests
 from datarobot_drum.drum.enum import TargetType
 from datarobot_drum.drum.root_predictors.drum_server_utils import DrumServerRun
 from openai import OpenAI
-from openai.types.chat import (
-    ChatCompletion,
-    ChatCompletionSystemMessageParam,
-    ChatCompletionUserMessageParam,
-)
-from openai.types.chat.completion_create_params import (
-    CompletionCreateParamsNonStreaming,
-)
+from openai.types.chat import ChatCompletion
 
 root = logging.getLogger()
 
 parser = argparse.ArgumentParser()
-parser.add_argument("--user_prompt", type=str, default="", help="user_prompt for chat endpoint")
-parser.add_argument("--extra_body", type=str, default="", help="extra_body for chat endpoint")
+parser.add_argument(
+    "--chat_completion",
+    type=str,
+    default="{}",
+    help="OpenAI ChatCompletion dict as json string",
+)
+parser.add_argument(
+    "--default_headers",
+    type=str,
+    default="{}",
+    help="OpenAI default_headers as json string",
+)
 parser.add_argument(
     "--custom_model_dir",
     type=str,
@@ -72,29 +73,8 @@ def setup_logging(logger: logging.Logger, output_path: str, log_level: int = log
     logger.addHandler(handler_file)
 
 
-def construct_prompt(user_prompt: str, extra_body: str) -> Any:
-    extra_body_params = json.loads(extra_body) if extra_body else {}
-    completion_create_params = CompletionCreateParamsNonStreaming(
-        model="datarobot-deployed-llm",
-        messages=[
-            ChatCompletionSystemMessageParam(
-                content="You are a helpful assistant",
-                role="system",
-            ),
-            ChatCompletionUserMessageParam(
-                content=user_prompt,
-                role="user",
-            ),
-        ],
-        n=1,
-        temperature=0.01,
-        extra_body=extra_body_params,  # type: ignore[typeddict-unknown-key]
-    )
-    return completion_create_params
-
-
 def execute_drum(
-    user_prompt: str, extra_body: str, custom_model_dir: str, output_path: str
+    chat_completion: str, default_headers: str, custom_model_dir: str, output_path: str
 ) -> ChatCompletion:
     root.info("Executing agent as [chat] endpoint. DRUM Executor.")
     root.info("Starting DRUM server.")
@@ -111,22 +91,30 @@ def execute_drum(
         port=8191,
         stream_output=True,
     ) as drum_runner:
-        root.info("Verifying DRUM server.")
+        root.info("Verifying DRUM server")
         response = requests.get(drum_runner.url_server_address)
         if not response.ok:
-            raise RuntimeError("Server failed to start")
+            root.error("Server failed to start")
+            try:
+                root.error(response.text)
+                root.error(response.json())
+            finally:
+                raise RuntimeError("Server failed to start")
+
+        root.info("Parsing OpenAI request")
+        completion_params = json.loads(chat_completion)
+        header_params = json.loads(default_headers)
 
         # Use a standard OpenAI client to call the DRUM server. This mirrors the behavior of a deployed agent.
-        root.info("Building prompt.")
+        # Using the `chat.completions.create` method ensures the parameters are OpenAI compatible.
+        root.info("Executing Agent")
         client = OpenAI(
             base_url=drum_runner.url_server_address,
             api_key="not-required",
+            default_headers=header_params,
             max_retries=0,
         )
-        completion_create_params = construct_prompt(user_prompt, extra_body)
-
-        root.info("Executing Agent.")
-        completion = client.chat.completions.create(**completion_create_params)
+        completion = client.chat.completions.create(**completion_params)
 
     # Continue outside the context manager to ensure the server is stopped and logs
     # are flushed before we write the output
@@ -145,8 +133,8 @@ def execute_drum(
     args.custom_model_dir = os.path.join(os.getcwd(), "custom_model")
 setup_logging(logger=root, output_path=args.output_path, log_level=logging.INFO)
 result = execute_drum(
-    user_prompt=args.user_prompt,
-    extra_body=args.extra_body,
+    chat_completion=args.chat_completion,
+    default_headers=args.default_headers,
     custom_model_dir=args.custom_model_dir,
     output_path=args.output_path,
 )