diff --git a/src/shared/key-management/anthropic/provider.ts b/src/shared/key-management/anthropic/provider.ts index 04fdf7b..7b7dc14 100644 --- a/src/shared/key-management/anthropic/provider.ts +++ b/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 -> & - Pick; +type SerializableAnthropicKey = BaseSerializableKey & + Partial>; export type AnthropicKeyUpdate = Omit< Partial, diff --git a/src/shared/key-management/index.ts b/src/shared/key-management/index.ts index 97ed10c..f425f19 100644 --- a/src/shared/key-management/index.ts +++ b/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 diff --git a/src/shared/key-management/stores/firebase.ts b/src/shared/key-management/stores/firebase.ts index ba9c113..6edc4d7 100644 --- a/src/shared/key-management/stores/firebase.ts +++ b/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> +export class FirebaseKeyStore implements KeyStore { 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> 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 = {}; + + 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); + // }); + } } diff --git a/src/shared/key-management/stores/index.ts b/src/shared/key-management/stores/index.ts index 1c13c87..a2a494f 100644 --- a/src/shared/key-management/stores/index.ts +++ b/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> { 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}`); + } +} diff --git a/src/shared/key-management/stores/memory.ts b/src/shared/key-management/stores/memory.ts index da2b298..0838b18 100644 --- a/src/shared/key-management/stores/memory.ts +++ b/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> implements KeyStore { +export class MemoryKeyStore + implements KeyStore +{ private env: string; + private deserializer: KeyDeserializer; constructor(service: APIFormat) { switch (service) { @@ -20,6 +23,7 @@ export class MemoryKeyStore> implements KeyStore { const never: never = service; throw new Error(`Unknown service: ${never}`); } + this.deserializer = getDeserializer(service); } public async load() { @@ -27,7 +31,7 @@ export class MemoryKeyStore> implements KeyStore { 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) {}