Skip to content

Commit 36cc9f9

Browse files
committed
refactor: integrate telemetry configuration for AI processing across categorization, deduplication, and document parsing
1 parent a9597e6 commit 36cc9f9

7 files changed

Lines changed: 188 additions & 2 deletions

File tree

.env.example

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,11 @@ WXT_SUPABASE_PUBLISHABLE_KEY=your-supabase-publishable-key
1919
# Set this to 'true' to enable debug logs in production builds
2020
# WXT_DEBUG=true
2121

22-
# Langfuse Observability (optional)
23-
# Get these from https://us.cloud.langfuse.com
22+
# Langfuse Observability (dev-only, optional)
23+
# AI telemetry is automatically enabled in dev mode (console logging).
24+
# To also send traces to Langfuse, set these keys.
25+
# Get them from https://us.cloud.langfuse.com (free tier available).
26+
# These are NEVER included in production builds.
2427
# WXT_LANGFUSE_PUBLIC_KEY=pk-lf-...
2528
# WXT_LANGFUSE_SECRET_KEY=sk-lf-...
2629
# WXT_LANGFUSE_BASEURL=https://us.cloud.langfuse.com

src/lib/ai/bulk-categorizer.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { generateText, Output } from "ai";
22
import { z } from "zod";
3+
import { getTelemetryConfig } from "@/lib/ai/telemetry";
34
import { allowedCategories } from "@/lib/copies";
45
import { createLogger } from "@/lib/logger";
56
import { getAIModel, getProviderOptions } from "@/lib/providers/model-factory";
@@ -92,6 +93,7 @@ export class BulkCategorizer {
9293
system: systemPrompt,
9394
prompt: userPrompt,
9495
providerOptions: getProviderOptions(provider),
96+
...getTelemetryConfig("bulk-categorization"),
9597
});
9698

9799
const categorized = this.mapResultsToFields(result, fields);

src/lib/ai/categorization.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { generateText, Output } from "ai";
22
import { z } from "zod";
3+
import { getTelemetryConfig } from "@/lib/ai/telemetry";
34
import { createLogger } from "@/lib/logger";
45
import { getAIModel, getProviderOptions } from "@/lib/providers/model-factory";
56
import type { AIProvider } from "@/lib/providers/registry";
@@ -152,6 +153,7 @@ Be precise and consider context. For example:
152153
prompt: userPrompt,
153154
temperature: 0.3,
154155
providerOptions: getProviderOptions(provider),
156+
...getTelemetryConfig("categorization"),
155157
});
156158

157159
return result.output;
@@ -243,6 +245,7 @@ Rephrase the following answer based on the provided context.
243245
prompt: userPrompt,
244246
temperature: 0.4,
245247
providerOptions: getProviderOptions(provider),
248+
...getTelemetryConfig("rephrase-context"),
246249
});
247250

248251
return output.rephrasedAnswer;
@@ -284,6 +287,7 @@ export const rephraseAgent = async (
284287
prompt: userPrompt,
285288
temperature: 0.5,
286289
providerOptions: getProviderOptions(provider),
290+
...getTelemetryConfig("rephrase"),
287291
});
288292

289293
return output;

src/lib/ai/deduplication-categorizer.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { generateText, Output } from "ai";
22
import { z } from "zod";
33
import { CategoryEnum, TagSchema } from "@/lib/ai/categorization";
4+
import { getTelemetryConfig } from "@/lib/ai/telemetry";
45
import { createLogger } from "@/lib/logger";
56
import type { AIProvider } from "@/lib/providers/registry";
67
import { storage } from "@/lib/storage";
@@ -141,6 +142,7 @@ export class DeduplicationCategorizer {
141142
prompt: userPrompt,
142143
temperature: 0.3,
143144
providerOptions: getProviderOptions(provider),
145+
...getTelemetryConfig("deduplication"),
144146
});
145147

