diff --git a/src/index.ts b/src/index.ts index be5a86a..e64490e 100644 --- a/src/index.ts +++ b/src/index.ts @@ -7,16 +7,9 @@ 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) @@ -24,7 +17,6 @@ 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 \ No newline at end of file