wip broken shit

src/shared/key-management/anthropic/provider.ts

@@ -1,5 +1,5 @@
 import crypto from "crypto";
-import { Key, KeyProvider } from "..";
+import { BaseSerializableKey, Key, KeyProvider } from "..";
 import { config } from "../../../config";
 import { logger } from "../../../logger";
 import type { AnthropicModelFamily } from "../../models";
@@ -21,10 +21,8 @@ type AnthropicKeyUsage = {
 };
 
 const SERIALIZABLE_FIELDS = ["key", "service", "hash", "claudeTokens"] as const;
-type SerializableAnthropicKey = Partial<
-  Pick<AnthropicKey, (typeof SERIALIZABLE_FIELDS)[number]>
-> &
-  Pick<AnthropicKey, "key">;
+type SerializableAnthropicKey = BaseSerializableKey &
+  Partial<Pick<AnthropicKey, (typeof SERIALIZABLE_FIELDS)[number]>>;
 
 export type AnthropicKeyUpdate = Omit<
   Partial<AnthropicKey>,

src/shared/key-management/index.ts

@@ -39,6 +39,10 @@ export interface Key {
   hash: string;
 }
 
+export interface BaseSerializableKey {
+  key: string;
+}
+
 /*
 KeyPool and KeyProvider's similarities are a relic of the old design where
 there was only a single KeyPool for OpenAI keys. Now that there are multiple

src/shared/key-management/stores/firebase.ts

@@ -1,15 +1,22 @@
 import type firebase from "firebase-admin";
-import { KeyStore } from ".";
-import { AIService, Key } from "..";
 import { getFirebaseApp } from "../../../config";
+import { logger } from "../../../logger";
+import { KeyDeserializer, KeyStore, MemoryKeyStore, getDeserializer } from ".";
+import { AIService, BaseSerializableKey } from "..";
 
-export class FirebaseKeyStore<K extends Pick<Key, "key">>
+export class FirebaseKeyStore<K extends BaseSerializableKey>
   implements KeyStore<K>
 {
   private db: firebase.database.Database;
+  private service: AIService;
+  private log: typeof logger;
+  private deserializer: KeyDeserializer;
 
   constructor(service: AIService, app = getFirebaseApp()) {
     this.db = app.database();
+    this.service = service;
+    this.log = logger.child({ module: "key-store", service });
+    this.deserializer = getDeserializer(service);
   }
 
   public async load() {
@@ -24,4 +31,24 @@ export class FirebaseKeyStore<K extends Pick<Key, "key">>
   public update(key: K) {
     throw new Error("Method not implemented.");
   }
+
+  private async migrate() {
+    this.log.info("Migrating keys from environment to Firebase.");
+    const envStore = new MemoryKeyStore(this.service);
+    const keysRef = this.db.ref(`keys/${this.service}`);
+    const updates: Record<string, K> = {};
+
+    const keys = await envStore.load();
+
+    keys.forEach((key) => {
+      updates[key.key] = this.deserializer(key);
+    });
+
+    // envStore.load().then((keys) => {
+    //   keys.forEach((key) => {
+    //     updates[key.key] = key;
+    //   });
+    //   keysRef.update(updates);
+    // });
+  }
 }

src/shared/key-management/stores/index.ts

@@ -1,4 +1,6 @@
-import { Key } from "..";
+import { AIService, Key } from "..";
+import { AnthropicKeyProvider } from "../anthropic/provider";
+import { OpenAIKeyProvider } from "../openai/provider";
 
 export { FirebaseKeyStore } from "./firebase";
 export { MemoryKeyStore } from "./memory";
@@ -8,3 +10,23 @@ export interface KeyStore<T extends Pick<Key, "key">> {
   add(key: T): void;
   update(key: T): void;
 }
+
+interface BaseSerializableKey {
+  key: string;
+}
+
+export type KeyDeserializer =
+  | typeof AnthropicKeyProvider.deserialize
+  | typeof OpenAIKeyProvider.deserialize;
+
+export function getDeserializer(service: AIService): KeyDeserializer {
+  switch (service) {
+    case "anthropic":
+      return AnthropicKeyProvider.deserialize;
+    case "openai":
+      return OpenAIKeyProvider.deserialize;
+    default:
+      const never: never = service;
+      throw new Error(`Unknown service: ${never}`);
+  }
+}

src/shared/key-management/stores/memory.ts

@@ -1,8 +1,11 @@
-import { KeyStore } from ".";
-import { APIFormat, Key } from "..";
+import { KeyDeserializer, KeyStore, getDeserializer } from ".";
+import { APIFormat, BaseSerializableKey } from "..";
 
-export class MemoryKeyStore<K extends Pick<Key, "key">> implements KeyStore<K> {
+export class MemoryKeyStore<K extends BaseSerializableKey>
+  implements KeyStore<K>
+{
   private env: string;
+  private deserializer: KeyDeserializer;
 
   constructor(service: APIFormat) {
     switch (service) {
@@ -20,6 +23,7 @@ export class MemoryKeyStore<K extends Pick<Key, "key">> implements KeyStore<K> {
         const never: never = service;
         throw new Error(`Unknown service: ${never}`);
     }
+    this.deserializer = getDeserializer(service);
   }
 
   public async load() {
@@ -27,7 +31,7 @@ export class MemoryKeyStore<K extends Pick<Key, "key">> implements KeyStore<K> {
     bareKeys = [
       ...new Set(process.env[this.env]?.split(",").map((k) => k.trim())),
     ];
-    return bareKeys.map((key) => ({ key } as K)); // TODO: remove assertion
+    return bareKeys.map((key) => this.deserializer({ key }));
   }
 
   public add(_key: K) {}