add metrics for token

Signed-off-by: Mengxin Liu <[email protected]>

src/index.ts

@@ -7,24 +7,16 @@ const app = new Hono<{ Bindings: Bindings }>()
 const azureOpenAI = new Hono()
 
 azureOpenAI.use(async (c: Context, next) => {
-  const start = Date.now()
+  c.set('start', Date.now())
+  c.set('endpoint', 'azure-openai')
   await next()
-  const end = Date.now()
-
-  c.env.MALACCA.writeDataPoint({
-    'blobs': [c.get('endpoint'), c.req.path, c.res.status],
-    'doubles': [end-start],
-    'indexes': ['azure'],
-  })
-
 })
 
 azureOpenAI.post('/*', handleChat)
 app.get('/', (c) => c.text('Welcome to Malacca!'))
 app.route('/azure-openai/:resource_name/:deployment_name', azureOpenAI).onError((err, c) => c.text(err.message, 500))
 
 async function handleChat(c: Context) {
-  c.set("endpoint", "azure-openai")
   const resourceName = c.req.param('resource_name')
   const deploymentName = c.req.param('deployment_name')
   const functionName = c.req.path.slice(`/azure-openai/${resourceName}/${deploymentName}/`.length)
@@ -40,28 +32,38 @@ async function handleChat(c: Context) {
     headers: c.req.header()
   })
 
-  console.log(response.headers)
-  const { readable, writable } = new TransformStream();
-  const writer = writable.getWriter();
-  const reader = response.body?.getReader();
+  const { readable, writable } = new TransformStream()
+  const writer = writable.getWriter()
+  const reader = response.body?.getReader()
   if (!reader) {
     return c.text('Internal Server Error', 500)
   }
   const decoder = new TextDecoder('utf-8');
   (async () => {
     while (true) {
-      const { done, value } = await reader.read();
+      const { done, value } = await reader.read()
 
       if (done) {
         break;
       }
-      console.log(decoder.decode(value));
-      await writer.write(value);
+      if (response.headers.get('content-type') === 'application/json') {
+        const usage = JSON.parse(decoder.decode(value))['usage']
+
+        await writer.write(value)
+        const duration = Date.now() - c.get('start')
+        c.env.MALACCA.writeDataPoint({
+          'blobs': [c.get('endpoint'), c.req.path, c.res.status],
+          'doubles': [duration, usage['prompt_tokens'] | 0, usage['completion_tokens'] | 0],
+          'indexes': ['azure'],
+        })
+        console.log('finish writing data point')
+      }
     }
-    await writer.close();
+
+    await writer.close()
   })();
 
-  return new Response(readable, response);
+  return new Response(readable, response)
 }
 
 export default app