feat: report custom pre/post processing

traceloop · Jan 16, 2024 · ad63d87 · ad63d87
1 parent 018b242
commit ad63d87
Show file tree

Hide file tree

Showing 7 changed files with 79 additions and 10 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/ai-semantic-conventions/src/SemanticAttributes.ts b/packages/ai-semantic-conventions/src/SemanticAttributes.ts
@@ -40,6 +40,11 @@ export const SpanAttributes = {
   TRACELOOP_WORKFLOW_NAME: "traceloop.workflow.name",
   TRACELOOP_ENTITY_NAME: "traceloop.entity.name",
   TRACELOOP_ASSOCIATION_PROPERTIES: "traceloop.association.properties",
+
+  // LLM Pre/Post Processing
+  TRACELOOP_CUSTOM_PROCESSING_STEP_TYPE: "traceloop.processing.type",
+  TRACELOOP_CUSTOM_PROCESSING_INPUT: "traceloop.processing.input",
+  TRACELOOP_CUSTOM_PROCESSING_OUTPUT: "traceloop.processing.output",
 };
 
 export enum LLMRequestTypeValues {
@@ -56,3 +61,8 @@ export enum TraceloopSpanKindValues {
   TOOL = "tool",
   UNKNOWN = "unknown",
 }
+
+export enum TraceloopCustomProcessingStepTypeValues {
+  PRE_PROCESSING = "pre_process",
+  POST_PROCESSING = "post_process",
+}
diff --git a/packages/sample-app/src/sample_with.ts b/packages/sample-app/src/sample_with.ts
@@ -12,12 +12,24 @@ async function chat() {
   return await traceloop.withWorkflow("sample_chat", async () => {
     const chatCompletion = await openai.chat.completions.create({
       messages: [
-        { role: "user", content: "Tell me a joke about OpenTelemetry" },
+        {
+          role: "user",
+          content:
+            "Tell me a joke about OpenTelemetry and explain in a separate line",
+        },
       ],
       model: "gpt-3.5-turbo",
     });
 
-    return chatCompletion.choices[0].message.content;
+    const jokeWithExplanation = chatCompletion.choices[0].message.content;
+    const onlyExplanation = jokeWithExplanation?.split("Explanation:")[1];
+
+    traceloop.reportPostProcessing(
+      jokeWithExplanation || "",
+      onlyExplanation || "",
+    );
+
+    return onlyExplanation;
   });
 }
 

diff --git a/packages/traceloop-sdk/package.json b/packages/traceloop-sdk/package.json
@@ -34,8 +34,9 @@
   "dependencies": {
     "@opentelemetry/exporter-trace-otlp-proto": "^0.44.0",
     "@opentelemetry/sdk-node": "^0.44.0",
-    "@traceloop/instrumentation-llamaindex": "^0.0.31",
-    "@traceloop/instrumentation-openai": "^0.0.31",
+    "@traceloop/ai-semantic-conventions": "*",
+    "@traceloop/instrumentation-llamaindex": "*",
+    "@traceloop/instrumentation-openai": "*",
     "@types/nunjucks": "^3.2.5",
     "fetch-retry": "^5.0.6",
     "nunjucks": "^3.2.4",

diff --git a/packages/traceloop-sdk/src/lib/node-server-sdk.ts b/packages/traceloop-sdk/src/lib/node-server-sdk.ts
@@ -3,7 +3,11 @@ import { initInstrumentations } from "./tracing";
 export * from "./errors";
 export { InitializeOptions } from "./interfaces";
 export { initialize } from "./configuration";
-export { forceFlush } from "./tracing";
+export {
+  forceFlush,
+  reportPreProcessing,
+  reportPostProcessing,
+} from "./tracing";
 export * from "./tracing/decorators";
 export * from "./tracing/association";
 export * from "./tracing/score";

diff --git a/packages/traceloop-sdk/src/lib/tracing/index.ts b/packages/traceloop-sdk/src/lib/tracing/index.ts
@@ -11,8 +11,15 @@ import { Instrumentation } from "@opentelemetry/instrumentation";
 import { InitializeOptions } from "../interfaces";
 import { OpenAIInstrumentation } from "@traceloop/instrumentation-openai";
 import { LlamaIndexInstrumentation } from "@traceloop/instrumentation-llamaindex";
-import { SpanAttributes } from "@traceloop/ai-semantic-conventions";
-import { ASSOCATION_PROPERTIES_KEY, WORKFLOW_NAME_KEY } from "./tracing";
+import {
+  SpanAttributes,
+  TraceloopCustomProcessingStepTypeValues,
+} from "@traceloop/ai-semantic-conventions";
+import {
+  ASSOCATION_PROPERTIES_KEY,
+  WORKFLOW_NAME_KEY,
+  reportCustomProcessing,
+} from "./tracing";
 import { Telemetry } from "../telemetry/telemetry";
 import { TraceloopSampler } from "./sampler";
 
@@ -113,3 +120,19 @@ export const startTracing = (options: InitializeOptions) => {
 export const forceFlush = async () => {
   await _spanProcessor.forceFlush();
 };
+
+export const reportPreProcessing = (input: string, output: string) => {
+  reportCustomProcessing(
+    TraceloopCustomProcessingStepTypeValues.PRE_PROCESSING,
+    input,
+    output,
+  );
+};
+
+export const reportPostProcessing = (input: string, output: string) => {
+  reportCustomProcessing(
+    TraceloopCustomProcessingStepTypeValues.POST_PROCESSING,
+    input,
+    output,
+  );
+};
diff --git a/packages/traceloop-sdk/src/lib/tracing/tracing.ts b/packages/traceloop-sdk/src/lib/tracing/tracing.ts
@@ -1,5 +1,8 @@
 import { trace, createContextKey } from "@opentelemetry/api";
-
+import {
+  SpanAttributes,
+  TraceloopCustomProcessingStepTypeValues,
+} from "@traceloop/ai-semantic-conventions";
 const TRACER_NAME = "traceloop.tracer";
 export const WORKFLOW_NAME_KEY = createContextKey("workflow_name");
 export const ASSOCATION_PROPERTIES_KEY = createContextKey(
@@ -9,3 +12,18 @@ export const ASSOCATION_PROPERTIES_KEY = createContextKey(
 export const getTracer = () => {
   return trace.getTracer(TRACER_NAME);
 };
+
+export const reportCustomProcessing = (
+  kind: TraceloopCustomProcessingStepTypeValues,
+  input: string,
+  output: string,
+) => {
+  const s = getTracer().startSpan(`custom.${kind}`, {
+    attributes: {
+      [SpanAttributes.TRACELOOP_CUSTOM_PROCESSING_STEP_TYPE]: kind,
+      [SpanAttributes.TRACELOOP_CUSTOM_PROCESSING_INPUT]: input,
+      [SpanAttributes.TRACELOOP_CUSTOM_PROCESSING_OUTPUT]: output,
+    },
+  });
+  s.end();
+};