Added new events for chat completions (#11)

* Added new events for chat completions

* store headers by completion key

* Add "is_patched" attribute to patched functions

* add request.model and response.model to the events

* change rate-limit values to numbers from strings

* generate events for requests with errors

* Added LlmEmbedding

* add support for async create

* clean up get_rate_limit_data

* log exception on the root level instead of raising

* Change error handling

* move handling to decorator
* raise exceptions from openai's original methods

* added ingest source to events

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* bugfix:  stream=true was breaking client apps

* fix stream=true for async functions

* calculate response time internally

* Limit message content to 4095 chars

* add message uuid

* added support for "engine" instead of "model"

---------

Co-authored-by: ykriger-newrelic <[email protected]>

.gitignore

@@ -0,0 +1,150 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# poetry
+poetry.lock
+
+# vscode
+.vscode
+
+#pyenv
+.python-version
+
+#git
+*.patch

README.md

@@ -39,7 +39,9 @@ os.environ["NEW_RELIC_LICENSE_KEY"] = "<license key>"
 
 ```python
 from nr_openai_observability import monitor
-monitor.initialization()
+monitor.initialization(
+    application_name="OpenAI observability example"
+)
 ```
 
 #### Code example:
@@ -51,15 +53,21 @@ import os
 import openai
 from nr_openai_observability import monitor
 
-monitor.initialization()
+monitor.initialization(
+    application_name="OpenAI observability example"
+)
 
 openai.api_key = os.getenv("OPENAI_API_KEY")
-openai.Completion.create(
-    model="text-davinci-003",
-    prompt="What is Observability?",
-    max_tokens=20,
-    temperature=0 
+response = openai.ChatCompletion.create(
+    model="gpt-3.5-turbo",
+    messages=[
+        {
+            "role": "user",
+            "content": "Write a rhythm about observability",
+        },
+    ],
 )
+print(response["choices"][0]["message"]["content"])
 ```
 
 #### STEP 3: Follow the instruction [here](https://one.newrelic.com/launcher/catalog-pack-details.launcher/?pane=eyJuZXJkbGV0SWQiOiJjYXRhbG9nLXBhY2stZGV0YWlscy5jYXRhbG9nLXBhY2stY29udGVudHMiLCJxdWlja3N0YXJ0SWQiOiI1ZGIyNWRiZC1hNmU5LTQ2ZmMtYTcyOC00Njk3ZjY3N2ZiYzYifQ==) to add the dashboard to your New Relic account.

src/nr_openai_observability/build_events.py

@@ -0,0 +1,150 @@
+import uuid
+from datetime import datetime
+
+import openai
+
+
+def _build_messages_events(messages, completion_id, model):
+    message_id = str(uuid.uuid4())
+    events = []
+    for index, message in enumerate(messages):
+        currMessage = {
+            "id": message_id,
+            "content": message.get("content")[:4095],
+            "role": message.get("role"),
+            "completion_id": completion_id,
+            "sequence": index,
+            "model": model,
+            "vendor": "openAI",
+            "ingest_source": "PythonSDK",
+        }
+
+        events.append(currMessage)
+
+    return events
+
+
+def _get_rate_limit_data(response_headers):
+    def _get_numeric_header(name):
+        header = response_headers.get(name)
+        return int(header) if header and header.isdigit() else None
+
+    return {
+        "ratelimit_limit_requests": _get_numeric_header("ratelimit_limit_requests"),
+        "ratelimit_limit_tokens": _get_numeric_header("ratelimit_limit_tokens"),
+        "ratelimit_reset_tokens": response_headers.get("x-ratelimit-reset-tokens"),
+        "ratelimit_reset_requests": response_headers.get("x-ratelimit-reset-requests"),
+        "ratelimit_remaining_tokens": _get_numeric_header("ratelimit_remaining_tokens"),
+        "ratelimit_remaining_requests": _get_numeric_header(
+            "ratelimit_remaining_requests"
+        ),
+    }
+
+
+def build_completion_events(response, request, response_headers, response_time):
+    completion_id = str(uuid.uuid4())
+
+    completion = {
+        "id": completion_id,
+        "api_key_last_four_digits": f"sk-{response.api_key[-4:]}",
+        "timestamp": datetime.now(),
+        "response_time": int(response_time * 1000),
+        "request.model": request.get("model") or request.get("engine"),
+        "response.model": response.model,
+        "usage.completion_tokens": response.usage.completion_tokens,
+        "usage.total_tokens": response.usage.total_tokens,
+        "usage.prompt_tokens": response.usage.prompt_tokens,
+        "temperature": request.get("temperature"),
+        "max_tokens": request.get("max_tokens"),
+        "finish_reason": response.choices[0].finish_reason,
+        "api_type": response.api_type,
+        "vendor": "openAI",
+        "ingest_source": "PythonSDK",
+        "number_of_messages": len(request.get("messages", [])) + len(response.choices),
+        "organization": response.organization,
+        "api_version": response_headers.get("openai-version"),
+    }
+
+    completion.update(_get_rate_limit_data(response_headers))
+
+    messages = _build_messages_events(
+        request.get("messages", []) + [response.choices[0].message],
+        completion_id,
+        response.model,
+    )
+
+    return {"messages": messages, "completion": completion}
+
+
+def build_completion_error_events(request, error):
+    completion_id = str(uuid.uuid4())
+
+    completion = {
+        "id": completion_id,
+        "api_key_last_four_digits": f"sk-{openai.api_key[-4:]}",
+        "timestamp": datetime.now(),
+        "request.model": request.get("model") or request.get("engine"),
+        "temperature": request.get("temperature"),
+        "max_tokens": request.get("max_tokens"),
+        "vendor": "openAI",
+        "ingest_source": "PythonSDK",
+        "organization": error.organization,
+        "number_of_messages": len(request.get("messages", [])),
+        "error_status": error.http_status,
+        "error_message": error.error.message,
+        "error_type": error.error.type,
+        "error_code": error.error.code,
+        "error_param": error.error.param,
+    }
+
+    messages = _build_messages_events(
+        request.get("messages", []),
+        completion_id,
+        request.get("model") or request.get("engine"),
+    )
+
+    return {"messages": messages, "completion": completion}
+
+
+def build_embedding_event(response, request, response_headers, response_time):
+    embedding_id = str(uuid.uuid4())
+
+    embedding = {
+        "id": embedding_id,
+        "api_key_last_four_digits": f"sk-{response.api_key[-4:]}",
+        "timestamp": datetime.now(),
+        "response_time": int(response_time * 1000),
+        "request.model": request.get("model") or request.get("engine"),
+        "response.model": response.model,
+        "usage.total_tokens": response.usage.total_tokens,
+        "usage.prompt_tokens": response.usage.prompt_tokens,
+        "api_type": response.api_type,
+        "vendor": "openAI",
+        "ingest_source": "PythonSDK",
+        "organization": response.organization,
+        "api_version": response_headers.get("openai-version"),
+    }
+
+    embedding.update(_get_rate_limit_data(response_headers))
+    return embedding
+
+
+def build_embedding_error_event(request, error):
+    embedding_id = str(uuid.uuid4())
+
+    embedding = {
+        "id": embedding_id,
+        "api_key_last_four_digits": f"sk-{openai.api_key[-4:]}",
+        "timestamp": datetime.now(),
+        "request.model": request.get("model") or request.get("engine"),
+        "vendor": "openAI",
+        "ingest_source": "PythonSDK",
+        "organization": error.organization,
+        "error_status": error.http_status,
+        "error_message": error.error.message,
+        "error_type": error.error.type,
+        "error_code": error.error.code,
+        "error_param": error.error.param,
+    }
+
+    return embedding

src/nr_openai_observability/error_handling_decorator.py

@@ -0,0 +1,14 @@
+import logging
+from functools import wraps
+
+logger = logging.getLogger("nr_openai_observability")
+
+
+def handle_errors(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as err:
+            logger.error(f"An error occurred in {func.__name__}: {err}")
+    return wrapper