Add telemetry abstraction to remove Sentry dep (round 2) (#47)

* Add telemetry abstraction to remove Sentry dep

This is a simplified version of #46. It was easier to copy the important
parts and start clean than try and delete all of the generics code.

Replace the direct dependency on Sentry with the abstract `Telemetry`
interface from @transitive-bullshit. This should have no impact on
users who weren't using the Sentry integration, and requires passing
Sentry as an argument when creating a model instance.

```ts
import { ChatModel } from '@dexaai/dexter'
import * as Sentry from '@sentry/node'

const model = new ChatModel({ telemetry: Sentry })
```

* Rename `Telemetry.Base` to `Telemetry.Provider`

.github/actions/install-node-pnpm/action.yml

@@ -8,7 +8,7 @@ inputs:
     required: false
   pnpm-version:
     description: 'The version of pnpm to install'
-    default: '8'
+    default: '9'
     required: false
 
 runs:

examples/instrument.ts

@@ -0,0 +1,11 @@
+import * as Sentry from '@sentry/node';
+
+/**
+ * Configure Sentry for tracing.
+ * Ensure to call this before importing any other modules!
+ */
+Sentry.init({
+  dsn: process.env.SENTRY_DSN!,
+  tracesSampleRate: 1.0,
+  debug: true,
+});

examples/with-telemetry.ts

@@ -0,0 +1,19 @@
+import './instrument.js';
+import 'dotenv/config';
+import { ChatModel } from '@dexaai/dexter';
+import * as Sentry from '@sentry/node';
+
+const chatModel = new ChatModel({
+  // Send tracing data to Sentry
+  telemetry: Sentry,
+  params: { model: 'gpt-4o-mini' },
+});
+
+async function main() {
+  const result = await chatModel.run({
+    messages: [{ role: 'user', content: 'Tell me a short joke' }],
+  });
+  console.log(result);
+}
+
+main().catch(console.error);

package.json

@@ -98,8 +98,5 @@
     "vite": "^5.2.11",
     "vitest": "^1.6.0"
   },
-  "peerDependencies": {
-    "@sentry/node": "8.x"
-  },
   "packageManager": "[email protected]"
 }

src/index.ts

@@ -4,4 +4,5 @@ export * from './datastore/index.js';
 export * from './datastore/pinecone/index.js';
 export * from './model/index.js';
 export * from './prompt/index.js';
+export * from './telemetry/index.js';
 export * from './utils/index.js';

src/model/chat.ts

@@ -254,6 +254,7 @@ export class ChatModel<
       cache: this.cache,
       client: this.client,
       debug: this.debug,
+      telemetry: this.telemetry,
       ...args,
       params: deepMerge(this.params, args?.params),
       context:

src/model/completion.ts

@@ -87,6 +87,7 @@ export class CompletionModel<
       cache: this.cache,
       client: this.client,
       debug: this.debug,
+      telemetry: this.telemetry,
       ...args,
       context: deepMerge(this.context, args?.context),
       params: deepMerge(this.params, args?.params),

src/model/embedding.ts

@@ -179,6 +179,7 @@ export class EmbeddingModel<
       cache: this.cache,
       client: this.client,
       debug: this.debug,
+      telemetry: this.telemetry,
       ...args,
       context: deepMerge(this.context, args?.context),
       params: deepMerge(this.params, args?.params),

src/model/model.ts

@@ -1,5 +1,4 @@
 import type { PartialDeep } from 'type-fest';
-import * as Sentry from '@sentry/node';
 import {
   extractAttrsFromContext,
   extractAttrsFromParams,
@@ -14,6 +13,8 @@ import {
   type CacheStorage,
   defaultCacheKey,
 } from '../utils/cache.js';
+import type { Telemetry } from '../telemetry/types.js';
+import { DefaultTelemetry } from '../telemetry/default-telemetry.js';
 
 export interface ModelArgs<
   MClient extends Model.Base.Client,
@@ -40,6 +41,7 @@ export interface ModelArgs<
   context?: Ctx;
   params: MConfig & Partial<MRun>;
   events?: Model.Events<MRun & MConfig, MResponse, Ctx>;
+  telemetry?: Telemetry.Provider;
   /** Whether or not to add default `console.log` event handlers */
   debug?: boolean;
 }
