Home/llm-observability-mcp
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor config loading (in progress 3)

Browse Source
This commit is contained in:
Stefano Fiorini
2025-07-15 00:15:56 -05:00
parent 0a045762bc
commit 4b2d63335a
19 changed files with 380 additions and 360 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/cli/posthog-llm.cli.ts
View File

@@ -1,7 +1,7 @@
import { Command } from 'commander';
import { Logger } from '../utils/logger.util.js';
import { handleCliError } from '../utils/error.util.js';
import postHogLlmController from '../controllers/posthog-llm.controller.js';
//import { PostHogController } from '../controllers/posthog-llm.controller.js';
const logger = Logger.forContext('cli/posthog-llm.cli.ts');
@@ -48,7 +48,9 @@ function register(program: Command) {
}
const args = { eventName, distinctId, properties };
await postHogLlmController.capture(args);
actionLogger.debug(`CLI posthog-llm capture args`, args);
// TODO: Implement the controller capture with correct properties
//await PostHogController.capture(args);
} catch (error) {
handleCliError(error);
}

42
src/config/base.config.ts
View File

@@ -1,42 +0,0 @@
import { z } from 'zod';
import { config } from './global.config.js';
/**
* Base configuration class with common functionality
*/
export abstract class BaseConfig<T> {
protected abstract schema: z.ZodSchema<T>;
protected abstract configKey: string;
protected getValue(key: string, defaultValue?: string): string | undefined {
return config.get(key, defaultValue);
}
protected getBoolean(key: string, defaultValue: boolean = false): boolean {
return config.getBoolean(key, defaultValue);
}
protected getNumber(key: string, defaultValue: number): number {
const value = this.getValue(key);
if (value === undefined) return defaultValue;
return Number(value);
}
protected parseHeaders(headersString: string): Record<string, string> {
const headers: Record<string, string> = {};
const pairs = headersString.split(',');
for (const pair of pairs) {
const [key, value] = pair.split('=');
if (key && value) {
headers[key.trim()] = value.trim();
}
}
return headers;
}
public load(): T {
return this.schema.parse(this.loadFromSources());
}
protected abstract loadFromSources(): Record<string, unknown>;
}

14
src/config/common.schema.ts Normal file
View File

@@ -0,0 +1,14 @@
import { z } from 'zod';
/**
* Common/shared configuration options schema (now includes global options)
*/
export const CommonConfigSchema = z.object({
serviceName: z.string().default('llm-observability-mcp'),
serviceVersion: z.string().default('1.0.0'),
environment: z.string().default('development'),
debug: z.boolean().default(false),
logLevel: z.enum(['debug', 'info', 'warn', 'error']).default('info'),
});
export type CommonConfigType = z.infer<typeof CommonConfigSchema>;

202
src/config/config-loader.ts Normal file
View File

