feat: refactor API keys management

- add general class SecureStorage for keyring one and file based one
- implement file based secure storage
- add --show-keys command to see added API keys for all providers
- update readme with some issues resolved and some more precice simitation description
- add getting embedding model provider api separately from main model provider
- intoduce 'weak model' that may handle some easy LLM tasks

Author: erik-balfe

src/ai/agent.ts

@@ -1,7 +1,7 @@
 import type { ChatMessage } from "llamaindex";
 import { OpenAIAgent } from "llamaindex";
 import { DynamicContextData, gatherContextData } from "../cli/contextUtils";
-import { LLM_ID, MESSAGE_ROLES, MODEL_PRICES } from "../constants";
+import { MESSAGE_ROLES, MODELS, WEAK_MODEL_ID } from "../constants";
 import { getCorrectness, getFaithfulness, getRelevancy } from "../features/userScore/evaluations/evaluations";
 import { getUserEvaluationScore } from "../features/userScore/getUserEvaluationScore";
 import { askUserCallback } from "../tools/askUser";
@@ -16,7 +16,7 @@ import {
 } from "../utils/database";
 import { formatAgentMessage } from "../utils/formatting";
 import { generateConversationTitle } from "../utils/generateConversationTitle";
-import { getApiKeyForModel, getOrPromptForAPIKey } from "../utils/getOrPromptForAPIKey";
+import { getOrPromptForAPIKey } from "../utils/getOrPromptForAPIKey";
 import { UsageCostResult } from "../utils/interface";
 import { logger, LogLevel, LogLevelType } from "../utils/logger";
 import { parseLLMResponse } from "../utils/parseLLMResponse";