@@ -100,6 +102,7 @@ export abstract class AbstractModel<
     AResponse
   >;
   public readonly tokenizer: Model.ITokenizer;
+  public readonly telemetry: Telemetry.Provider;
 
   constructor(args: ModelArgs<MClient, MConfig, MRun, MResponse, CustomCtx>) {
     this.cacheKey = args.cacheKey ?? defaultCacheKey;
@@ -110,6 +113,7 @@ export abstract class AbstractModel<
     this.params = args.params;
     this.events = args.events || {};
     this.tokenizer = createTokenizer(args.params.model);
+    this.telemetry = args.telemetry ?? DefaultTelemetry;
   }
 
   async run(
@@ -129,56 +133,81 @@ export abstract class AbstractModel<
       ...mergedContext,
     });
 
-    return Sentry.startSpan({ name: spanName, op: 'llm' }, async (span) => {
-      const start = Date.now();
-
-      // Add tags for the prompt name and version if present
-      Sentry.setTags({
-        ...(typeof mergedContext?.promptName === 'string'
-          ? { promptName: mergedContext.promptName }
-          : {}),
-        ...(typeof mergedContext?.promptVersion === 'string'
-          ? { promptVersion: Number(mergedContext.promptVersion) }
-          : {}),
-      });
-
-      span.setAttributes({
-        ...extractAttrsFromContext(mergedContext),
-        ...extractAttrsFromParams({
-          modelType: this.modelType,
-          modelProvider: this.modelProvider,
-          ...mergedParams,
-        }),
-      });
-
-      await Promise.allSettled(
-        this.events.onStart?.map((event) =>
-          Promise.resolve(
-            event({
-              timestamp: new Date().toISOString(),
-              modelType: this.modelType,
-              modelProvider: this.modelProvider,
-              params: mergedParams,
-              context: mergedContext,
-            })
-          )
-        ) ?? []
-      );
-
-      const cacheKey = await this.cacheKey(mergedParams);
-
-      try {
-        // Check the cache
-        const cachedResponse =
-          this.cache && (await Promise.resolve(this.cache.get(cacheKey)));
-        if (cachedResponse) {
-          const response: MResponse = {
-            ...cachedResponse,
-            cached: true,
-            cost: 0,
-            latency: Date.now() - start,
-          };
+    return this.telemetry.startSpan(
+      { name: spanName, op: 'llm' },
+      async (span) => {
+        const start = Date.now();
+
+        // Add tags for the prompt name and version if present
+        this.telemetry.setTags({
+          ...(typeof mergedContext?.promptName === 'string'
+            ? { promptName: mergedContext.promptName }
+            : {}),
+          ...(typeof mergedContext?.promptVersion === 'string'
+            ? { promptVersion: Number(mergedContext.promptVersion) }
+            : {}),
+        });
+
+        span.setAttributes({
+          ...extractAttrsFromContext(mergedContext),
+          ...extractAttrsFromParams({
+            modelType: this.modelType,
+            modelProvider: this.modelProvider,
+            ...mergedParams,
+          }),
+        });
+
+        await Promise.allSettled(
+          this.events.onStart?.map((event) =>
+            Promise.resolve(
+              event({
+                timestamp: new Date().toISOString(),
+                modelType: this.modelType,
+                modelProvider: this.modelProvider,
+                params: mergedParams,
+                context: mergedContext,
+              })
+            )
+          ) ?? []
+        );
+
+        const cacheKey = await this.cacheKey(mergedParams);
+
+        try {
+          // Check the cache
+          const cachedResponse =
+            this.cache && (await Promise.resolve(this.cache.get(cacheKey)));
+          if (cachedResponse) {
+            const response: MResponse = {
+              ...cachedResponse,
+              cached: true,
+              cost: 0,
+              latency: Date.now() - start,
+            };
+            span.setAttributes(extractAttrsFromResponse(response));
+            await Promise.allSettled(
+              this.events.onComplete?.map((event) =>
+                Promise.resolve(
+                  event({
+                    timestamp: new Date().toISOString(),
+                    modelType: this.modelType,
+                    modelProvider: this.modelProvider,
+                    params: mergedParams,
+                    response,
+                    context: mergedContext,
+                    cached: true,
+                  })
+                )
+              ) ?? []
+            );
+            return response;
+          }
+
+          // Run the model (e.g. make the API request)
+          const response = await this.runModel(mergedParams, mergedContext);
+
           span.setAttributes(extractAttrsFromResponse(response));
+
           await Promise.allSettled(
             this.events.onComplete?.map((event) =>
               Promise.resolve(
@@ -189,56 +218,34 @@ export abstract class AbstractModel<
                   params: mergedParams,
                   response,
                   context: mergedContext,
-                  cached: true,
+                  cached: false,
                 })
               )
             ) ?? []
           );
-          return response;
-        }
-
-        // Run the model (e.g. make the API request)
-        const response = await this.runModel(mergedParams, mergedContext);
-
-        span.setAttributes(extractAttrsFromResponse(response));
-
-        await Promise.allSettled(
-          this.events.onComplete?.map((event) =>
-            Promise.resolve(
-              event({
-                timestamp: new Date().toISOString(),
-                modelType: this.modelType,
-                modelProvider: this.modelProvider,
-                params: mergedParams,
-                response,
-                context: mergedContext,
-                cached: false,
-              })
-            )
-          ) ?? []
-        );
 
-        // Update the cache
-        await Promise.resolve(this.cache?.set(cacheKey, response));
+          // Update the cache
+          await Promise.resolve(this.cache?.set(cacheKey, response));
 
-        return response;
-      } catch (error) {
-        await Promise.allSettled(
-          this.events?.onError?.map((event) =>
-            Promise.resolve(
-              event({
-                timestamp: new Date().toISOString(),
-                modelType: this.modelType,
-                modelProvider: this.modelProvider,
-                params: mergedParams,
-                error,
-                context: mergedContext,
-              })
-            )
-          ) ?? []
-        );
-        throw error;
+          return response;
+        } catch (error) {
+          await Promise.allSettled(
+            this.events?.onError?.map((event) =>
+              Promise.resolve(
+                event({
+                  timestamp: new Date().toISOString(),
+                  modelType: this.modelType,
+                  modelProvider: this.modelProvider,
+                  params: mergedParams,
+                  error,
+                  context: mergedContext,
+                })
+              )
+            ) ?? []
+          );
+          throw error;
+        }
       }
-    });
+    );
   }
 }