@@ -0,0 +1,202 @@
import fs from 'fs';
import path from 'path';
import { Logger } from '../utils/logger.util.js';
import dotenv from 'dotenv';
import os from 'os';
import { CommonConfigSchema, CommonConfigType } from './common.schema.js';
import { PostHogConfigSchema, PostHogConfigType } from './posthog.schema.js';
import {
OpenTelemetryConfigSchema,
OpenTelemetryConfigType,
} from './opentelemetry-llm.schema.js';
import { z } from 'zod';
/**
* Configuration loader that handles multiple sources with priority:
* 1. Direct ENV pass (process.env)
* 2. .env file in project root
* 3. Global config file at $HOME/.mcp/configs.json
*/
class ConfigLoader {
private packageName: string;
private configLoaded: boolean = false;
constructor(packageName: string) {
this.packageName = packageName;
}
load(): void {
const logger = Logger.forContext('config-loader', 'load');
if (this.configLoaded) {
logger.debug('Configuration already loaded, skipping');
return;
}
logger.debug('Loading configuration...');
// Load from global config file
this.loadFromGlobalConfig();
// Load from .env file
this.loadFromEnvFile();
this.configLoaded = true;
logger.debug('Configuration loaded successfully');
}
private loadFromEnvFile(): void {
const logger = Logger.forContext('config-loader', 'loadFromEnvFile');
try {
const result = dotenv.config();
if (result.error) {
logger.debug('No .env file found or error reading it');
return;
}
logger.debug('Loaded configuration from .env file');
} catch (error) {
logger.error('Error loading .env file', error);
}
}
private loadFromGlobalConfig(): void {
const logger = Logger.forContext(
'config-loader',
'loadFromGlobalConfig',
);
try {
const homedir = os.homedir();
const globalConfigPath = path.join(homedir, '.mcp', 'configs.json');
if (!fs.existsSync(globalConfigPath)) {
logger.debug('Global config file not found');
return;
}
const configContent = fs.readFileSync(globalConfigPath, 'utf8');
const config = JSON.parse(configContent);
const shortKey = 'llm-observability-mcp';
const fullPackageName = this.packageName;
const unscopedPackageName =
fullPackageName.split('/')[1] || fullPackageName;
const potentialKeys = [
shortKey,
fullPackageName,
unscopedPackageName,
];
let foundConfig = null;
for (const key of potentialKeys) {
if (
config[key] &&
typeof config[key] === 'object' &&
config[key].environments
) {
foundConfig = config[key];
logger.debug(`Found configuration using key: ${key}`);
break;
}
}
if (!foundConfig || !foundConfig.environments) {
logger.debug(
`No config found for ${this.packageName} using keys: ${potentialKeys.join(', ')}`,
);
return;
}
for (const [key, value] of Object.entries(
foundConfig.environments,
)) {
if (process.env[key] === undefined) {
process.env[key] = String(value);
}
}
logger.debug(`Loaded configuration from global config file`);
} catch (error) {
logger.error('Error loading global config file', error);
}
}
private getEnvObject(schema: z.ZodTypeAny): Record<string, unknown> {
const shape =
typeof schema._def.shape === 'function'
? schema._def.shape()
: schema._def.shape;
const result: Record<string, unknown> = {};
for (const key in shape) {
if (Object.prototype.hasOwnProperty.call(shape, key)) {
const envKey = this.toEnvKey(key);
const value = process.env[envKey];
if (value !== undefined) {
const type = shape[key]._def.typeName;
if (type === 'ZodNumber') {
result[key] = Number(value);
} else if (type === 'ZodBoolean') {
result[key] = value.toLowerCase() === 'true';
} else {
result[key] = value;
}
}
}
}
return result;
}
private toEnvKey(key: string): string {
// Map schema keys to environment variable names
// e.g. serviceName -> OTEL_SERVICE_NAME, apiKey -> POSTHOG_API_KEY, etc.
const map: Record<string, string> = {
serviceName: 'OTEL_SERVICE_NAME',
serviceVersion: 'OTEL_SERVICE_VERSION',
environment: 'OTEL_ENVIRONMENT',
metricsEndpoint: 'OTEL_EXPORTER_OTLP_METRICS_ENDPOINT',
tracesEndpoint: 'OTEL_EXPORTER_OTLP_TRACES_ENDPOINT',
logsEndpoint: 'OTEL_EXPORTER_OTLP_LOGS_ENDPOINT',
headers: 'OTEL_EXPORTER_OTLP_HEADERS',
exportIntervalMillis: 'OTEL_METRIC_EXPORT_INTERVAL',
exportTimeoutMillis: 'OTEL_METRIC_EXPORT_TIMEOUT',
samplingRatio: 'OTEL_TRACES_SAMPLER_ARG',
apiKey: 'POSTHOG_API_KEY',
host: 'POSTHOG_HOST',
debug: 'DEBUG',
logLevel: 'LOG_LEVEL',
};
return map[key] || key;
}
getCommonConfig(): CommonConfigType {
return CommonConfigSchema.parse(this.getEnvObject(CommonConfigSchema));
}
getPosthogConfig(): PostHogConfigType {
return PostHogConfigSchema.parse(
this.getEnvObject(PostHogConfigSchema),
);
}
getOpenTelemetryConfig(): OpenTelemetryConfigType {
const raw = this.getEnvObject(OpenTelemetryConfigSchema);
// Special handling for headers (parse comma-separated string to object)
if (typeof raw.headers === 'string') {
const headers: Record<string, string> = {};
const pairs = raw.headers.split(',');
for (const pair of pairs) {
const [key, value] = pair.split('=');
if (key && value) {
headers[key.trim()] = value.trim();
}
}
raw.headers = headers;
}
return OpenTelemetryConfigSchema.parse(raw);
}
}
const configLoader = new ConfigLoader('@sfiorini/llm-observability-mcp');
export default configLoader;

