feat: observability for Python with OpenTelemetry

run-llama · Apr 5, 2024 · e55fc26 · e55fc26
1 parent 83e7b65
commit e55fc26
Show file tree

Hide file tree

Showing 5 changed files with 50 additions and 21 deletions.
diff --git a/helpers/python.ts b/helpers/python.ts
@@ -206,6 +206,7 @@ export const installPythonTemplate = async ({
   tools,
   postInstallAction,
   useLlamaParse,
+  observability,
 }: Pick<
   InstallTemplateArgs,
   | "root"
@@ -216,6 +217,7 @@ export const installPythonTemplate = async ({
   | "tools"
   | "useLlamaParse"
   | "postInstallAction"
+  | "observability"
 >) => {
   console.log("\nInitializing Python project with template:", template, "\n");
   const templatePath = path.join(templatesDir, "types", template, framework);
@@ -258,6 +260,26 @@ export const installPythonTemplate = async ({
   const addOnDependencies = dataSources
     .map((ds) => getAdditionalDependencies(vectorDb, ds, tools))
     .flat();
+
+  if (observability === "opentelemetry") {
+    addOnDependencies.push({
+      name: "traceloop-sdk",
+      version: "^0.5.11",
+    });
+
+    const observabilityPath = path.join(
+      root,
+      templatesDir,
+      "components",
+      "observability",
+      "python",
+      "opentelemetry",
+    );
+    await copy("**", observabilityPath, {
+      cwd: path.join(root, "app", "observability"),
+    });
+  }
+
   await addDependencies(root, addOnDependencies);
 
   if (postInstallAction === "runApp" || postInstallAction === "dependencies") {

diff --git a/questions.ts b/questions.ts
@@ -510,28 +510,26 @@ export const askQuestions = async (
     }
   }
 
-  if (program.framework === "express" || program.framework === "nextjs") {
-    if (!program.observability) {
-      if (ciInfo.isCI) {
-        program.observability = getPrefOrDefault("observability");
-      } else {
-        const { observability } = await prompts(
-          {
-            type: "select",
-            name: "observability",
-            message: "Would you like to set up observability?",
-            choices: [
-              { title: "No", value: "none" },
-              { title: "OpenTelemetry", value: "opentelemetry" },
-            ],
-            initial: 0,
-          },
-          handlers,
-        );
+  if (!program.observability) {
+    if (ciInfo.isCI) {
+      program.observability = getPrefOrDefault("observability");
+    } else {
+      const { observability } = await prompts(
+        {
+          type: "select",
+          name: "observability",
+          message: "Would you like to set up observability?",
+          choices: [
+            { title: "No", value: "none" },
+            { title: "OpenTelemetry", value: "opentelemetry" },
+          ],
+          initial: 0,
+        },
+        handlers,
+      );
 
-        program.observability = observability;
-        preferences.observability = observability;
-      }
+      program.observability = observability;
+      preferences.observability = observability;
     }
   }
 

diff --git a/templates/components/observability/python/opentelemetry/observability.py b/templates/components/observability/python/opentelemetry/observability.py
@@ -0,0 +1,5 @@
+from traceloop.sdk import Traceloop
+
+
+def init_observability():
+    Traceloop.init()
diff --git a/templates/types/streaming/fastapi/app/observability.py b/templates/types/streaming/fastapi/app/observability.py
@@ -0,0 +1,2 @@
+def init_observability():
+    pass
diff --git a/templates/types/streaming/fastapi/main.py b/templates/types/streaming/fastapi/main.py
@@ -10,11 +10,13 @@
 from fastapi.responses import RedirectResponse
 from app.api.routers.chat import chat_router
 from app.settings import init_settings
+from app.observability import init_observability
 
 
 app = FastAPI()
 
 init_settings()
+init_observability()
 
 environment = os.getenv("ENVIRONMENT", "dev")  # Default to 'development' if not set