openai: adds examples and flattens tests

This adds an examples directory which include two main features this
instrumentation supports: chat and embeddings

It also flattens test recordings and infrastructure around unit tests.
Notably, instrumentation code doesn't change depending on if azure is
used or not. Unit testing only openai (platform) simplifies tests and
reduces infrastructure significantly.

Signed-off-by: Adrian Cole <[email protected]>

Author: codefromthecrypt

instrumentation/elastic-opentelemetry-instrumentation-openai/README.md

@@ -22,32 +22,31 @@ pip install elastic-opentelemetry-instrumentation-openai
 
 This instrumentation supports *zero-code* / *autoinstrumentation*:
 
-```
-opentelemetry-instrument python use_openai.py
+You can see telemetry from this package if you have an OpenTelemetry collector started, for example
+as documented in the root [examples](../../examples/) folder.
 
-# You can record more information about prompts as log events by enabling content capture.
-OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT=true opentelemetry-instrument python use_openai.py
+Set up a virtual environment with this package, the dependencies it requires
+and `dotenv` (a portable way to load environment variables).
+```
+python3 -m venv .venv
+source .venv/bin/activate
+pip install -r test-requirements.txt
+pip install python-dotenv[cli]
 ```
 
-Or manual instrumentation:
-
-```python
-import openai
-from opentelemetry.instrumentation.openai import OpenAIInstrumentor
-
-OpenAIInstrumentor().instrument()
-
-# assumes at least the OPENAI_API_KEY environment variable set
-client = openai.Client()
+Run the script with telemetry setup to use the instrumentation. [ollama.env](ollama.env)
+includes variables to point to Ollama instead of OpenAI, which allows you to
+run examples without a cloud account:
 
-messages = [
-    {
-        "role": "user",
-        "content": "Answer in up to 3 words: Which ocean contains the canarian islands?",
-    }
-]
+```
+dotenv -f ollama.env run -- \
+opentelemetry-instrument python examples/chat.py
+```
 
-chat_completion = client.chat.completions.create(model="gpt-4o-mini", messages=messages)
+You can record more information about prompts as log events by enabling content capture.
+```
+OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT=true dotenv -f ollama.env run -- \
+opentelemetry-instrument python examples/chat.py
 ```
 
 ### Instrumentation specific environment variable configuration
@@ -110,20 +109,22 @@ response without querying the LLM.
 
 ### Azure OpenAI Environment Variables
 
-Azure is different from OpenAI primarily in that a URL has an implicit model. This means it ignores
-the model parameter set by the OpenAI SDK. The implication is that one endpoint cannot serve both
-chat and embeddings at the same time. Hence, we need separate environment variables for chat and
-embeddings. In either case, the `DEPLOYMENT_URL` is the "Endpoint Target URI" and the `API_KEY` is
-the `Endpoint Key` for a corresponding deployment in https://oai.azure.com/resource/deployments
-
-* `AZURE_CHAT_COMPLETIONS_DEPLOYMENT_URL`
-  * It should look like https://endpoint.com/openai/deployments/my-deployment/chat/completions?api-version=2023-05-15
-* `AZURE_CHAT_COMPLETIONS_API_KEY`
-  * It should be in hex like `abc01...` and possibly the same as `AZURE_EMBEDDINGS_API_KEY`
-* `AZURE_EMBEDDINGS_DEPLOYMENT_URL`
-  * It should look like https://endpoint.com/openai/deployments/my-deployment/embeddings?api-version=2023-05-15
-* `AZURE_EMBEDDINGS_API_KEY`
-  * It should be in hex like `abc01...` and possibly the same as `AZURE_CHAT_COMPLETIONS_API_KEY`
+The `AzureOpenAI` client extends `OpenAI` with parameters specific to the Azure OpenAI Service.
+
+* `AZURE_OPENAI_ENDPOINT` - "Azure OpenAI Endpoint" in https://oai.azure.com/resource/overview
+  * It should look like `https://<your-resource-name>.openai.azure.com/`
+* `AZURE_OPENAI_API_KEY` - "API key 1 (or 2)" in https://oai.azure.com/resource/overview
+  * It should look be a hex string like `abc01...`
+* `OPENAI_API_VERSION` = "Inference version" from https://learn.microsoft.com/en-us/azure/ai-services/openai/api-version-deprecation
+  * It should look like `2024-10-01-preview`
+* `TEST_CHAT_MODEL` = "Name" from https://oai.azure.com/resource/deployments that deployed a model
+  that supports tool calling, such as "gpt-4o-mini".
+* `TEST_EMBEDDINGS_MODEL` = "Name" from https://oai.azure.com/resource/deployments that deployed a
+  model that supports embeddings, such as "text-embedding-3-small".
+
+Note: The model parameter of a chat completion or embeddings request is substituted for an identical
+deployment name. As deployment names are arbitrary they may have no correlation with a real model
+like `gpt-4o`
 
 ## License
 

instrumentation/elastic-opentelemetry-instrumentation-openai/examples/chat.py