130
src/config/global.config.ts
View File

@@ -1,130 +0,0 @@
import fs from 'fs';
import path from 'path';
import { Logger } from '../utils/logger.util.js';
import dotenv from 'dotenv';
import os from 'os';
/**
* Configuration loader that handles multiple sources with priority:
* 1. Direct ENV pass (process.env)
* 2. .env file in project root
* 3. Global config file at $HOME/.mcp/configs.json
*/
class ConfigLoader {
private packageName: string;
private configLoaded: boolean = false;
constructor(packageName: string) {
this.packageName = packageName;
}
load(): void {
const logger = Logger.forContext('global.config', 'load');
if (this.configLoaded) {
logger.debug('Configuration already loaded, skipping');
return;
}
logger.debug('Loading configuration...');
// Load from global config file
this.loadFromGlobalConfig();
// Load from .env file
this.loadFromEnvFile();
this.configLoaded = true;
logger.debug('Configuration loaded successfully');
}
private loadFromEnvFile(): void {
const logger = Logger.forContext('global.config', 'loadFromEnvFile');
try {
const result = dotenv.config();
if (result.error) {
logger.debug('No .env file found or error reading it');
return;
}
logger.debug('Loaded configuration from .env file');
} catch (error) {
logger.error('Error loading .env极 file', error);
}
}
private loadFromGlobalConfig(): void {
const logger = Logger.forContext(
'global.config',
'loadFromGlobalConfig',
);
try {
const homedir = os.homedir();
const globalConfigPath = path.join(homedir, '.mcp', 'configs.json');
if (!fs.existsSync(globalConfigPath)) {
logger.debug('Global config file not found');
return;
}
const configContent = fs.readFileSync(globalConfigPath, 'utf8');
const config = JSON.parse(configContent);
const shortKey = 'llm-observability-mcp';
const fullPackageName = this.packageName;
const unscopedPackageName =
fullPackageName.split('/')[1] || fullPackageName;
const potentialKeys = [
shortKey,
fullPackageName,
unscopedPackageName,
];
let foundConfig = null;
for (const key of potentialKeys) {
if (
config[key] &&
typeof config[key] === 'object' &&
config[key].environments
) {
foundConfig = config[key];
logger.debug(`Found configuration using key: ${key}`);
break;
}
}
if (!foundConfig || !foundConfig.environments) {
logger.debug(
`No config found for ${this.packageName} using keys: ${potentialKeys.join(', ')}`,
);
return;
}
for (const [key, value] of Object.entries(
foundConfig.environments,
)) {
if (process.env[key] === undefined) {
process.env[key] = String(value);
}
}
logger.debug(`Loaded configuration from global config file`);
} catch (error) {
logger.error('Error loading global config file', error);
}
}
get(key: string, defaultValue?: string): string | undefined {
return process.env[key] || defaultValue;
}
getBoolean(key: string, defaultValue: boolean = false): boolean {
const value = this.get(key);
if (value === undefined) return defaultValue;
return value.toLowerCase() === 'true';
}
}
export const config = new ConfigLoader('@sfiorini/llm-observability-mcp');