src/model/sparse-vector.ts

@@ -136,6 +136,7 @@ export class SparseVectorModel<
       cacheKey: this.cacheKey,
       cache: this.cache,
       debug: this.debug,
+      telemetry: this.telemetry,
       serviceUrl: this.serviceUrl,
       ...args,
       context: deepMerge(this.context, args?.context),

src/model/utils/telemetry.ts

@@ -1,4 +1,5 @@
 import type { ChatMessage, ChatParams } from 'openai-fetch';
+import type { Telemetry } from '../../telemetry/types.js';
 
 const SpanAttrs = {
   LLM_MODEL: 'llm.model',
@@ -21,8 +22,7 @@ const SpanAttrs = {
   LLM_TOKENS_TOTAL: 'llm.tkns.total',
 };
 
-type AttrValue = string | number | boolean | undefined;
-type AttrMap = Record<string, AttrValue>;
+type AttrMap = Telemetry.SpanAttributes;
 
 export function extractAttrsFromContext(context: any): AttrMap {
   try {

src/telemetry/default-telemetry.ts

@@ -0,0 +1,23 @@
+import type { Telemetry } from './types.js';
+
+export const DefaultTelemetry: Telemetry.Provider = {
+  startSpan<T>(
+    _options: Telemetry.SpanOptions,
+    callback: (span: Telemetry.Span) => T
+  ): T {
+    return callback({
+      setAttribute() {
+        // no-op
+        return this;
+      },
+      setAttributes() {
+        // no-op
+        return this;
+      },
+    });
+  },
+
+  setTags() {
+    // no-op
+  },
+};

src/telemetry/index.ts

@@ -0,0 +1,2 @@
+export type { Telemetry } from './types.js';
+export { DefaultTelemetry } from './default-telemetry.js';

src/telemetry/sentry.test.ts

@@ -0,0 +1,6 @@
+import * as Sentry from '@sentry/node';
+import type { Telemetry } from './types.js';
+
+test('Sentry Telemetry Provider', () => {
+  expectTypeOf(Sentry).toMatchTypeOf<Telemetry.Provider>();
+});