From 21af866fd9c13cb1aa809b3fe60b511e1a7f80c2 Mon Sep 17 00:00:00 2001
From: nai-degen <khoners@protonmail.com>
Date: Sun, 10 Sep 2023 19:07:56 -0500
Subject: [PATCH] moves keystore interface

---
 src/shared/key-management/anthropic/provider.ts | 3 ++-
 src/shared/key-management/index.ts              | 6 ------
 src/shared/key-management/openai/provider.ts    | 3 ++-
 src/shared/key-management/stores/firebase.ts    | 3 ++-
 src/shared/key-management/stores/index.ts       | 8 ++++++++
 src/shared/key-management/stores/memory.ts      | 5 +++--
 6 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/src/shared/key-management/anthropic/provider.ts b/src/shared/key-management/anthropic/provider.ts
index 266e533..04fdf7b 100644
--- a/src/shared/key-management/anthropic/provider.ts
+++ b/src/shared/key-management/anthropic/provider.ts
@@ -1,8 +1,9 @@
 import crypto from "crypto";
-import { Key, KeyProvider, KeyStore } from "..";
+import { Key, KeyProvider } from "..";
 import { config } from "../../../config";
 import { logger } from "../../../logger";
 import type { AnthropicModelFamily } from "../../models";
+import { KeyStore } from "../stores";
 import { AnthropicKeyChecker } from "./checker";
 
 // https://docs.anthropic.com/claude/reference/selecting-a-model
diff --git a/src/shared/key-management/index.ts b/src/shared/key-management/index.ts
index cf31fe8..97ed10c 100644
--- a/src/shared/key-management/index.ts
+++ b/src/shared/key-management/index.ts
@@ -65,12 +65,6 @@ export interface KeyProvider<T extends Key = Key> {
   recheck(): void;
 }
 
-export interface KeyStore<T extends Pick<Key, "key">> {
-  load(): Promise<T[]>;
-  add(key: T): void;
-  update(key: T): void;
-}
-
 export let keyPool: KeyPool;
 export async function init() {
   keyPool = new KeyPool();
diff --git a/src/shared/key-management/openai/provider.ts b/src/shared/key-management/openai/provider.ts
index 80f4fe9..6bc44c5 100644
--- a/src/shared/key-management/openai/provider.ts
+++ b/src/shared/key-management/openai/provider.ts
@@ -3,10 +3,11 @@ round-robin access to keys. Keys are stored in the OPENAI_KEY environment
 variable as a comma-separated list of keys. */
 import crypto from "crypto";
 import http from "http";
-import { Key, KeyProvider, Model, KeyStore } from "../index";
+import { Key, KeyProvider, Model } from "../index";
 import { config } from "../../../config";
 import { logger } from "../../../logger";
 import { OpenAIKeyChecker } from "./checker";
+import { KeyStore } from "../stores";
 import { getOpenAIModelFamily, OpenAIModelFamily } from "../../models";
 
 export type OpenAIModel =
diff --git a/src/shared/key-management/stores/firebase.ts b/src/shared/key-management/stores/firebase.ts
index c97377f..ba9c113 100644
--- a/src/shared/key-management/stores/firebase.ts
+++ b/src/shared/key-management/stores/firebase.ts
@@ -1,5 +1,6 @@
 import type firebase from "firebase-admin";
-import { AIService, Key, KeyStore } from "..";
+import { KeyStore } from ".";
+import { AIService, Key } from "..";
 import { getFirebaseApp } from "../../../config";
 
 export class FirebaseKeyStore<K extends Pick<Key, "key">>
diff --git a/src/shared/key-management/stores/index.ts b/src/shared/key-management/stores/index.ts
index eee6299..1c13c87 100644
--- a/src/shared/key-management/stores/index.ts
+++ b/src/shared/key-management/stores/index.ts
@@ -1,2 +1,10 @@
+import { Key } from "..";
+
 export { FirebaseKeyStore } from "./firebase";
 export { MemoryKeyStore } from "./memory";
+
+export interface KeyStore<T extends Pick<Key, "key">> {
+  load(): Promise<T[]>;
+  add(key: T): void;
+  update(key: T): void;
+}
diff --git a/src/shared/key-management/stores/memory.ts b/src/shared/key-management/stores/memory.ts
index 04d63bb..da2b298 100644
--- a/src/shared/key-management/stores/memory.ts
+++ b/src/shared/key-management/stores/memory.ts
@@ -1,4 +1,5 @@
-import { APIFormat, Key, KeyStore } from "..";
+import { KeyStore } from ".";
+import { APIFormat, Key } from "..";
 
 export class MemoryKeyStore<K extends Pick<Key, "key">> implements KeyStore<K> {
   private env: string;
@@ -26,7 +27,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));
+    return bareKeys.map((key) => ({ key } as K)); // TODO: remove assertion
   }
 
   public add(_key: K) {}