146148
logger.info("Deduplication + Categorization result:", {

src/lib/ai/matcher.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { generateText, Output } from "ai";
22
import { z } from "zod";
3+
import { getTelemetryConfig } from "@/lib/ai/telemetry";
34
import { getAuthService } from "@/lib/auth/auth-service";
45
import { FallbackMatcher } from "@/lib/autofill/fallback-matcher";
56
import {
@@ -222,6 +223,7 @@ export class AIMatcher {
222223
prompt: userPrompt,
223224
temperature: 0.3,
224225
providerOptions: getProviderOptions(provider),
226+
...getTelemetryConfig("field-matching"),
225227
});
226228

227229
return result.output;

src/lib/ai/telemetry.ts

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
import type { TelemetryIntegration } from "ai";
2+
import { bindTelemetryIntegration } from "ai";
3+
import { createLogger } from "@/lib/logger";
4+
5+
const logger = createLogger("ai:telemetry");
6+
7+
const LANGFUSE_PUBLIC_KEY = import.meta.env.WXT_LANGFUSE_PUBLIC_KEY as
8+
| string
9+
| undefined;
10+
const LANGFUSE_SECRET_KEY = import.meta.env.WXT_LANGFUSE_SECRET_KEY as
11+
| string
12+
| undefined;
13+
const LANGFUSE_BASEURL =
14+
(import.meta.env.WXT_LANGFUSE_BASEURL as string | undefined) ||
15+
"https://us.cloud.langfuse.com";
16+
17+
const langfuseEnabled = !!(LANGFUSE_PUBLIC_KEY && LANGFUSE_SECRET_KEY);
18+
19+
class DevTelemetryIntegration implements TelemetryIntegration {
20+
private startTime = 0;
21+
private startTimestamp = "";
22+
private inputData: unknown = undefined;
23+
24+
async onStart(event: {
25+
model: { provider: string; modelId: string };
26+
system: unknown;
27+
prompt: unknown;
28+
messages: unknown;
29+
}) {
30+
this.startTime = performance.now();
31+
this.startTimestamp = new Date().toISOString();
32+
this.inputData = event.messages ?? event.prompt ?? event.system;
33+
logger.info(
34+
`AI call started | provider: ${event.model.provider} | model: ${event.model.modelId}`,
35+
);
36+
}
37+
38+
async onFinish(event: {
39+
model: { provider: string; modelId: string };
40+
totalUsage: {
41+
inputTokens: number | undefined;
42+
outputTokens: number | undefined;
43+
};
44+
finishReason: string;
45+
functionId: string | undefined;
46+
text: string;
47+
}) {
48+
const duration = performance.now() - this.startTime;
49+
const inputTokens = event.totalUsage.inputTokens ?? 0;
50+
const outputTokens = event.totalUsage.outputTokens ?? 0;
51+
52+
logger.info(
53+
`AI call completed | fn: ${event.functionId ?? "unknown"} | model: ${event.model.modelId} | tokens: ${inputTokens}+${outputTokens}=${inputTokens + outputTokens} | duration: ${(duration / 1000).toFixed(2)}s | finish: ${event.finishReason}`,
54+
);
55+
56+
if (langfuseEnabled) {
57+
this.sendToLangfuse(event, duration).catch((err) =>
58+
logger.debug("Langfuse send failed:", err),
59+
);
60+
}
61+
}
62+
63+
private async sendToLangfuse(
64+
event: {
65+
model: { provider: string; modelId: string };
66+
totalUsage: {
67+
inputTokens: number | undefined;
68+
outputTokens: number | undefined;
69+
};
70+
finishReason: string;
71+
functionId: string | undefined;
72+
text: string;
73+
},
74+
durationMs: number,
75+
) {
76+
const traceId = crypto.randomUUID();
77+
const generationId = crypto.randomUUID();
78+
const endTime = new Date().toISOString();
79+
const startTime = this.startTimestamp || endTime;
80+
const input = this.inputData;
81+
const output = event.text;
82+
83+
const batch = [
84+
{
85+
id: crypto.randomUUID(),
86+
type: "trace-create" as const,
87+
timestamp: startTime,
88+
body: {
89+
id: traceId,
90+
name: event.functionId ?? "ai-call",
91+
input,
92+
output,
93+
metadata: {
94+
provider: event.model.provider,
95+
model: event.model.modelId,
96+
source: "superfill-extension-dev",
97+
},
98+
},
99+
},
100+
{
101+
id: crypto.randomUUID(),
102+
type: "generation-create" as const,
103+
timestamp: startTime,
104+
body: {
105+
id: generationId,
106+
traceId,
107+
name: event.functionId ?? "ai-call",
108+
model: event.model.modelId,
109+
input,
110+
output,
111+
startTime,
112+
endTime,
113+
modelParameters: {
114+
provider: event.model.provider,
115+
},
116+
usage: {
117+
input: event.totalUsage.inputTokens ?? 0,
118+
output: event.totalUsage.outputTokens ?? 0,
119+
unit: "TOKENS",
120+
},
121+
metadata: {
122+
durationMs: Math.round(durationMs),
123+
finishReason: event.finishReason,
124+
},
125+
},
126+
},
127+
];
128+
129+
const credentials = btoa(`${LANGFUSE_PUBLIC_KEY}:${LANGFUSE_SECRET_KEY}`);
130+
131+
await fetch(`${LANGFUSE_BASEURL}/api/public/ingestion`, {
132+
method: "POST",
133+
headers: {
134+
"Content-Type": "application/json",
135+
Authorization: `Basic ${credentials}`,
136+
},
137+
body: JSON.stringify({ batch }),
138+
});
139+
140+
logger.debug(`Langfuse trace sent: ${traceId}`);
141+
}
142+
}
143+
144+
function devTelemetryIntegration(): TelemetryIntegration {
145+
return bindTelemetryIntegration(new DevTelemetryIntegration());
146+
}
147+
148+
/**
149+
* Returns telemetry config to spread into generateText/streamText calls.
150+
* In production, returns an empty object (no-op when spread).
151+
* In dev, enables telemetry with console logging + optional Langfuse.
152+
*/
153+
export function getTelemetryConfig(functionId: string): {
154+
experimental_telemetry?: {
155+
isEnabled: boolean;
156+
functionId: string;
157+
integrations: TelemetryIntegration[];
158+
};
159+
} {
160+
if (!import.meta.env.DEV) {
161+
return {};
162+
}
163+
164+
return {
165+
experimental_telemetry: {
166+
isEnabled: true,
167+
functionId,
168+
integrations: [devTelemetryIntegration()],
169+
},
170+
};
171+
}

src/lib/document/document-parser.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { generateText, Output } from "ai";
22
import { z } from "zod";
3+
import { getTelemetryConfig } from "@/lib/ai/telemetry";
34
import { createLogger } from "@/lib/logger";
45
import { getAIModel, getProviderOptions } from "@/lib/providers/model-factory";
56
import { getKeyVaultService } from "@/lib/security/key-vault-service";
@@ -209,6 +210,7 @@ async function parseDocumentWithAI(text: string): Promise<ExtractedItem[]> {
209210
prompt: `Extract all useful personal and professional information from this document:\n\n${text}`,
210211
temperature: 0.1,
211212
providerOptions: getProviderOptions(selectedProvider),
213+
...getTelemetryConfig("document-parsing"),
212214
});
213215

214216
logger.debug("AI extracted items:", output.items.length);

0 commit comments

Comments
 (0)