move metrics to middleware

Signed-off-by: oilbeater <[email protected]>
oilbeater · Oct 1, 2024 · 1e9eb6d · 1e9eb6d
1 parent 2011047
commit 1e9eb6d
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 27 deletions.
diff --git a/src/providers/azureOpenAI.ts b/src/providers/azureOpenAI.ts
@@ -1,15 +1,18 @@
 import { Hono, Context } from 'hono';
 import { AIProvider, AIRequestParams } from '../types';
 import { generateCacheKey } from '../utils/cache';
-import { recordAnalytics, timingMiddleware } from '../utils/analytics';
+import { metricsMiddleware } from '../utils/analytics';
 import { bufferMiddleware } from '../utils/buffer';
+import { loggingMiddleware } from '../utils/logging';
+
 const BasePath = '/azure-openai/:resource_name/deployments/:deployment_name';
 const ProviderName = 'azure-openai';
-
 const azureOpenAIRoute = new Hono();
 
-azureOpenAIRoute.use(bufferMiddleware, timingMiddleware);
-azureOpenAIRoute.post('/*', async (c) => {
+azureOpenAIRoute.use(bufferMiddleware, metricsMiddleware, loggingMiddleware);
+
+azureOpenAIRoute.post('/*', async (c: Context) => {
+    c.set('endpoint', ProviderName);
     const resourceName = c.req.param('resource_name') || '';
     const deploymentName = c.req.param('deployment_name') || '';
     const functionName = c.req.path.slice(`/azure-openai/${resourceName}/deployments/${deploymentName}/`.length);
@@ -59,10 +62,7 @@ export const azureOpenAIProvider: AIProvider = {
             return c.text('Internal Server Error', 500);
         }
 
-        const startTime = Date.now();
         let buf = '';
-        let prompt_tokens = 0;
-        let completion_tokens = 0;
 
         (async () => {
             const decoder = new TextDecoder('utf-8');
@@ -74,21 +74,7 @@ export const azureOpenAIProvider: AIProvider = {
             }
 
             if (response.status === 200) {
-                if (response.headers.get('content-type') === 'application/json') {
-                    const usage = JSON.parse(buf)['usage'];
-                    if (usage) {
-                        prompt_tokens = usage['prompt_tokens'] | 0;
-                        completion_tokens = usage['completion_tokens'] | 0;
-                    }
-                } else {
-                    completion_tokens = buf.split('\n\n').length - 1;
-                }
-            }
-
-            const duration = Date.now() - startTime;
-            recordAnalytics(c, ProviderName, duration, prompt_tokens, completion_tokens);
-            if (response.status === 200) {
-                c.executionCtx.waitUntil(c.env.MALACCA_CACHE.put(cacheKeyHex, buf, {expirationTtl: 3600}));
+                c.executionCtx.waitUntil(c.env.MALACCA_CACHE.put(cacheKeyHex, buf, { expirationTtl: 3600 }));
             }
             await writer.close();
         })();

diff --git a/src/utils/analytics.ts b/src/utils/analytics.ts
@@ -16,16 +16,31 @@ export function recordAnalytics(
   }
 }
 
-export const timingMiddleware: MiddlewareHandler = async (c, next) => {
+export const metricsMiddleware: MiddlewareHandler = async (c, next) => {
   const startTime = Date.now();
 
   await next();
-  
+
   c.executionCtx.waitUntil((async () => {
     await c.get('bufferPromise')
+    const buf = c.get('buffer')
     const endTime = Date.now();
     const duration = endTime - startTime;
-    console.log(`Request duration: ${duration}ms`);
+    const endpoint = c.get('endpoint') || 'unknown';
+    let prompt_tokens = 0;
+    let completion_tokens = 0;
+    if (c.res.status === 200) {
+      if (c.res.headers.get('content-type') === 'application/json') {
+        const usage = JSON.parse(buf)['usage'];
+        if (usage) {
+          prompt_tokens = usage['prompt_tokens'] | 0;
+          completion_tokens = usage['completion_tokens'] | 0;
+        }
+      } else {
+        completion_tokens = buf.split('\n\n').length - 1;
+      }
+    }
+    recordAnalytics(c, endpoint, duration, prompt_tokens, completion_tokens);
   })());
 };
 

diff --git a/src/utils/buffer.ts b/src/utils/buffer.ts
@@ -15,7 +15,11 @@ export const bufferMiddleware: MiddlewareHandler = async (c, next) => {
   const writer = writable.getWriter();
   c.executionCtx.waitUntil((async () => {
     const reader = originalResponse.body?.getReader();
-    if (!reader) return;
+    if (!reader) {
+      c.set('buffer', buffer);
+      resolveBuffer();
+      return;
+    }
 
     try {
       while (true) {
@@ -26,7 +30,7 @@ export const bufferMiddleware: MiddlewareHandler = async (c, next) => {
         await writer.write(value);
       }
     } finally {
-      c.set('buffer', buffer)
+      c.set('buffer', buffer);
       resolveBuffer();
       await writer.close();
     }