@@ -0,0 +1,23 @@
+import os
+
+import openai
+
+CHAT_MODEL = os.environ.get("TEST_CHAT_MODEL", "gpt-4o-mini")
+
+
+def main():
+    client = openai.Client()
+
+    messages = [
+        {
+            "role": "user",
+            "content": "Answer in up to 3 words: Which ocean contains Bouvet Island?",
+        }
+    ]
+
+    chat_completion = client.chat.completions.create(model=CHAT_MODEL, messages=messages)
+    print(chat_completion.choices[0].message.content)
+
+
+if __name__ == "__main__":
+    main()

instrumentation/elastic-opentelemetry-instrumentation-openai/examples/embeddings.py

@@ -0,0 +1,49 @@
+import os
+
+import numpy as np
+import openai
+
+EMBEDDINGS_MODEL = os.environ.get("TEST_EMBEDDINGS_MODEL", "text-embedding-3-small")
+
+
+def main():
+    client = openai.Client()
+
+    products = [
+        "Search: Ingest your data, and explore Elastic's machine learning and retrieval augmented generation (RAG) capabilities."
+        "Observability: Unify your logs, metrics, traces, and profiling at scale in a single platform.",
+        "Security: Protect, investigate, and respond to cyber threats with AI-driven security analytics."
+        "Elasticsearch: Distributed, RESTful search and analytics.",
+        "Kibana: Visualize your data. Navigate the Stack.",
+        "Beats: Collect, parse, and ship in a lightweight fashion.",
+        "Connectors: Connect popular databases, file systems, collaboration tools, and more.",
+        "Logstash: Ingest, transform, enrich, and output.",
+    ]
+
+    # Generate embeddings for each product. Keep them in an array instead of a vector DB.
+    product_embeddings = []
+    for product in products:
+        product_embeddings.append(create_embedding(client, product))
+
+    query_embedding = create_embedding(client, "What can help me connect to a database?")
+
+    # Calculate cosine similarity between the query and document embeddings
+    similarities = []
+    for product_embedding in product_embeddings:
+        similarity = np.dot(query_embedding, product_embedding) / (
+            np.linalg.norm(query_embedding) * np.linalg.norm(product_embedding)
+        )
+        similarities.append(similarity)
+
+    # Get the index of the most similar document
+    most_similar_index = np.argmax(similarities)
+
+    print(products[most_similar_index])
+
+
+def create_embedding(client, text):
+    return client.embeddings.create(input=[text], model=EMBEDDINGS_MODEL, encoding_format="float").data[0].embedding
+
+
+if __name__ == "__main__":
+    main()

instrumentation/elastic-opentelemetry-instrumentation-openai/ollama.env

@@ -0,0 +1,10 @@
+# Env to run the integration tests against a local Ollama.
+OPENAI_BASE_URL=http://127.0.0.1:11434/v1
+OPENAI_API_KEY=notused
+
+# These models may be substituted in the future with inexpensive to run, newer
+# variants.
+TEST_CHAT_MODEL=qwen2.5:0.5b
+TEST_EMBEDDINGS_MODEL=all-minilm:33m
+
+OTEL_SERVICE_NAME=elastic-opentelemetry-instrumentation-openai

instrumentation/elastic-opentelemetry-instrumentation-openai/tests/base.py

