add model name and cache status to metrics

Signed-off-by: oilbeater <[email protected]>

src/middlewares/analytics.ts

@@ -8,10 +8,12 @@ export function recordAnalytics(
   completion_tokens: number
 ) {
   if (c.env.MALACCA) {
+    const getModelName = c.get('getModelName');
+    const modelName = typeof getModelName === 'function' ? getModelName(c) : 'unknown';
     c.env.MALACCA.writeDataPoint({
-      'blobs': [endpoint, c.req.path, c.res.status],
+      'blobs': [endpoint, c.req.path, c.res.status, c.get('malacca-cache-status') || 'miss', modelName],
       'doubles': [duration, prompt_tokens, completion_tokens],
-      'indexes': ['azure'],
+      'indexes': [endpoint],
     });
   }
 }

src/middlewares/cache.ts

@@ -20,6 +20,7 @@ export const cacheMiddleware: MiddlewareHandler = async (c: Context, next: Next)
   if (response) {
     const { _, metadata } = await c.env.MALACCA_CACHE.getWithMetadata(cacheKeyHex, "stream");
     const contentType = metadata['contentType'] || 'application/octet-stream';
+    c.set('malacca-cache-status', 'hit');
     return new Response(response, { headers: { 'malacca-cache-status': 'hit', 'content-type': contentType } });
   }
 

src/providers/azureOpenAI.ts

@@ -14,11 +14,12 @@ const azureOpenAIRoute = new Hono();
 
 const initMiddleware = async (c: Context, next: Next) => {
     c.set('endpoint', ProviderName);
+    c.set('getModelName', getModelName);
     await next();
 };
 
 
-azureOpenAIRoute.use(initMiddleware, bufferMiddleware, metricsMiddleware, loggingMiddleware, virtualKeyMiddleware, cacheMiddleware);
+azureOpenAIRoute.use(initMiddleware, metricsMiddleware, loggingMiddleware, bufferMiddleware, virtualKeyMiddleware, cacheMiddleware);
 
 azureOpenAIRoute.post('/*', async (c: Context) => {
     return azureOpenAIProvider.handleRequest(c);
@@ -28,6 +29,7 @@ export const azureOpenAIProvider: AIProvider = {
     name: ProviderName,
     basePath: BasePath,
     route: azureOpenAIRoute,
+    getModelName: getModelName,
     handleRequest: async (c: Context) => {
         const resourceName = c.req.param('resource_name') || '';
         const deploymentName = c.req.param('deployment_name') || '';
@@ -50,3 +52,30 @@ export const azureOpenAIProvider: AIProvider = {
         return response;
     }
 };
+
+function getModelName(c: Context): string {
+    if (c.res.status === 200) {
+        const buf = c.get('buffer') || ""
+        if (c.res.headers.get('content-type') === 'application/json') {
+            const model = JSON.parse(buf)['model'];
+            if (model) {
+                return model
+            }
+        } else {
+            const chunks = buf.split('\n\n');
+            for (const chunk of chunks) {
+                if (chunk.startsWith('data: ')) {
+                    const jsonStr = chunk.slice(6);
+                    try {
+                        const jsonData = JSON.parse(jsonStr);
+                        if (jsonData.model != "") {
+                            return jsonData.model;
+                        }
+                    } catch (error) {
+                    }
+                }
+            }
+        }
+    }
+    return ""
+}

src/types.ts

@@ -9,6 +9,7 @@ export interface Bindings {
 export interface AIProvider {
     name: string;
     handleRequest: (c: Context) => Promise<Response>;
+    getModelName: (c: Context) => string;
     basePath: string;
     route: Hono;
 }