Add: Elastic apm

weni-ai · Jun 13, 2024 · c80b265 · c80b265
1 parent c9e1d65
commit c80b265
Show file tree

Hide file tree

Showing 3 changed files with 137 additions and 4 deletions.
diff --git a/app/main.py b/app/main.py
@@ -1,5 +1,7 @@
 import sentry_sdk
+import os
 from elasticsearch import Elasticsearch
+from elasticapm.contrib.starlette import ElasticAPM, make_apm_client
 from fastapi import FastAPI
 from langchain.embeddings import HuggingFaceHubEmbeddings, CohereEmbeddings
 from langchain.embeddings.base import Embeddings
@@ -88,6 +90,18 @@ def __init__(self, config: AppConfig):
         self.content_base_handler = ContentBaseHandler(self.content_base_indexer)
         self.api.include_router(self.content_base_handler.router)
 
+        # APM Configuration
+
+        apm_config = {
+            'SERVICE_NAME': os.environ.get('APM_SERVICE_NAME'),
+            'SECRET_TOKEN': os.environ.get('APM_SECRET_TOKEN'),
+            'SERVER_URL': os.environ.get('APM_SERVER_URL'),
+            'ENVIRONMENT': os.environ.get('APM_ENVIRONMENT'),
+        }
+
+        apm_CLIENT = make_apm_client(apm_config)
+        self.api.add_middleware(ElasticAPM, client=apm_CLIENT)
+
 
 config = AppConfig()
 main_app = App(config)

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -30,6 +30,7 @@ sentry-sdk = {extras = ["fastapi"], version = "^1.35.0"}
 xlrd = "^2.0.1"
 pdfminer-six = "^20231228"
 cohere = "^4.37"
+elastic-apm = "^6.22.3"
 
 
 [tool.poetry.group.dev.dependencies]