From 09810ecd050326f7539aa3ad9b7149d9b1190ed0 Mon Sep 17 00:00:00 2001
From: Rathijit Paul <30369246+rathijitpapon@users.noreply.github.com>
Date: Thu, 27 Jun 2024 20:31:02 +0600
Subject: [PATCH] :construction: added jaeger in agency

---
 agency/.env.template        |  4 +++-
 agency/app/settings.py      |  4 +++-
 agency/app/tracing/utils.py | 38 ++++++++++++++++++++++++++++++-------
 3 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/agency/.env.template b/agency/.env.template
index b8b53960..dead636e 100644
--- a/agency/.env.template
+++ b/agency/.env.template
@@ -81,7 +81,9 @@ TRACING__SENTRY_DSN=
 # TRACING__ENABLE_TRACING=
 # TRACING__ENVIRONMENT=
 # TRACING__PHOENIX_API=
-# TRACING__PHOENIX_PROJECT_NAME=
+# TRACING__PROJECT_NAME=
+# TRACING__SERVIVE_NAME=
+# TRACING__JAEGER_ENDPOINT=
 
 # GROQ API Configuration
 GROQ__API_KEY=
diff --git a/agency/app/settings.py b/agency/app/settings.py
index 53208f9c..50f03260 100644
--- a/agency/app/settings.py
+++ b/agency/app/settings.py
@@ -115,10 +115,12 @@ class WandbSettings(BaseSettings):
 
 class TracingSettings(BaseSettings):
     sentry_dsn: SecretStr
+    jaeger_endpoint: str = "http://127.0.0.1:4317"
     enable_tracing: bool = False
     environment: str = "development"
     phoenix_api: str = "http://127.0.0.1:6006/v1/traces"
-    phoenix_project_name: str = "Curieo Search Agency"
+    project_name: str = "Curieo Search"
+    service_name: str = "agency-service"
 
 
 class GroqSettings(BaseSettings):
diff --git a/agency/app/tracing/utils.py b/agency/app/tracing/utils.py
index 51befbe3..88bd23c3 100644
--- a/agency/app/tracing/utils.py
+++ b/agency/app/tracing/utils.py
@@ -3,11 +3,19 @@
 import sentry_sdk
 from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
 from openinference.semconv.resource import ResourceAttributes
+from opentelemetry import trace
 from opentelemetry import trace as trace_api
-from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+from opentelemetry.instrumentation.grpc import GrpcAioInstrumentorServer
+from opentelemetry.instrumentation.llamaindex import (
+    LlamaIndexInstrumentor as OpentelemetryLlamaIndexInstrumentor,
+)
+from opentelemetry.propagate import set_global_textmap
 from opentelemetry.sdk import trace as trace_sdk
 from opentelemetry.sdk.resources import Resource
-from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 from sentry_sdk.integrations.asyncio import AsyncioIntegration
 from sentry_sdk.integrations.grpc import GRPCIntegration
 from sentry_sdk.integrations.logging import LoggingIntegration
@@ -17,6 +25,26 @@
 
 def setup_tracing(settings: TracingSettings) -> None:
     if settings.environment == "production":
+        # jaeger opentelemetry tracing
+        resource = Resource(
+            attributes={
+                ResourceAttributes.PROJECT_NAME: settings.project_name,
+                "service.name": settings.service_name,
+            },
+        )
+        trace.set_tracer_provider(TracerProvider(resource=resource))
+        trace.get_tracer_provider().add_span_processor(
+            BatchSpanProcessor(
+                OTLPSpanExporter(
+                    endpoint=settings.jaeger_endpoint,
+                )
+            )
+        )
+        set_global_textmap(TraceContextTextMapPropagator())
+        GrpcAioInstrumentorServer().instrument()
+        OpentelemetryLlamaIndexInstrumentor().instrument()
+
+        # sentry tracing
         sentry_sdk.init(
             dsn=settings.sentry_dsn.get_secret_value(),
             enable_tracing=settings.enable_tracing,
@@ -27,11 +55,7 @@ def setup_tracing(settings: TracingSettings) -> None:
             ],
         )
 
-        resource = Resource(
-            attributes={
-                ResourceAttributes.PROJECT_NAME: settings.phoenix_project_name,
-            },
-        )
+        # phoenix tracing
         tracer_provider = trace_sdk.TracerProvider(resource=resource)
         span_exporter = OTLPSpanExporter(endpoint=settings.phoenix_api)
         span_processor = SimpleSpanProcessor(span_exporter=span_exporter)