From 12b6712b281d77d2a92df29dc45cfdf684acad8a Mon Sep 17 00:00:00 2001
From: Trayan Azarov <tazarov@users.noreply.github.com>
Date: Mon, 13 Nov 2023 08:50:27 -0800
Subject: [PATCH] [ENH]: FastAPI Instrumentation for improved traceability
 (#1281)

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
	 - ...
 - New functionality
- Added FastAPI OTEL instrumentation; this leads to much more consistent
and cohesive traces (before, we had some disjoint traces as the context
was not contiguous, from top to bottom)

## Test plan
*How are these changes tested?*

- [ ] Tests pass locally with `pytest` for python, `yarn test` for js

## Documentation Changes
N/A


![image](https://github.com/chroma-core/chroma/assets/1157440/74fe9603-2f57-42ff-b1c4-d93744aba985)
---
 chromadb/server/fastapi/__init__.py         |  2 ++
 chromadb/telemetry/opentelemetry/fastapi.py | 10 ++++++++++
 pyproject.toml                              |  1 +
 requirements.txt                            |  1 +
 4 files changed, 14 insertions(+)
 create mode 100644 chromadb/telemetry/opentelemetry/fastapi.py

diff --git a/chromadb/server/fastapi/__init__.py b/chromadb/server/fastapi/__init__.py
index 6fed3088d9b..ba56f7fddb2 100644
--- a/chromadb/server/fastapi/__init__.py
+++ b/chromadb/server/fastapi/__init__.py
@@ -50,6 +50,7 @@
 from starlette.requests import Request
 
 import logging
+from chromadb.telemetry.opentelemetry.fastapi import instrument_fastapi
 from chromadb.types import Database, Tenant
 from chromadb.telemetry.product import ServerContext, ProductTelemetryClient
 from chromadb.telemetry.opentelemetry import (
@@ -266,6 +267,7 @@ def __init__(self, settings: Settings):
         self._app.include_router(self.router)
 
         use_route_names_as_operation_ids(self._app)
+        instrument_fastapi(self._app)
 
     def app(self) -> fastapi.FastAPI:
         return self._app
diff --git a/chromadb/telemetry/opentelemetry/fastapi.py b/chromadb/telemetry/opentelemetry/fastapi.py
new file mode 100644
index 00000000000..34825794555
--- /dev/null
+++ b/chromadb/telemetry/opentelemetry/fastapi.py
@@ -0,0 +1,10 @@
+from typing import List, Optional
+from fastapi import FastAPI
+from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+
+
+def instrument_fastapi(app: FastAPI, excluded_urls: Optional[List[str]] = None) -> None:
+    """Instrument FastAPI to emit OpenTelemetry spans."""
+    FastAPIInstrumentor.instrument_app(
+        app, excluded_urls=",".join(excluded_urls) if excluded_urls else None
+    )
diff --git a/pyproject.toml b/pyproject.toml
index d223ea338e0..21a795e936d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -28,6 +28,7 @@ dependencies = [
   'onnxruntime >= 1.14.1',
   'opentelemetry-api>=1.2.0',
   'opentelemetry-exporter-otlp-proto-grpc>=1.2.0',
+  'opentelemetry-instrumentation-fastapi>=0.41b0',
   'opentelemetry-sdk>=1.2.0',
   'tokenizers >= 0.13.2',
   'pypika >= 0.48.9',
diff --git a/requirements.txt b/requirements.txt
index 338083ef3d0..c23ddcc412e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,6 +10,7 @@ numpy>=1.22.4; python_version >= '3.8'
 onnxruntime>=1.14.1
 opentelemetry-api>=1.2.0
 opentelemetry-exporter-otlp-proto-grpc>=1.2.0
+opentelemetry-instrumentation-fastapi>=0.41b0
 opentelemetry-sdk>=1.2.0
 overrides==7.3.1
 posthog==2.4.0