@@ -0,0 +1,191 @@
+import os
+from dataclasses import dataclass
+
+import openai
+from opentelemetry.instrumentation.openai import OpenAIInstrumentor
+from opentelemetry.metrics import Histogram
+from vcr.unittest import VCRMixin
+
+# Use the same model for tools as for chat completion
+OPENAI_API_KEY = "test_openai_api_key"
+OPENAI_ORG_ID = "test_openai_org_key"
+OPENAI_PROJECT_ID = "test_openai_project_id"
+
+LOCAL_MODEL = "qwen2.5:0.5b"
+
+
+@dataclass
+class OpenAIEnvironment:
+    # TODO: add system
+    operation_name: str = "chat"
+    model: str = "gpt-4o-mini"
+    response_model: str = "gpt-4o-mini-2024-07-18"
+    server_address: str = "api.openai.com"
+    server_port: int = 443
+
+
+class OpenaiMixin(VCRMixin):
+    def _get_vcr_kwargs(self, **kwargs):
+        """
+        This scrubs sensitive data and gunzips bodies when in recording mode.
+
+        Without this, you would leak cookies and auth tokens in the cassettes.
+        Also, depending on the request, some responses would be binary encoded
+        while others plain json. This ensures all bodies are human-readable.
+        """
+        return {
+            "decode_compressed_response": True,
+            "filter_headers": [
+                ("authorization", "Bearer " + OPENAI_API_KEY),
+                ("openai-organization", OPENAI_ORG_ID),
+                ("openai-project", OPENAI_PROJECT_ID),
+                ("cookie", None),
+            ],
+            "before_record_response": self.scrub_response_headers,
+        }
+
+    @staticmethod
+    def scrub_response_headers(response):
+        """
+        This scrubs sensitive response headers. Note they are case-sensitive!
+        """
+        response["headers"]["openai-organization"] = OPENAI_ORG_ID
+        response["headers"]["Set-Cookie"] = "test_set_cookie"
+        return response
+
+    @classmethod
+    def setup_client(cls):
+        # Control the arguments
+        return openai.Client(
+            api_key=os.getenv("OPENAI_API_KEY", OPENAI_API_KEY),
+            organization=os.getenv("OPENAI_ORG_ID", OPENAI_ORG_ID),
+            project=os.getenv("OPENAI_PROJECT_ID", OPENAI_PROJECT_ID),
+            max_retries=1,
+        )
+
+    @classmethod
+    def setup_environment(cls):
+        return OpenAIEnvironment()
+
+    @classmethod
+    def setUpClass(cls):
+        cls.client = cls.setup_client()
+        cls.openai_env = cls.setup_environment()
+
+    def setUp(self):
+        super().setUp()
+        OpenAIInstrumentor().instrument()
+
+    def tearDown(self):
+        super().tearDown()
+        OpenAIInstrumentor().uninstrument()
+
+    def assertOperationDurationMetric(self, metric: Histogram):
+        self.assertEqual(metric.name, "gen_ai.client.operation.duration")
+        self.assert_metric_expected(
+            metric,
+            [
+                self.create_histogram_data_point(
+                    count=1,
+                    sum_data_point=0.006543334107846022,
+                    max_data_point=0.006543334107846022,
+                    min_data_point=0.006543334107846022,
+                    attributes={
+                        "gen_ai.operation.name": self.openai_env.operation_name,
+                        "gen_ai.request.model": self.openai_env.model,
+                        "gen_ai.response.model": self.openai_env.response_model,
+                        "gen_ai.system": "openai",
+                        "server.address": self.openai_env.server_address,
+                        "server.port": self.openai_env.server_port,
+                    },
+                ),
+            ],
+            est_value_delta=0.2,
+        )
+
+    def assertErrorOperationDurationMetric(self, metric: Histogram, attributes: dict, data_point: float = None):
+        self.assertEqual(metric.name, "gen_ai.client.operation.duration")
+        default_attributes = {
+            "gen_ai.operation.name": self.openai_env.operation_name,
+            "gen_ai.request.model": self.openai_env.model,
+            "gen_ai.system": "openai",
+            "error.type": "APIConnectionError",
+            "server.address": "localhost",
+            "server.port": 9999,
+        }
+        if data_point is None:
+            data_point = 0.8643839359283447
+        self.assert_metric_expected(
+            metric,
+            [
+                self.create_histogram_data_point(
+                    count=1,
+                    sum_data_point=data_point,
+                    max_data_point=data_point,
+                    min_data_point=data_point,
+                    attributes={**default_attributes, **attributes},
+                ),
+            ],
+            est_value_delta=0.5,
+        )
+
+    def assertTokenUsageInputMetric(self, metric: Histogram, input_data_point=4):
+        self.assertEqual(metric.name, "gen_ai.client.token.usage")
+        self.assert_metric_expected(
+            metric,
+            [
+                self.create_histogram_data_point(
+                    count=1,
+                    sum_data_point=input_data_point,
+                    max_data_point=input_data_point,
+                    min_data_point=input_data_point,
+                    attributes={
+                        "gen_ai.operation.name": self.openai_env.operation_name,
+                        "gen_ai.request.model": self.openai_env.model,
+                        "gen_ai.response.model": self.openai_env.response_model,
+                        "gen_ai.system": "openai",
+                        "server.address": self.openai_env.server_address,
+                        "server.port": self.openai_env.server_port,
+                        "gen_ai.token.type": "input",
+                    },
+                ),
+            ],
+        )
+
+    def assertTokenUsageMetric(self, metric: Histogram, input_data_point=24, output_data_point=4):
+        self.assertEqual(metric.name, "gen_ai.client.token.usage")
+        self.assert_metric_expected(
+            metric,
+            [
+                self.create_histogram_data_point(
+                    count=1,
+                    sum_data_point=input_data_point,
+                    max_data_point=input_data_point,
+                    min_data_point=input_data_point,
+                    attributes={
+                        "gen_ai.operation.name": self.openai_env.operation_name,
+                        "gen_ai.request.model": self.openai_env.model,
+                        "gen_ai.response.model": self.openai_env.response_model,
+                        "gen_ai.system": "openai",
+                        "server.address": self.openai_env.server_address,
+                        "server.port": self.openai_env.server_port,
+                        "gen_ai.token.type": "input",
+                    },
+                ),
+                self.create_histogram_data_point(
+                    count=1,
+                    sum_data_point=output_data_point,
+                    max_data_point=output_data_point,
+                    min_data_point=output_data_point,
+                    attributes={
+                        "gen_ai.operation.name": self.openai_env.operation_name,
+                        "gen_ai.request.model": self.openai_env.model,
+                        "gen_ai.response.model": self.openai_env.response_model,
+                        "gen_ai.system": "openai",
+                        "server.address": self.openai_env.server_address,
+                        "server.port": self.openai_env.server_port,
+                        "gen_ai.token.type": "output",
+                    },
+                ),
+            ],
+        )