add more type hints

Signed-off-by: oilbeater <[email protected]>
oilbeater · Oct 2, 2024 · ed8ab0f · ed8ab0f
1 parent 09ac428
commit ed8ab0f
Show file tree

Hide file tree

Showing 7 changed files with 36 additions and 16 deletions.
diff --git a/src/middlewares/analytics.ts b/src/middlewares/analytics.ts
@@ -1,7 +1,8 @@
-import { Context, MiddlewareHandler } from 'hono';
+import { Context, MiddlewareHandler, Next } from 'hono';
+import { AppContext } from './index';
 
 export function recordAnalytics(
-  c: Context,
+  c: Context<AppContext>,
   endpoint: string,
   duration: number,
 ) {
@@ -16,20 +17,20 @@ export function recordAnalytics(
 
   if (c.env.MALACCA) {
     c.env.MALACCA.writeDataPoint({
-      'blobs': [endpoint, c.req.path, c.res.status, c.get('malacca-cache-status') || 'miss', modelName],
+      'blobs': [endpoint, c.req.path, c.res.status.toString(), c.get('malacca-cache-status') || 'miss', modelName],
       'doubles': [duration, input_tokens, output_tokens],
       'indexes': [endpoint],
     });
   }
 }
 
-export const metricsMiddleware: MiddlewareHandler = async (c, next) => {
+export const metricsMiddleware: MiddlewareHandler = async (c: Context<AppContext>, next: Next) => {
   const startTime = Date.now();
-
+  c.env.MALACCA
   await next();
 
   c.executionCtx.waitUntil((async () => {
-    await c.get('bufferPromise')
+    await c.get('bufferPromise');
     const endTime = Date.now();
     const duration = endTime - startTime;
     const endpoint = c.get('endpoint') || 'unknown';

diff --git a/src/middlewares/buffer.ts b/src/middlewares/buffer.ts
@@ -1,4 +1,5 @@
 import { Context, MiddlewareHandler, Next } from 'hono'
+import { AppContext } from './index';
 
 export const bufferMiddleware: MiddlewareHandler = async (c: Context, next: Next) => {
   let buffer = ''

diff --git a/src/middlewares/cache.ts b/src/middlewares/cache.ts
@@ -1,8 +1,5 @@
 import { Context, MiddlewareHandler, Next } from "hono";
-
-interface Env {
-  MALACCA_CACHE: KVNamespace;
-}
+import { AppContext } from './index';
 
 export async function generateCacheKey(urlWithQueryParams: string, body: string): Promise<string> {
   const cacheKey = await crypto.subtle.digest(
@@ -14,13 +11,14 @@ export async function generateCacheKey(urlWithQueryParams: string, body: string)
     .join('');
 }
 
-export const cacheMiddleware: MiddlewareHandler = async (c: Context, next: Next) => {
+export const cacheMiddleware: MiddlewareHandler = async (c: Context<AppContext>, next: Next) => {
   const cacheKeyHex = await generateCacheKey(c.req.url, await c.req.text());
   const response = await c.env.MALACCA_CACHE.get(cacheKeyHex, "stream");
   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');
+    console.log(contentType);
     return new Response(response, { headers: { 'malacca-cache-status': 'hit', 'content-type': contentType } });
   }
 

diff --git a/src/middlewares/index.ts b/src/middlewares/index.ts
@@ -1,6 +1,23 @@
+import { Context } from 'hono';
+
 export { cacheMiddleware } from './cache';
 export { metricsMiddleware } from './analytics';
 export { bufferMiddleware } from './buffer';
 export { loggingMiddleware } from './logging';
 export { virtualKeyMiddleware } from './virtualKey';
 export { rateLimiterMiddleware } from './rateLimiter';
+import { Bindings } from '../types';
+
+export type AppContext = {
+    Bindings: Bindings,
+    Variables: {
+        endpoint: string,
+        'malacca-cache-status': string,
+        bufferPromise: Promise<any>,
+        buffer: string,
+        reqBuffer: string,
+        realKey: string,
+        getModelName: (c: Context) => string,
+        getTokenCount: (c: Context) => { input_tokens: number, output_tokens: number },
+    }
+}
diff --git a/src/middlewares/logging.ts b/src/middlewares/logging.ts
@@ -1,6 +1,7 @@
 import { Context, Next } from 'hono';
+import { AppContext } from './index';
 
-export const loggingMiddleware = async (c: Context, next: Next) => {
+export const loggingMiddleware = async (c: Context<AppContext>, next: Next) => {
     await next();
 
     // Log request and response

diff --git a/src/middlewares/rateLimiter.ts b/src/middlewares/rateLimiter.ts
@@ -1,7 +1,8 @@
 import { Context, Next } from "hono";
+import { AppContext } from './index';
 
-export const rateLimiterMiddleware = async (c: Context, next: Next) => {
-    const key = c.req.header('api-key');
+export const rateLimiterMiddleware = async (c: Context<AppContext>, next: Next) => {
+    const key = c.req.header('api-key') || '';
     const { success } = await c.env.MY_RATE_LIMITER.limit({ key: key })
     if (!success) {
       return new Response(`429 Failure – rate limit exceeded`, { status: 429 })

diff --git a/src/middlewares/virtualKey.ts b/src/middlewares/virtualKey.ts
@@ -1,7 +1,8 @@
 import { Context, Next } from "hono";
+import { AppContext } from './index';
 
-export const virtualKeyMiddleware = async (c: Context, next: Next) => {
-    const apiKey = c.req.header('api-key');
+export const virtualKeyMiddleware = async (c: Context<AppContext>, next: Next) => {
+    const apiKey = c.req.header('api-key') || '';
     const realKey = await c.env.MALACCA_USER.get(apiKey);
     if (!realKey) {
         return c.text('Unauthorized', 401);