add middleware support for deepseek

Signed-off-by: oilbeater <[email protected]>
oilbeater · Oct 11, 2024 · 169b234 · 169b234
1 parent 8bb13a4
commit 169b234
Show file tree

Hide file tree

Showing 13 changed files with 240 additions and 21 deletions.
diff --git a/src/middlewares/analytics.ts b/src/middlewares/analytics.ts
@@ -1,5 +1,5 @@
 import { Context, MiddlewareHandler, Next } from 'hono';
-import { AppContext } from '.';
+import { AppContext, setMiddlewares } from '.';
 
 export function recordAnalytics(
   c: Context<AppContext>,
@@ -13,7 +13,7 @@ export function recordAnalytics(
   const getTokenCount = c.get('getTokenCount');
   const { input_tokens, output_tokens } = typeof getTokenCount === 'function' ? getTokenCount(c) : { input_tokens: 0, output_tokens: 0 };
 
-  // console.log(endpoint, c.req.path, modelName, input_tokens, output_tokens, c.get('malacca-cache-status') || 'miss', c.res.status);
+  console.log(endpoint, c.req.path, modelName, input_tokens, output_tokens, c.get('malacca-cache-status') || 'miss', c.res.status);
 
   if (c.env.MALACCA) {
     c.env.MALACCA.writeDataPoint({
@@ -25,6 +25,7 @@ export function recordAnalytics(
 }
 
 export const metricsMiddleware: MiddlewareHandler = async (c: Context<AppContext>, next: Next) => {
+  setMiddlewares(c, 'metrics');
   const startTime = Date.now();
   await next();
 

diff --git a/src/middlewares/buffer.ts b/src/middlewares/buffer.ts
@@ -1,7 +1,8 @@
 import { Context, MiddlewareHandler, Next } from 'hono'
-import { AppContext } from '.';
+import { AppContext, setMiddlewares } from '.';
 
 export const bufferMiddleware: MiddlewareHandler = async (c: Context<AppContext>, next: Next) => {
+  setMiddlewares(c, 'buffer');
   let buffer = ''
   let resolveBuffer!: () => void
   const bufferPromise = new Promise<void>((resolve) => {

diff --git a/src/middlewares/cache.ts b/src/middlewares/cache.ts
@@ -1,5 +1,5 @@
 import { Context, MiddlewareHandler, Next } from "hono";
-import { AppContext } from '.';
+import { AppContext, setMiddlewares } from '.';
 
 export async function generateCacheKey(urlWithQueryParams: string, body: string): Promise<string> {
   const cacheKey = await crypto.subtle.digest(
@@ -12,6 +12,7 @@ export async function generateCacheKey(urlWithQueryParams: string, body: string)
 }
 
 export const cacheMiddleware: MiddlewareHandler = async (c: Context<AppContext>, next: Next) => {
+  setMiddlewares(c, 'cache');
   const cacheKeyHex = await generateCacheKey(c.req.url, await c.req.text());
   const response = await c.env.MALACCA_CACHE.get(cacheKeyHex, "stream");
   if (response) {

diff --git a/src/middlewares/fallback.ts b/src/middlewares/fallback.ts
@@ -1,7 +1,8 @@
 import { Context, Next } from 'hono';
-import { AppContext } from '.';
+import { AppContext, setMiddlewares } from '.';
 
 export const fallbackMiddleware = async (c: Context<AppContext>, next: Next) => {
+    setMiddlewares(c, 'fallback');
     try {
         await next();
 

diff --git a/src/middlewares/guard.ts b/src/middlewares/guard.ts
@@ -1,5 +1,5 @@
 import { Context, MiddlewareHandler, Next } from "hono";
-import { AppContext } from ".";
+import { AppContext, setMiddlewares } from ".";
 
 const denyRequestPatterns = [
     'password',
@@ -13,6 +13,7 @@ const denyResponsePatterns = [
 // The guard middleware is used to protect the API by checking if the request match the specific regex. 
 // If so it returns message "Rejected due to inappropriate content" with 403 status code.
 export const guardMiddleware: MiddlewareHandler = async (c: Context<AppContext>, next: Next) => {
+    setMiddlewares(c, 'guard');
     const requestText = await c.req.text();
     if (denyRequestPatterns.some(pattern => new RegExp(pattern, 'i').test(requestText))) {
         return c.text('Rejected due to inappropriate content', 403);

diff --git a/src/middlewares/index.ts b/src/middlewares/index.ts
@@ -11,6 +11,7 @@ export { fallbackMiddleware } from './fallback';
 export interface AppContext {
     Bindings: Env,
     Variables: {
+        middlewares: string[],
         endpoint: string,
         'malacca-cache-status': string,
         bufferPromise: Promise<any>,
@@ -19,5 +20,14 @@ export interface AppContext {
         realKey: string,
         getModelName: (c: Context) => string,
         getTokenCount: (c: Context) => { input_tokens: number, output_tokens: number },
+        getVirtualKey: (c: Context) => string,
+    }
+}
+
+export function setMiddlewares(c: Context, name: string) {
+    if (!c.get('middlewares')) {
+        c.set('middlewares', [name]);
+    } else {
+        c.set('middlewares', [...c.get('middlewares'), name]);
     }
 }
diff --git a/src/middlewares/logging.ts b/src/middlewares/logging.ts
@@ -1,7 +1,8 @@
 import { Context, Next } from 'hono';
-import { AppContext } from '.';
+import { AppContext, setMiddlewares } from '.';
 
 export const loggingMiddleware = async (c: Context<AppContext>, next: Next) => {
+    setMiddlewares(c, 'logging');
     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 '.';
+import { AppContext, setMiddlewares } from '.';
 
 export const rateLimiterMiddleware = async (c: Context<AppContext>, next: Next) => {
+    setMiddlewares(c, 'rateLimiter');
     const key = c.req.header('api-key') || '';
     const { success } = await c.env.MY_RATE_LIMITER.limit({ key: key })
     if (!success) {

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

diff --git a/src/providers/azureOpenAI.ts b/src/providers/azureOpenAI.ts
@@ -16,13 +16,29 @@ const ProviderName = 'azure-openai';
 const azureOpenAIRoute = new Hono();
 
 const initMiddleware = async (c: Context, next: Next) => {
+    if (!c.get('middlewares')) {
+        c.set('middlewares', ['init']);
+    } else {
+        c.set('middlewares', [...c.get('middlewares'), 'init']);
+    }
     c.set('endpoint', ProviderName);
     c.set('getModelName', getModelName);
     c.set('getTokenCount', getTokenCount);
+    c.set('getVirtualKey', getVirtualKey);
     await next();
 };
 
-azureOpenAIRoute.use(initMiddleware, metricsMiddleware, loggingMiddleware, bufferMiddleware, virtualKeyMiddleware, rateLimiterMiddleware, guardMiddleware, cacheMiddleware, fallbackMiddleware);
+azureOpenAIRoute.use(
+    initMiddleware,
+    metricsMiddleware,
+    loggingMiddleware,
+    bufferMiddleware,
+    virtualKeyMiddleware,
+    rateLimiterMiddleware,
+    guardMiddleware,
+    cacheMiddleware,
+    fallbackMiddleware
+);
 
 azureOpenAIRoute.post('/*', async (c: Context) => {
     return azureOpenAIProvider.handleRequest(c);
@@ -43,9 +59,11 @@ export const azureOpenAIProvider: AIProvider = {
         const urlWithQueryParams = `${azureEndpoint}?${queryParams}`;
 
         const headers = new Headers(c.req.header());
-        const apiKey: string = c.get('realKey');
-        if (apiKey) {
-            headers.set('api-key', apiKey);
+        if (c.get('middlewares')?.includes('virtualKey')) {
+            const apiKey: string = c.get('realKey');
+            if (apiKey) {
+                headers.set('api-key', apiKey);
+            }
         }
         const response = await fetch(urlWithQueryParams, {
             method: c.req.method,
@@ -107,4 +125,8 @@ function getTokenCount(c: Context): { input_tokens: number, output_tokens: numbe
         }
     }
     return { input_tokens: 0, output_tokens: 0 }
+}
+
+function getVirtualKey(c: Context): string {
+    return c.req.header('api-key') || '';
 }
diff --git a/src/providers/deepseek.ts b/src/providers/deepseek.ts
@@ -19,9 +19,22 @@ const initMiddleware = async (c: Context, next: Next) => {
     c.set('endpoint', ProviderName);
     c.set('getModelName', getModelName);
     c.set('getTokenCount', getTokenCount);
+    c.set('getVirtualKey', getVirtualKey);
     await next();
 };
 
+deepseekRoute.use(
+    initMiddleware,
+    metricsMiddleware,
+    loggingMiddleware,
+    bufferMiddleware,
+    virtualKeyMiddleware,
+    rateLimiterMiddleware,
+    guardMiddleware,
+    cacheMiddleware,
+    fallbackMiddleware
+);
+
 deepseekRoute.post('/*', async (c: Context) => {
     return deepseekProvider.handleRequest(c);
 });
@@ -39,28 +52,37 @@ export const deepseekProvider: AIProvider = {
     handleRequest: async (c: Context) => {
         const functionName = c.req.path.slice(`/deepseek/`.length);
         const deepseekEndpoint = `https://api.deepseek.com/${functionName}`;
-        console.log(`DeepSeek endpoint: ${deepseekEndpoint}`);
+
+        const headers = new Headers(c.req.header());
+        if (c.get('middlewares')?.includes('virtualKey')) {
+            const apiKey: string = c.get('realKey');
+            if (apiKey) {
+                headers.set('Authorization', `Bearer ${apiKey}`);
+            }
+        }
 
         const response = await fetch(deepseekEndpoint, {
             method: c.req.method,
             body: JSON.stringify(await c.req.json()),
-            headers: c.req.header()
+            headers: headers
         });
 
         return response;
     }
 };
 
 function getModelName(c: Context): string {
-    const model = c.req.param('model');
+    const body = c.get('reqBuffer') || '{}';
+    const model = JSON.parse(body).model;
     return model || "unknown";
 }
 
 function getTokenCount(c: Context): { input_tokens: number, output_tokens: number } {
     const buf = c.get('buffer') || "";
     if (c.res.status === 200) {
-        try {
-            const jsonResponse = JSON.parse(buf);
+        if (c.res.headers.get('content-type') === 'application/json') {
+            try {
+                const jsonResponse = JSON.parse(buf);
             const usage = jsonResponse.usage;
             if (usage) {
                 return {
@@ -69,8 +91,25 @@ function getTokenCount(c: Context): { input_tokens: number, output_tokens: numbe
                 };
             }
         } catch (error) {
-            console.error("Error parsing response:", error);
+                console.error("Error parsing response:", error);
+            }
         }
-    }
+        else {
+            const output = buf.trim().split('\n\n').at(-2);
+            if (output && output.startsWith('data: ')) {
+                const usage_message = JSON.parse(output.slice('data: '.length));
+                return {
+                    input_tokens: usage_message.usage.prompt_tokens || 0,
+                    output_tokens: usage_message.usage.completion_tokens || 0
+                };
+            }
+        }
+    } 
     return { input_tokens: 0, output_tokens: 0 };
 }
+
+function getVirtualKey(c: Context): string {
+    const authHeader = c.req.header('Authorization') || '';
+    return authHeader.startsWith('Bearer ') ? authHeader.slice(7) : '';
+}
+
diff --git a/src/providers/workersAI.ts b/src/providers/workersAI.ts
@@ -37,3 +37,7 @@ function getModelName(c: Context) {
 function getTokenCount(c: Context) {
     return { input_tokens: 0, output_tokens: 0 };
 }
+
+function getVirtualKey(c: Context) {
+    return '';
+}