Merge pull request #18 from grafana/fix-17

Add fixes for loki connection issue #17

Author: ishanjainn

.github/workflows/python-tests.yml

@@ -10,7 +10,7 @@ on:
     branches: [ "main" ]
   workflow_dispatch:
   schedule:
-    - cron: '*/180 * * * *'
+    - cron: '1 1 * * *'
 
 defaults:
   run:

python/src/grafana_openai_monitoring/__handlers.py

@@ -18,9 +18,12 @@
     __send_logs(logs_url, metrics_username, access_token, metrics):
         Send logs to the specified Loki URL.
 """
-
+import logging
 import requests
 
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.ERROR)
+
 # Function to check if all required arguments are provided and modify metrics and logs URLs
 def __check(metrics_url, logs_url, metrics_username, logs_username, access_token):
     # Check if all required parameters are provided
@@ -92,9 +95,8 @@ def __send_logs(logs_url, logs_username, access_token, logs):
                                  timeout=60
                                 )
         response.raise_for_status()  # Raise an exception for non-2xx status codes
-        return response
     except requests.exceptions.RequestException as err:
-        raise requests.exceptions.RequestException(f"Error sending Logs: {err}")
+        logger.error("Error sending Logs: %s", err)
 
 # Function to send metrics to the specified metrics URL
 def __send_metrics(metrics_url, metrics_username, access_token, metrics):
@@ -107,6 +109,5 @@ def __send_metrics(metrics_url, metrics_username, access_token, metrics):
                                  timeout=60
                             )
         response.raise_for_status()  # Raise an exception for non-2xx status codes
-        return response
     except requests.exceptions.RequestException as err:
-        raise requests.exceptions.RequestException(f"Error sending Metrics: {err}")
+        logger.error("Error sending Metrics: %s", err)

python/src/grafana_openai_monitoring/chat_v1.py

@@ -12,7 +12,9 @@
 from .__handlers import __send_metrics, __send_logs, __calculate_cost, __check
 
 # Decorator function to monitor chat completion
-def monitor(func, metrics_url, logs_url, metrics_username, logs_username, access_token): # pylint: disable=too-many-arguments
+# pylint: disable=too-many-arguments
+def monitor(func, metrics_url, logs_url, metrics_username, logs_username, access_token,
+            disable_content=False, environment="default"):
     """
     A decorator function to monitor chat completions using the OpenAI API.
 
@@ -72,7 +74,8 @@ def wrapper(*args, **kwargs):
             {
                     "stream": {
                         "job": "integrations/openai", 
-                        "prompt": prompt, 
+                        "prompt": prompt,
+                        "environment": environment,
                         "model": response.model, 
                         "finish_reason": response.choices[0].finish_reason,
                         "prompt_tokens": str(response.usage.prompt_tokens), 
@@ -90,38 +93,39 @@ def wrapper(*args, **kwargs):
             ]
         }
 
-        # Send logs to the specified logs URL
-        __send_logs(logs_url=logs_url,
-                    logs_username=logs_username,
-                    access_token=access_token,
-                    logs=logs
-        )
+        if disable_content is True:
+            # Send logs to the specified logs URL
+            __send_logs(logs_url=logs_url,
+                        logs_username=logs_username,
+                        access_token=access_token,
+                        logs=logs
+            )
 
         # Prepare metrics to be sent
         metrics = [
             # Metric to track the number of completion tokens used in the response
             f'openai,job=integrations/openai,'
-            f'source=python_chatv1,model={response.model} '
+            f'source=python_chatv1,model={response.model},environment={environment} '
             f'completionTokens={response.usage.completion_tokens}',
 
             # Metric to track the number of prompt tokens used in the response
             f'openai,job=integrations/openai,'
-            f'source=python_chatv1,model={response.model} '
+            f'source=python_chatv1,model={response.model},environment={environment} '
             f'promptTokens={response.usage.prompt_tokens}',
 
             # Metric to track the total number of tokens used in the response
             f'openai,job=integrations/openai,'
-            f'source=python_chatv1,model={response.model} '
+            f'source=python_chatv1,model={response.model},environment={environment} '
             f'totalTokens={response.usage.total_tokens}',
 
             # Metric to track the duration of the API request and response cycle
             f'openai,job=integrations/openai,'
-            f'source=python_chatv1,model={response.model} '
+            f'source=python_chatv1,model={response.model},environment={environment} '
             f'requestDuration={duration}',
 
             # Metric to track the usage cost based on the model and token usage
             f'openai,job=integrations/openai,'
-            f'source=python_chatv1,model={response.model} '
+            f'source=python_chatv1,model={response.model},environment={environment} '
             f'usageCost={cost}',
         ]
 

python/src/grafana_openai_monitoring/chat_v2.py

@@ -11,7 +11,7 @@
 from .__handlers import __send_metrics, __send_logs, __calculate_cost, __check
 
 # Decorator function to monitor chat completion