@@ -107,7 +107,7 @@ export async function runAgent(
   if (input === "<input_aborted_by_user />") {
     input = "<service_message>User expressed explicit intent to exit the program.</service_message>";
   }
-  const apiKey = await getApiKeyForModel(model);
+  const apiKey = await getOrPromptForAPIKey(model);
   if (!apiKey) {
     logger.error("No API key found");
     throw new Error("LLM API key not found. Please run the application again to set it up.");
@@ -278,8 +278,8 @@ async function finalizeAgentRun(
 
   logger.debug("Generating conversation title");
   const title = await generateConversationTitle(fullConversation, {
-    apiKey: await getOrPromptForAPIKey(LLM_ID),
-    modelName: LLM_ID,
+    apiKey: await getOrPromptForAPIKey(WEAK_MODEL_ID),
+    modelName: WEAK_MODEL_ID,
   });
 
   logger.debug("Getting user evaluation score");
@@ -306,7 +306,7 @@ async function finalizeAgentRun(
 
 function countUsageCost(usage: Record<string, number>, model: string): UsageCostResult {
   // Get model pricing from the current model being used
-  const modelConfig = MODEL_PRICES[model] ?? MODEL_PRICES["gpt-4o-mini"]; // fallback to a default model
+  const modelConfig = MODELS[model] ?? MODELS["gpt-4o-mini"]; // fallback to a default model
 
   // Normalize token counts from different possible field names
   const inputTokens = usage.input_tokens || usage.prompt_tokens || 0;

src/ai/retrieval/vectorStore.ts

@@ -7,7 +7,7 @@ import {
   storageContextFromDefaults,
   VectorStoreIndex,
 } from "llamaindex";
-import { EMBEDDINGS_MODEL_ID, LLM_ID, VECTOR_STORE_PATH } from "../../constants";
+import { EMBEDDINGS_MODEL_ID, VECTOR_STORE_PATH, WEAK_MODEL_ID } from "../../constants";
 import { getOrPromptForAPIKey } from "../../utils/getOrPromptForAPIKey";
 import { logger } from "../../utils/logger";
 
@@ -16,16 +16,17 @@ let vectorStoreIndex: VectorStoreIndex | null = null;
 export async function initializeVectorStoreIndex(): Promise<VectorStoreIndex> {
   const storageContext = await storageContextFromDefaults({ persistDir: VECTOR_STORE_PATH });
 
-  const apiKey = await getOrPromptForAPIKey(LLM_ID);
+  const embeddingProviderApiKey = await getOrPromptForAPIKey(EMBEDDINGS_MODEL_ID);
+  const weakModelApiKey = await getOrPromptForAPIKey(WEAK_MODEL_ID);
   let index: VectorStoreIndex;
 
   try {
     logger.debug("Attempting to load existing index");
     index = await VectorStoreIndex.init({
       storageContext,
       serviceContext: serviceContextFromDefaults({
-        embedModel: new OpenAIEmbedding({ apiKey, model: EMBEDDINGS_MODEL_ID }),
-        llm: new OpenAI({ apiKey, model: LLM_ID }),
+        embedModel: new OpenAIEmbedding({ apiKey: embeddingProviderApiKey, model: EMBEDDINGS_MODEL_ID }),
+        llm: new OpenAI({ apiKey: weakModelApiKey, model: WEAK_MODEL_ID }),
       }),
     });
     logger.debug("Existing index loaded successfully");

src/constants.ts

@@ -1,8 +1,6 @@
 import path from "path";
 import { getUserHomeDir } from "./utils/getUserHomeDir";
 
-export const LLM_ID = "gpt-4o-mini";
-
 export const API_KEY_PROMPTS = {
   OPENAI: "Please enter your OpenAI API key: ",
   GROQ: "Please enter your Groq API key: ",
@@ -30,6 +28,7 @@ export const USER_PREFS_FILE_NAME = "user_preferences.json";
 export const APP_CONFIG_FILE_PATH = path.join(CONFIG_DIR_PATH, APP_CONFIG_FILE_NAME);
 export const USER_PREFS_FILE_PATH = path.join(CONFIG_DIR_PATH, USER_PREFS_FILE_NAME);
 export const EMBEDDINGS_MODEL_ID = "text-embedding-3-small";
+export const EMBEDDING_MODELS = [EMBEDDINGS_MODEL_ID];
 export const VECTOR_STORE_PATH = path.join(CONFIG_DIR_PATH, "vector_store");
 
 export const MESSAGE_ROLES = {
@@ -69,7 +68,7 @@ interface AIModelConfig {
   default?: boolean;
 }
 
-export const MODEL_PRICES: Record<string, AIModelConfig> = {
+export const MODELS: Record<string, AIModelConfig> = {
   "gpt-4o-mini": {
     id: "gpt-4o-mini",
     friendlyName: "openAIGpt4oMiniModel",
@@ -171,3 +170,5 @@ export const CONTEXT_ALLOCATION = {
     maxChars: 20000,
   },
 };
+
+export const WEAK_MODEL_ID = "gpt-4o-mini";

src/index.ts

@@ -4,7 +4,7 @@ import { APIError } from "openai";
 import { agentLoop } from "./ai/agent";
 import { formatApiKey, getAllStoredKeys } from "./cli/getAllKeys";
 import { parseArguments, printHelp } from "./cli/interface";
-import { MAX_INPUT_LENGTH } from "./constants";
+import { EMBEDDINGS_MODEL_ID, MAX_INPUT_LENGTH } from "./constants";
 import { deleteAPIKey } from "./utils/apiKeyManager";
 import { loadConfig, saveConfig } from "./utils/config";
 import { initializeDatabase } from "./utils/database";
@@ -45,6 +45,11 @@ async function main() {
       appConfig.model = model;
     }
 
+    await getOrPromptForAPIKey(EMBEDDINGS_MODEL_ID, {
+      prePromptText:
+        "Please enter your OpenAI API key. It is used for embeddings and will consume small amount of tokens",
+    });
+
     logger.info(`running using model "${appConfig.model}"`);
 
     logger.setLevel(appConfig.logLevel);
@@ -62,7 +67,7 @@ async function main() {
 
     if (resetKey) {
       const provider = getProviderFromModel(appConfig.model);
-      deleteAPIKey(provider);
+      await deleteAPIKey(provider);
       console.log(
         chalk.green("API key has been deleted. You will be prompted for a new key on the next run."),
       );
@@ -90,8 +95,8 @@ async function main() {
     } catch (error) {
       if (error instanceof APIError) {
         if (error.status === 401) {
-          console.error(chalk.red("Invalid OpenAI API key." + error.message));
-          deleteAPIKey(appConfig.model);
+          console.error(chalk.red("Invalid API key." + error.message));
+          await deleteAPIKey(appConfig.model);
           await getOrPromptForAPIKey(appConfig.model);
         }
       }

src/utils/apiKeyManager.ts

@@ -127,9 +127,9 @@ class StorageManager {
 const storageManager = new StorageManager();
 
 // Export the methods
-export const getAPIKey = (provider: APIProvider) => storageManager.getAPIKey(provider);
-export const storeAPIKey = (apiKey: string, provider: APIProvider) =>
+export const getAPIKey = async (provider: APIProvider) => storageManager.getAPIKey(provider);
+export const storeAPIKey = async (apiKey: string, provider: APIProvider) =>
   storageManager.storeAPIKey(apiKey, provider);
-export const deleteAPIKey = (modelId: string) => storageManager.deleteAPIKey(modelId);
-export const checkStorageAvailability = () => storageManager.checkStorageAvailability();
-export const getStorageDiagnostics = () => storageManager.getStorageDiagnostics();
+export const deleteAPIKey = async (modelId: string) => storageManager.deleteAPIKey(modelId);
+export const checkStorageAvailability = async () => storageManager.checkStorageAvailability();
+export const getStorageDiagnostics = async () => storageManager.getStorageDiagnostics();

src/utils/config.ts

@@ -1,6 +1,6 @@
 import { execSync } from "child_process";
 import fs from "fs";
-import { APP_CONFIG_FILE_PATH, CONFIG_DIR_PATH, MODEL_PRICES, USER_PREFS_FILE_PATH } from "../constants";
+import { APP_CONFIG_FILE_PATH, CONFIG_DIR_PATH, MODELS, USER_PREFS_FILE_PATH } from "../constants";
 import { LogLevel, LogLevelType } from "./logger";
 
 export interface AppConfig {
@@ -17,7 +17,7 @@ export interface UserPreferences {
 
 const DEFAULT_APP_CONFIG: AppConfig = {
   logLevel: LogLevel.WARN,
-  model: Object.values(MODEL_PRICES).find((model) => model.default)?.id || "",
+  model: Object.values(MODELS).find((model) => model.default)?.id || "",
 };
 
 const DEFAULT_USER_PREFS: UserPreferences = {

src/utils/embeddings.ts

@@ -1,5 +1,5 @@
 import { OpenAIEmbedding } from "llamaindex";
-import { LLM_ID } from "../constants";
+import { WEAK_MODEL_ID } from "../constants";
 import { getOrPromptForAPIKey } from "./getOrPromptForAPIKey";
 import { logger } from "./logger";
 
@@ -13,7 +13,7 @@ const embeddingModelId = "text-embedding-3-small";
 let embeddingModel: OpenAIEmbedding | null = null;
 
 export async function initializeEmbeddingModel(): Promise<void> {
-  const apiKey = await getOrPromptForAPIKey(LLM_ID);
+  const apiKey = await getOrPromptForAPIKey(WEAK_MODEL_ID);
   if (!apiKey) {
     throw new Error("OpenAI API key not found. Please set up your API key.");
   }

src/utils/getOrPromptForAPIKey.ts

@@ -1,12 +1,18 @@
 import chalk from "chalk";
 import { getFreeformInput } from "../cli/interface";
-import { API_KEY_PREFIXES, API_KEY_PROMPTS, APIProvider } from "../constants";
+import { API_KEY_PREFIXES, API_KEY_PROMPTS, APIProvider, EMBEDDINGS_MODEL_ID } from "../constants";
 import { getAPIKey, storeAPIKey } from "./apiKeyManager";
 import { logger } from "./logger";
 
 const MAX_RETRIES = 5;
 
-export async function getOrPromptForAPIKey(modelId: string, forceNew: boolean = false): Promise<string> {
+interface GetAPIKeyOptions {
+  forceNew?: boolean;
+  prePromptText?: string;
+}
+
+export async function getOrPromptForAPIKey(modelId: string, options: GetAPIKeyOptions = {}): Promise<string> {
+  const { forceNew = false, prePromptText } = options;
   const provider = getProviderFromModel(modelId);
   let retries = 0;
 
@@ -32,7 +38,11 @@ export async function getOrPromptForAPIKey(modelId: string, forceNew: boolean =
       if (!apiKey) {
         try {
           logger.debug(`Prompting user for ${provider} API key`);
-          console.log(chalk.yellow(`${provider} API key not found or invalid.`));
+
+          const text = prePromptText
+            ? chalk.yellow(prePromptText)
+            : `${provider} API key not found or invalid.`;
+          console.log(chalk.yellow(text));
           apiKey = await getFreeformInput(API_KEY_PROMPTS[provider], true);
         } catch (inputError) {
           if (inputError.message.includes("non-interactive mode")) {
@@ -72,6 +82,17 @@ export async function getOrPromptForAPIKey(modelId: string, forceNew: boolean =
   process.exit(1);
 }
 
+export async function getApiKeyForModel(modelId: string): Promise<string> {
+  const apiKey = await getOrPromptForAPIKey(modelId);
+
+  if (!apiKey) {
+    logger.error(`No ${modelId} API key found`);
+    throw new Error(`${modelId} API key not found. Please run the application again to set it up.`);
+  }
+
+  return apiKey;
+}
+
 function isValidAPIKey(apiKey: string, provider: APIProvider): boolean {
   return apiKey.trim() !== "" && apiKey.startsWith(API_KEY_PREFIXES[provider]);
 }
@@ -82,18 +103,8 @@ export function getProviderFromModel(modelId: string): APIProvider {
   if (modelId.startsWith("gpt")) result = "OPENAI";
   if (modelId.startsWith("llama")) result = "GROQ";
   if (modelId.startsWith("claude")) result = "ANTHROPIC";
+  if (EMBEDDINGS_MODEL_ID === modelId) result = "OPENAI";
   if (!result) throw new Error(`Unsupported model ID: ${modelId}`);
   logger.debug(`Model ID ${modelId} matches a supported API provider, "${result}"`);
   return result;
 }
-
-export async function getApiKeyForModel(modelId: string): Promise<string> {
-  const apiKey = await getOrPromptForAPIKey(modelId);
-
-  if (!apiKey) {
-    logger.error(`No ${modelId} API key found`);
-    throw new Error(`${modelId} API key not found. Please run the application again to set it up.`);
-  }
-
-  return apiKey;
-}

src/utils/usageCost.ts

@@ -1,4 +1,4 @@
-import { MODEL_PRICES } from "../constants";
+import { MODELS } from "../constants";
 
 interface UsageCostResult {
   costUSD: number;
@@ -27,7 +27,7 @@ export function countUsageCost(
   },
   modelId: string,
 ): UsageCostResult {
-  const modelConfig = MODEL_PRICES[modelId];
+  const modelConfig = MODELS[modelId];
   if (!modelConfig) {
     throw new Error(`Unknown model ID: ${modelId}`);
   }