82
src/config/opentelemetry-llm.config.ts
View File

@@ -1,82 +0,0 @@
import { BaseConfig } from './base.config.js';
import {
OpenTelemetryConfigSchema,
type OpenTelemetryConfigType,
} from './opentelemetry-llm.schema.js';
import { z } from 'zod';
/**
* OpenTelemetry configuration class
*/
export class OpenTelemetryConfig extends BaseConfig<OpenTelemetryConfigType> {
protected schema =
OpenTelemetryConfigSchema as z.ZodSchema<OpenTelemetryConfigType>;
protected configKey = 'opentelemetry';
constructor() {
super();
}
/**
* Create configuration from environment variables
*/
public fromEnv(): OpenTelemetryConfigType {
return super.load();
}
/**
* Check if OpenTelemetry is enabled
*/
public isEnabled(config: OpenTelemetryConfigType): boolean {
return !!(
config.metricsEndpoint ||
config.tracesEndpoint ||
config.logsEndpoint
);
}
/**
* Get combined OTLP endpoint if only one is provided
*/
public getDefaultEndpoint(
config: OpenTelemetryConfigType,
): string | undefined {
if (
config.metricsEndpoint &&
config.tracesEndpoint === config.metricsEndpoint
) {
return config.metricsEndpoint;
}
return config.metricsEndpoint || config.tracesEndpoint;
}
protected loadFromSources(): Record<string, unknown> {
return {
serviceName: this.getValue(
'OTEL_SERVICE_NAME',
'llm-observability-mcp',
),
serviceVersion: this.getValue('OTEL_SERVICE_VERSION', '1.0.0'),
environment: this.getValue('OTEL_ENVIRONMENT', 'development'),
metricsEndpoint: this.getValue(
'OTEL_EXPORTER_OTLP_METRICS_ENDPOINT',
),
tracesEndpoint: this.getValue('OTEL_EXPORTER_OTLP_TRACES_ENDPOINT'),
logsEndpoint: this.getValue('OTEL_EXPORTER_OTLP_LOGS_ENDPOINT'),
headers: this.getValue('OTEL_EXPORTER_OTLP_HEADERS')
? this.parseHeaders(
this.getValue('OTEL_EXPORTER_OTLP_HEADERS')!,
)
: undefined,
exportIntervalMillis: this.getNumber(
'OTEL_METRIC_EXPORT_INTERVAL',
10000,
),
exportTimeoutMillis: this.getNumber(
'OTEL_METRIC_EXPORT_TIMEOUT',
5000,
),
samplingRatio: this.getNumber('OTEL_TRACES_SAMPLER_ARG', 1.0),
};
}
}

8
src/config/opentelemetry-llm.schema.ts
View File