-def monitor(func, metrics_url, logs_url, metrics_username, logs_username, access_token, use_async=False): # pylint: disable=too-many-arguments, line-too-long
+def monitor(func, metrics_url, logs_url, metrics_username, logs_username, access_token, use_async=False, disable_content=False, environment="default"): # pylint: disable=too-many-arguments, line-too-long
     """
     A decorator function to monitor chat completions using the OpenAI API.
 
@@ -72,7 +72,8 @@ async def async_wrapper(*args, **kwargs):
             {
                     "stream": {
                         "job": "integrations/openai", 
-                        "prompt": prompt, 
+                        "prompt": prompt,
+                        "environment": environment,
                         "model": response.model, 
                         "role": response.choices[0].message.role,
                         "finish_reason": response.choices[0].finish_reason,
@@ -90,39 +91,39 @@ async def async_wrapper(*args, **kwargs):
                 }
             ]
         }
-
-        # Send logs to the specified logs URL
-        __send_logs(logs_url=logs_url,
-                    logs_username=logs_username,
-                    access_token=access_token,
-                    logs=logs
-        )
+        if disable_content is True:
+            # Send logs to the specified logs URL
+            __send_logs(logs_url=logs_url,
+                        logs_username=logs_username,
+                        access_token=access_token,
+                        logs=logs
+            )
 
         # Prepare metrics to be sent
         metrics = [
             # Metric to track the number of completion tokens used in the response
             f'openai,job=integrations/openai,'
-            f'source=python_chatv2,model={response.model} '
+            f'source=python_chatv2,model={response.model},environment={environment} '
             f'completionTokens={response.usage.completion_tokens}',
 
             # Metric to track the number of prompt tokens used in the response
             f'openai,job=integrations/openai,'
-            f'source=python_chatv2,model={response.model} '
+            f'source=python_chatv2,model={response.model},environment={environment} '
             f'promptTokens={response.usage.prompt_tokens}',
 
             # Metric to track the total number of tokens used in the response
             f'openai,job=integrations/openai,'
-            f'source=python_chatv2,model={response.model} '
+            f'source=python_chatv2,model={response.model},environment={environment} '
             f'totalTokens={response.usage.total_tokens}',
 
             # Metric to track the usage cost based on the model and token usage
             f'openai,job=integrations/openai,'
-            f'source=python_chatv2,model={response.model} '
+            f'source=python_chatv2,model={response.model},environment={environment} '
             f'usageCost={cost}',
 
             # Metric to track the duration of the API request and response cycle
             f'openai,job=integrations/openai,'
-            f'source=python_chatv2,model={response.model} '
+            f'source=python_chatv2,model={response.model},environment={environment} '
             f'requestDuration={duration}',
         ]
 
@@ -168,7 +169,8 @@ def wrapper(*args, **kwargs):
                 {
                         "stream": {
                             "job": "integrations/openai", 
-                            "prompt": prompt, 
+                            "prompt": prompt,
+                            "environment": environment,
                             "model": response.model, 
                             "role": response.choices[0].message.role,
                             "finish_reason": response.choices[0].finish_reason,
@@ -186,39 +188,39 @@ def wrapper(*args, **kwargs):
                     }
                 ]
             }
-
-            # Send logs to the specified logs URL
-            __send_logs(logs_url=logs_url,
-                        logs_username=logs_username,
-                        access_token=access_token,
-                        logs=logs
-            )
+            if disable_content is True:
+                # Send logs to the specified logs URL
+                __send_logs(logs_url=logs_url,
+                            logs_username=logs_username,
+                            access_token=access_token,
+                            logs=logs
+                )
 
             # Prepare metrics to be sent
             metrics = [
                 # Metric to track the number of completion tokens used in the response
                 f'openai,job=integrations/openai,'
-                f'source=python_chatv2,model={response.model} '
+                f'source=python_chatv2,model={response.model},environment={environment} '
                 f'completionTokens={response.usage.completion_tokens}',
 
                 # Metric to track the number of prompt tokens used in the response
                 f'openai,job=integrations/openai,'
-                f'source=python_chatv2,model={response.model} '
+                f'source=python_chatv2,model={response.model},environment={environment} '
                 f'promptTokens={response.usage.prompt_tokens}',
 
                 # Metric to track the total number of tokens used in the response
                 f'openai,job=integrations/openai,'
-                f'source=python_chatv2,model={response.model} '
+                f'source=python_chatv2,model={response.model},environment={environment} '
                 f'totalTokens={response.usage.total_tokens}',
 
                 # Metric to track the usage cost based on the model and token usage
                 f'openai,job=integrations/openai,'
-                f'source=python_chatv2,model={response.model} '
+                f'source=python_chatv2,model={response.model},environment={environment} '
                 f'usageCost={cost}',
 
                 # Metric to track the duration of the API request and response cycle
                 f'openai,job=integrations/openai,'
-                f'source=python_chatv2,model={response.model} '
+                f'source=python_chatv2,model={response.model},environment={environment} '
                 f'requestDuration={duration}',
             ]
 

python/tests/test_chat_v1.py

@@ -38,6 +38,7 @@ def test_chat_v1():
     )
 
     # Now any call to openai.Completion.create will be automatically tracked
-    response = client.completions.create(model="davinci", prompt="Hello world", max_tokens=100)
+    response = client.completions.create(model="gpt-3.5-turbo-instruct",
+                                         prompt="Hello world", max_tokens=1)
 
     assert response.object == 'text_completion'

python/tests/test_chat_v2.py

@@ -39,7 +39,7 @@ def test_chat_v2():
 
     # Now any call to openai.ChatCompletion.create will be automatically tracked
     response = client.chat.completions.create(model="gpt-3.5-turbo",
-                                            max_tokens=100,
+                                            max_tokens=1,
                                             messages=
                                                 [
                                                     {