Added models handler to Albus

Author: roh26it

src/globals.ts

@@ -11,6 +11,7 @@ export const POSSIBLE_RETRY_STATUS_HEADERS = [
 ];
 
 export const HEADER_KEYS: Record<string, string> = {
+  API_KEY: `x-${POWERED_BY}-api-key`,
   MODE: `x-${POWERED_BY}-mode`,
   RETRIES: `x-${POWERED_BY}-retry-count`,
   PROVIDER: `x-${POWERED_BY}-provider`,
@@ -23,6 +24,7 @@ export const HEADER_KEYS: Record<string, string> = {
   REQUEST_TIMEOUT: `x-${POWERED_BY}-request-timeout`,
   STRICT_OPEN_AI_COMPLIANCE: `x-${POWERED_BY}-strict-open-ai-compliance`,
   CONTENT_TYPE: `Content-Type`,
+  VIRTUAL_KEY: `x-${POWERED_BY}-virtual-key`,
 };
 
 export const RESPONSE_HEADER_KEYS: Record<string, string> = {

src/handlers/modelsHandler.ts

@@ -1,37 +1,53 @@
-import { Context } from 'hono';
-import models from '../data/models.json';
-import providers from '../data/providers.json';
+import { Context, Next } from 'hono';
+import { HEADER_KEYS } from '../globals';
+import { env } from 'hono/adapter';
 
 /**
  * Handles the models request. Returns a list of models supported by the Ai gateway.
  * Allows filters in query params for the provider
  * @param c - The Hono context
  * @returns - The response
  */
-export async function modelsHandler(c: Context): Promise<Response> {
-  // If the request does not contain a provider query param, return all models. Add a count as well.
-  const provider = c.req.query('provider');
-  if (!provider) {
-    return c.json({
-      ...models,
-      count: models.data.length,
-    });
-  } else {
-    // Filter the models by the provider
-    const filteredModels = models.data.filter(
-      (model: any) => model.provider.id === provider
-    );
-    return c.json({
-      ...models,
-      data: filteredModels,
-      count: filteredModels.length,
-    });
+export const modelsHandler = async (context: Context, next: Next) => {
+  const fetchOptions: Record<string, any> = {};
+  fetchOptions['method'] = context.req.method;
+
+  const headers = Object.fromEntries(context.req.raw.headers);
+
+  const authHeader = headers['Authorization'] || headers['authorization'];
+
+  const apiKey =
+    headers[HEADER_KEYS.API_KEY] || authHeader?.replace('Bearer ', '');
+  let config: any = headers[HEADER_KEYS.CONFIG];
+  if (config && typeof config === 'string') {
+    try {
+      config = JSON.parse(config);
+    } catch {
+      config = {};
+    }
   }
-}
+  const providerHeader = headers[HEADER_KEYS.PROVIDER];
+  const virtualKey = headers[HEADER_KEYS.VIRTUAL_KEY];
+
+  const containsProvider =
+    providerHeader || virtualKey || config?.provider || config?.virtual_key;
+
+  if (containsProvider) {
+    return next();
+  }
+
+  // Strip gateway endpoint for models endpoint.
+  const urlObject = new URL(context.req.url);
+  const requestRoute = `${env(context).ALBUS_BASEPATH}${context.req.path.replace('/v1/', '/v2/')}${urlObject.search}`;
+  fetchOptions['headers'] = {
+    [HEADER_KEYS.API_KEY]: apiKey,
+  };
 
-export async function providersHandler(c: Context): Promise<Response> {
-  return c.json({
-    ...providers,
-    count: providers.data.length,
+  const resp = await fetch(requestRoute, fetchOptions);
+  return new Response(resp.body, {
+    status: resp.status,
+    headers: {
+      'content-type': 'application/json',
+    },
   });
-}
+};

src/index.ts

@@ -91,6 +91,9 @@ if (getRuntimeKey() === 'node') {
   app.use(logger());
 }
 
+// Support the /v1/models endpoint
+app.get('/v1/models', modelsHandler);
+
 // Use hooks middleware for all routes
 app.use('*', hooks);
 
@@ -252,9 +255,6 @@ app.post('/v1/prompts/*', requestValidator, (c) => {
   });
 });
 
-app.get('/v1/reference/models', modelsHandler);
-app.get('/v1/reference/providers', providersHandler);
-
 // WebSocket route
 if (runtime === 'workerd') {
   app.get('/v1/realtime', realTimeHandler);