@@ -1,14 +1,10 @@
import { z } from 'zod';
import { CommonConfigSchema } from './common.schema.js';
/**
* Configuration schema for OpenTelemetry
*/
export const OpenTelemetryConfigSchema = z.object({
// Service configuration
serviceName: z.string().default('llm-observability-mcp'),
serviceVersion: z.string().default('1.0.0'),
environment: z.string().default('development'),
export const OpenTelemetryConfigSchema = CommonConfigSchema.extend({
// OTLP endpoints
metricsEndpoint: z.string().optional(),
tracesEndpoint: z.string().optional(),

15
src/config/posthog.schema.ts Normal file
View File

@@ -0,0 +1,15 @@
import { z } from 'zod';
import { CommonConfigSchema } from './common.schema.js';
/**
* PostHog configuration options schema
*/
export const PostHogConfigSchema = CommonConfigSchema.extend({
apiKey: z.string().describe('PostHog API key'),
host: z
.string()
.default('https://app.posthog.com')
.describe('PostHog host URL'),
});
export type PostHogConfigType = z.infer<typeof PostHogConfigSchema>;

8
src/controllers/base.controller.ts Normal file
View File

@@ -0,0 +1,8 @@
import { ControllerResponse } from '../types/common.types.js';
export abstract class BaseController {
// Instance method (must be implemented by subclasses)
static capture(data: unknown): Promise<ControllerResponse> {
throw new Error(`Not implemented: capture: ${data}`);
}
}

25
src/controllers/opentelemetry-llm.controller.ts
View File

@@ -1,4 +1,5 @@
import { Logger } from '../utils/logger.util.js';
import { BaseController } from './base.controller.js';
import { OpenTelemetryService } from '../services/opentelemetry-llm.service.js';
import { OpenTelemetryLlmInputSchemaType } from '../types/opentelemetry-llm.types.js';
import { ControllerResponse } from '../types/common.types.js';
@@ -6,17 +7,11 @@ import { ControllerResponse } from '../types/common.types.js';
/**
* Controller for OpenTelemetry LLM observability
*/
export class OpenTelemetryController {
private service: OpenTelemetryService;
constructor(service: OpenTelemetryService) {
this.service = service;
}
export class OpenTelemetryController extends BaseController {
/**
* Capture LLM observability data
*/
public async captureLlmObservability(
static async capture(
data: OpenTelemetryLlmInputSchemaType,
): Promise<ControllerResponse> {
const logger = Logger.forContext(
@@ -26,8 +21,12 @@ export class OpenTelemetryController {
logger.debug('Capturing LLM observability data', data);
try {
// Get and initialize the service
const service = OpenTelemetryService.getInstance();
service.initialize();
// Record the LLM request
this.service.recordLlmRequest(data);
service.recordLlmRequest(data);
const content = `## OpenTelemetry LLM Observability Captured
@@ -64,11 +63,3 @@ The data has been sent to your configured OpenTelemetry collector and is now ava
}
}
}
// Export singleton instance
const openTelemetryService = OpenTelemetryService.getInstance();
const openTelemetryController = new OpenTelemetryController(
openTelemetryService,
);
export default openTelemetryController;

31
src/controllers/posthog-llm.controller.ts
View File

@@ -1,32 +1,35 @@
import { Logger } from '../utils/logger.util.js';
import { ControllerResponse } from '../types/common.types.js';
import postHogLlmService from '../services/posthog-llm.service.js';
import { PostHogLlmInputSchemaType } from '../types/posthog-llm.types.js';
const logger = Logger.forContext('controllers/posthog-llm.controller.ts');
/**
* Controller for PostHog LLM observability
*/
export class PostHogController {
/**
* Capture LLM observability data
*/
async function capture(args: {
eventName: string;
distinctId: string;
properties: Record<string, unknown>;
}): Promise<ControllerResponse> {
const methodLogger = logger.forMethod('capture');
methodLogger.debug('Capturing PostHog event...');
methodLogger.debug('Arguments:', args);
static async capture(
data: PostHogLlmInputSchemaType,
): Promise<ControllerResponse> {
const logger = Logger.forContext('posthog-llm.controller', 'capture');
logger.debug('Capturing LLM observability data', data);
try {
await postHogLlmService.capture(args);
await postHogLlmService.capture(data);
return {
content: 'Event captured successfully.',
};
} catch (error) {
methodLogger.error('Error capturing event:', error);
logger.error('Error capturing event:', error);
const errorMessage =
error instanceof Error ? error.message : 'Unknown error';
methodLogger.error('Error capturing event:', errorMessage);
logger.error('Error capturing event:', errorMessage);
return {
content: `Failed to capture event: ${errorMessage}`,
};
}
}
}
export default { capture };

4
src/index.ts
View File

@@ -1,6 +1,6 @@
#!/usr/bin/env node
import { Logger } from './utils/logger.util.js';
import { config } from './config/global.config.js';
import configLoader from './config/config-loader.js';
import { runCli } from './cli/index.js';
import { stdioTransport } from './server/stdio.js';
import { streamableHttpTransport } from './server/streamableHttp.js';
@@ -46,7 +46,7 @@ async function main() {
const mainLogger = Logger.forContext('index.ts', 'main');
// Load configuration
config.load();
configLoader.load();
// CLI mode - if any arguments are provided
if (process.argv.length > 2) {

5
src/resources/posthog-llm.resource.ts
View File

@@ -1,6 +1,6 @@
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { Logger } from '../utils/logger.util.js';
import posthogLlmController from '../controllers/posthog-llm.controller.js';
import { PostHogController } from '../controllers/posthog-llm.controller.js';
import { formatErrorForMcpResource } from '../utils/error.util.js';
const logger = Logger.forContext('resources/posthog-llm.resource.ts');
@@ -44,7 +44,8 @@ function registerResources(server: McpServer) {
}
// Call the controller to capture the event
const result = await posthogLlmController.capture({
// TODO: Implement the controller capture with correct properties
const result = await PostHogController.capture({
eventName,
distinctId,
properties,

7
src/server/mcpServer.ts
View File

@@ -3,7 +3,7 @@ import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import { Logger } from '../utils/logger.util';
import { config } from '../config/global.config';
import configLoader from '../config/config-loader';
import { PACKAGE_NAME, VERSION } from '../utils/constants.util';
import posthogLlmResources from '../resources/posthog-llm.resource.js';
import posthogLlmTools from '../tools/posthog-llm.tool.js';
@@ -14,9 +14,10 @@ export function createServer() {
// Load configuration
serverLogger.info('Starting MCP server initialization...');
config.load();
configLoader.load();
if (config.getBoolean('DEBUG')) {
const commonConfig = configLoader.getCommonConfig();
if (commonConfig.debug) {
serverLogger.debug('Debug mode enabled');
}

16
src/services/opentelemetry-llm.service.ts
View File

@@ -23,6 +23,7 @@ import {
metrics,
} from '@opentelemetry/api';
import { trace, Tracer, SpanStatusCode } from '@opentelemetry/api';
import configLoader from '../config/config-loader';
/**
* OpenTelemetry service for LLM observability
@@ -41,8 +42,8 @@ export class OpenTelemetryService {
private latencyHistogram: Histogram;
private activeRequests: UpDownCounter;
private constructor(config: OpenTelemetryConfigType) {
this.config = config;
private constructor() {
this.config = configLoader.getOpenTelemetryConfig();
this.tracer = trace.getTracer('llm-observability-mcp');
this.meter = this.initializeMetrics();
@@ -74,16 +75,9 @@ export class OpenTelemetryService {
/**
* Get singleton instance
*/
public static getInstance(
config?: OpenTelemetryConfigType,
): OpenTelemetryService {
public static getInstance(): OpenTelemetryService {
if (!OpenTelemetryService.instance) {
if (!config) {
throw new Error(
'OpenTelemetryService requires configuration on first initialization',
);
}
OpenTelemetryService.instance = new OpenTelemetryService(config);
OpenTelemetryService.instance = new OpenTelemetryService();
}
return OpenTelemetryService.instance;
}

12
src/services/posthog-llm.service.ts
View File

@@ -1,18 +1,18 @@
import { Logger } from '../utils/logger.util';
import { PostHog } from 'posthog-node';
import { config } from '../config/global.config';
import configLoader from '../config/config-loader';
// Ensure configuration is loaded before accessing environment variables
config.load();
configLoader.load();
const logger = Logger.forContext('services/posthog-llm.service.ts');
const posthogApiKey = config.get('POSTHOG_API_KEY');
const posthogConfig = configLoader.getPosthogConfig();
let posthogClient: PostHog | null = null;
if (posthogApiKey) {
posthogClient = new PostHog(posthogApiKey, {
host: config.get('POSTHOG_HOST') || 'https://app.posthog.com',
if (posthogConfig.apiKey) {
posthogClient = new PostHog(posthogConfig.apiKey, {
host: posthogConfig.host,
});
} else {
logger.warn('POSTHOG_API_KEY is not set. PostHog client not initialized.');

15
src/tools/opentelemetry-llm.tool.ts
View File

@@ -1,14 +1,12 @@
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { Logger } from '../utils/logger.util.js';
import { formatErrorForMcpTool } from '../utils/error.util.js';
import { OpenTelemetryService } from '../services/opentelemetry-llm.service.js';
import { OpenTelemetryController } from '../controllers/opentelemetry-llm.controller.js';
import { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
import {
OpenTelemetryLlmInputSchema,
OpenTelemetryLlmInputSchemaType,
} from '../types/opentelemetry-llm.types.js';
import { OpenTelemetryConfig } from '../config/opentelemetry-llm.config.js';
/**
* @function captureOpenTelemetryLlmObservability
@@ -31,19 +29,8 @@ async function captureOpenTelemetryLlmObservability(
// Parse and validate arguments
const validatedArgs = OpenTelemetryLlmInputSchema.parse(args);
// Ensure OpenTelemetry is configured
const configLoader = new OpenTelemetryConfig();
const config = configLoader.fromEnv();
const service = OpenTelemetryService.getInstance(config);
// Initialize if not already done
service.initialize();
// Create controller with the service
const controller = new OpenTelemetryController(service);
// Pass validated args to the controller
const result = await controller.captureLlmObservability(validatedArgs);
const result = await OpenTelemetryController.capture(validatedArgs);
methodLogger.debug('Got response from controller', result);
// Format the response for the MCP tool

36
src/tools/posthog-llm.tool.ts
View File

@@ -1,23 +1,25 @@
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { Logger } from '../utils/logger.util.js';
import { formatErrorForMcpTool } from '../utils/error.util.js';
import posthogLlmController from '../controllers/posthog-llm.controller.js';
import { PostHogController } from '../controllers/posthog-llm.controller.js';
import { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
import {
GetToolInputSchema,
GetToolInputSchemaType,
PostHogLlmInputSchemaType,
PostHogLlmPropertiesApiSchemaType,
PostHogLlmPropertiesPayloadSchema,
PostHogLlmPropertiesPayloadSchemaType,
} from '../types/posthog-llm.types.js';
/**
* @function capturePosthogLlmObservability
* @description MCP Tool handler to capture LLM observability events in PostHog.
* It calls the posthogLlmController to track the data and formats the response for the MCP.
* @param {GetToolInputSchemaType} args - Arguments provided to the tool.
* @param {PostHogLlmInputSchemaType} args - Arguments provided to the tool.
* @returns {Promise<CallToolResult>} Formatted response for the MCP.
* @throws {McpError} Formatted error if the controller or service layer encounters an issue.
*/
async function capturePosthogLlmObservability(
args: GetToolInputSchemaType,
args: PostHogLlmPropertiesPayloadSchemaType,
): Promise<CallToolResult> {
const methodLogger = Logger.forContext(
'tools/posthog-llm.tool.ts',
@@ -26,15 +28,16 @@ async function capturePosthogLlmObservability(
methodLogger.debug(`Capture LLM Observability in PostHog...`, args);
try {
const trackArgs = GetToolInputSchema.parse(args);
const trackArgs = PostHogLlmPropertiesPayloadSchema.parse(args);
const posthogProperties: Record<string, unknown> = {
const posthogProperties: PostHogLlmPropertiesApiSchemaType = {
userId: trackArgs.userId,
$ai_model: trackArgs.model,
$ai_provider: trackArgs.provider,
};
const toPostHogKey: Partial<
Record<keyof GetToolInputSchemaType, string>
Record<keyof PostHogLlmPropertiesPayloadSchemaType, string>
> = {
input: '$ai_input',
outputChoices: '$ai_output_choices',
@@ -47,22 +50,27 @@ async function capturePosthogLlmObservability(
};
for (const key of Object.keys(toPostHogKey) as Array<
keyof GetToolInputSchemaType
keyof PostHogLlmPropertiesPayloadSchemaType
>) {
if (trackArgs[key] !== undefined) {
const posthogKey = toPostHogKey[key];
if (posthogKey) {
posthogProperties[posthogKey] = trackArgs[key];
// Type assertion to satisfy TS: posthogKey is a key of posthogProperties
(posthogProperties as Record<string, unknown>)[posthogKey] =
trackArgs[key];
}
}
}
// Pass validated args to the controller
const result = await posthogLlmController.capture({
// validated input for the controller
const posthogInput: PostHogLlmInputSchemaType = {
eventName: '$ai_generation',
distinctId: trackArgs.userId,
properties: posthogProperties,
});
};
// Pass validated args to the controller
const result = await PostHogController.capture(posthogInput);
methodLogger.debug(`Got the response from the controller`, result);
// Format the response for the MCP tool
@@ -96,7 +104,7 @@ function registerTools(server: McpServer) {
server.tool(
'capture_llm_observability',
`Captures LLM usage in PostHog for observability, including requests, responses, and performance metrics`,
GetToolInputSchema.shape,
PostHogLlmPropertiesPayloadSchema.shape,
capturePosthogLlmObservability,
);

62
src/types/posthog-llm.types.ts
View File

@@ -1,9 +1,9 @@
import { z } from 'zod';
/**
* Zod schema for the PostHog LLM observability tool arguments.
* Zod schema for the PostHog LLM observability tool properties (payload).
*/
export const GetToolInputSchema = z.object({
export const PostHogLlmPropertiesPayloadSchema = z.object({
userId: z.string().describe('The distinct ID of the user'),
traceId: z.string().optional().describe('The trace ID to group AI events'),
model: z.string().describe('The model used (e.g., gpt-4, claude-3, etc.)'),
@@ -34,8 +34,60 @@ export const GetToolInputSchema = z.object({
baseUrl: z.string().optional().describe('The base URL of the LLM API'),
});
export type PostHogLlmPropertiesPayloadSchemaType = z.infer<
typeof PostHogLlmPropertiesPayloadSchema
>;
/**
* TypeScript type inferred from the GetToolInputSchema Zod schema.
* This represents the optional arguments passed to the tool handler and controller.
* Zod schema for the PostHog LLM observability tool properties (api payload).
*/
export type GetToolInputSchemaType = z.infer<typeof GetToolInputSchema>;
export const PostHogLlmPropertiesApiSchema = z.object({
userId: z.string().describe('The distinct ID of the user'),
$ai_trace_id: z
.string()
.optional()
.describe('The trace ID to group AI events'),
$ai_model: z
.string()
.describe('The model used (e.g., gpt-4, claude-3, etc.)'),
$ai_provider: z
.string()
.describe('The LLM provider (e.g., openai, anthropic, etc.)'),
$ai_input: z
.any()
.optional()
.describe('The input to the LLM (messages, prompt, etc.)'),
$ai_output_choices: z.any().optional().describe('The output from the LLM'),
$ai_input_tokens: z
.number()
.optional()
.describe('The number of tokens in the input'),
$ai_output_tokens: z
.number()
.optional()
.describe('The number of tokens in the output'),
$ai_latency: z
.number()
.optional()
.describe('The latency of the LLM call in seconds'),
$ai_http_status: z
.number()
.optional()
.describe('The HTTP status code of the LLM call'),
$ai_base_url: z.string().optional().describe('The base URL of the LLM API'),
});
export type PostHogLlmPropertiesApiSchemaType = z.infer<
typeof PostHogLlmPropertiesApiSchema
>;
/**
* Zod schema for the PostHog LLM observability tool arguments (top-level).
*/
export const PostHogLlmInputSchema = z.object({
eventName: z.string().describe('The name of the event to capture.'),
distinctId: z.string().describe('The distinct ID of the user.'),
properties: PostHogLlmPropertiesApiSchema.describe('The event properties.'),
});
export type PostHogLlmInputSchemaType = z.infer<typeof PostHogLlmInputSchema>;