Spaces:
Running
Running
| import { | |
| ANTHROPIC_DEFAULT_LARGE_MODEL, | |
| ANTHROPIC_DEFAULT_SMALL_MODEL, | |
| AZURE_DEFAULT_EMBEDDING_MODEL, | |
| AZURE_DEFAULT_LARGE_API_VERSION, | |
| AZURE_DEFAULT_LARGE_MODEL, | |
| AZURE_DEFAULT_SMALL_MODEL, | |
| BEDROCK_DEFAULT_EMBEDDING_MODEL, | |
| BEDROCK_DEFAULT_LARGE_MODEL, | |
| BEDROCK_DEFAULT_REGION, | |
| BEDROCK_DEFAULT_SMALL_MODEL, | |
| DEEPSEEK_DEFAULT_LARGE_MODEL, | |
| DEEPSEEK_DEFAULT_SMALL_MODEL, | |
| EMBEDDINGS_DEFAULTS, | |
| OPENAI_DEFAULT_EMBEDDING_MODEL, | |
| OPENAI_DEFAULT_LARGE_MODEL, | |
| OPENAI_DEFAULT_SMALL_MODEL, | |
| OLLAMA_EMBEDDING_DIMENSIONS, | |
| ENV_KEYS, | |
| GCP_DEFAULT_EMBEDDING_MODEL, | |
| GCP_DEFAULT_LARGE_MODEL, | |
| GCP_DEFAULT_SMALL_MODEL, | |
| OLLAMA_DEFAULT_HOST, | |
| OLLAMA_DEFAULT_LARGE_MODEL, | |
| OLLAMA_DEFAULT_PORT, | |
| OLLAMA_DEFAULT_SMALL_MODEL, | |
| SERVICE_DEFAULTS, | |
| } from "../constants.js"; | |
| import { applyEnvValues, loadOrSeedEnv, removeEnvValue, saveEnv } from "./env.js"; | |
| import type { InitChoices } from "../types.js"; | |
| import { resolveActiveEmbeddingModel } from "../flows/embedding-dimensions.js"; | |
| import { lookupEmbeddingDimensions } from "../constants.js"; | |
| const AZURE_ENDPOINT_PLACEHOLDER = "https://example.openai.azure.com"; | |
| const AZURE_EMBEDDING_PLACEHOLDER = | |
| "https://example.openai.azure.com/openai/deployments/text-embedding-3-large/embeddings?api-version=2023-05-15"; | |
| const AZURE_KEY_PLACEHOLDER = "azure-key-placeholder"; | |
| const OPENAI_KEY_PLACEHOLDER = "openai-api-key-placeholder"; | |
| const ANTHROPIC_KEY_PLACEHOLDER = "anthropic-api-key-placeholder"; | |
| const DEEPSEEK_KEY_PLACEHOLDER = "deepseek-api-key-placeholder"; | |
| function resolveModelsMode(choices: InitChoices): "local" | "remote" { | |
| const { llmSmallProvider, llmLargeProvider, embeddingsProvider } = choices.models; | |
| if ( | |
| llmSmallProvider === "ollama" && | |
| llmLargeProvider === "ollama" && | |
| embeddingsProvider === "ollama" | |
| ) { | |
| return "local"; | |
| } | |
| return choices.models.mode; | |
| } | |
| export async function writeEnvFromChoices(choices: InitChoices): Promise<void> { | |
| const env = await loadOrSeedEnv(); | |
| const values: Record<string, string | number | undefined> = {}; | |
| values[ENV_KEYS.graphDb] = choices.dbs.graphDb; | |
| values[ENV_KEYS.dataDb] = choices.dbs.dataDb; | |
| values[ENV_KEYS.vectorDb] = choices.dbs.vectorDb; | |
| values[ENV_KEYS.modelsMode] = resolveModelsMode(choices); | |
| values[ENV_KEYS.ocrMode] = choices.pipeline.ocrMode; | |
| values[ENV_KEYS.pipelineMode] = "accurate"; | |
| values[ENV_KEYS.agenticArchitecture] = "custom"; | |
| values[ENV_KEYS.runGraphConsolidator] = "true"; | |
| values[ENV_KEYS.celeryWorkerConcurrency] = 4; | |
| applySearchValues(values, choices); | |
| values[ENV_KEYS.brainpatToken] = choices.auth.brainpatToken; | |
| values[ENV_KEYS.enabledPlugins] = choices.plugins.map((plugin) => plugin.name).join(","); | |
| values["ENV"] = "development"; | |
| values["CONSOLE_ENABLED"] = "true"; | |
| values[ENV_KEYS.redisHost] = choices.connections.redis.host; | |
| values[ENV_KEYS.redisPort] = choices.connections.redis.port; | |
| applyPostgresValues(values, choices); | |
| applyNeo4jValues(values, choices); | |
| applyMilvusValues(values, choices, env); | |
| applyMongoValues(values, choices); | |
| applyOllamaValues(values, choices); | |
| applyRemoteModelValues(values, choices); | |
| applyEmbeddingDefaults(values, choices); | |
| applyEnvValues(env, values); | |
| await saveEnv(env); | |
| } | |
| function applySearchValues( | |
| values: Record<string, string | number | undefined>, | |
| choices: InitChoices, | |
| ): void { | |
| const search = choices.search ?? { enabled: false, retrieval: "hybrid" }; | |
| const enabled = search.enabled; | |
| const retrieval = search.retrieval ?? "hybrid"; | |
| values[ENV_KEYS.searchEnabled] = enabled ? "true" : "false"; | |
| values[ENV_KEYS.searchUseDense] = | |
| !enabled || retrieval === "hybrid" || retrieval === "dense" ? "true" : "false"; | |
| values[ENV_KEYS.searchUseBm25] = | |
| !enabled || retrieval === "hybrid" || retrieval === "bm25" ? "true" : "false"; | |
| values[ENV_KEYS.searchFusion] = "rrf"; | |
| values[ENV_KEYS.searchFusionAlpha] = "0.5"; | |
| values[ENV_KEYS.searchBm25K1] = "1.2"; | |
| values[ENV_KEYS.searchBm25B] = "0.75"; | |
| values[ENV_KEYS.contextPassageMode] = enabled ? "hybrid" : "hybrid"; | |
| } | |
| function applyPostgresValues( | |
| values: Record<string, string | number | undefined>, | |
| choices: InitChoices, | |
| ): void { | |
| const pg = choices.connections.postgresql; | |
| const defaults = SERVICE_DEFAULTS.postgresql; | |
| values[ENV_KEYS.postgresHost] = pg?.host ?? defaults.host; | |
| values[ENV_KEYS.postgresPort] = pg?.port ?? defaults.port; | |
| values[ENV_KEYS.postgresUsername] = pg?.username ?? defaults.username; | |
| values[ENV_KEYS.postgresPassword] = pg?.password ?? defaults.password; | |
| values[ENV_KEYS.postgresSystemDatabase] = | |
| pg?.systemDatabase ?? defaults.systemDatabase; | |
| values[ENV_KEYS.postgresMaintenanceDatabase] = | |
| pg?.maintenanceDatabase ?? defaults.maintenanceDatabase; | |
| } | |
| function applyNeo4jValues( | |
| values: Record<string, string | number | undefined>, | |
| choices: InitChoices, | |
| ): void { | |
| const neo = choices.connections.neo4j; | |
| const defaults = SERVICE_DEFAULTS.neo4j; | |
| values[ENV_KEYS.neo4jHost] = neo?.host ?? defaults.host; | |
| values[ENV_KEYS.neo4jPort] = neo?.port ?? defaults.port; | |
| values[ENV_KEYS.neo4jUsername] = neo?.username ?? defaults.username; | |
| values[ENV_KEYS.neo4jPassword] = neo?.password ?? defaults.password; | |
| } | |
| function applyMilvusValues( | |
| values: Record<string, string | number | undefined>, | |
| choices: InitChoices, | |
| env: Awaited<ReturnType<typeof loadOrSeedEnv>>, | |
| ): void { | |
| const mv = choices.connections.milvus; | |
| const defaults = SERVICE_DEFAULTS.milvus; | |
| values[ENV_KEYS.milvusHost] = mv?.host ?? defaults.host; | |
| values[ENV_KEYS.milvusPort] = mv?.port ?? defaults.port; | |
| if (mv?.uri) { | |
| values[ENV_KEYS.milvusUri] = mv.uri; | |
| } else { | |
| removeEnvValue(env, ENV_KEYS.milvusUri); | |
| } | |
| if (mv?.token) { | |
| values[ENV_KEYS.milvusToken] = mv.token; | |
| } else { | |
| removeEnvValue(env, ENV_KEYS.milvusToken); | |
| } | |
| } | |
| function applyMongoValues( | |
| values: Record<string, string | number | undefined>, | |
| choices: InitChoices, | |
| ): void { | |
| const mg = choices.connections.mongo; | |
| const defaults = SERVICE_DEFAULTS.mongo; | |
| values[ENV_KEYS.mongoHost] = mg?.host ?? defaults.host; | |
| values[ENV_KEYS.mongoPort] = mg?.port ?? defaults.port; | |
| values[ENV_KEYS.mongoUsername] = mg?.username ?? defaults.username; | |
| values[ENV_KEYS.mongoPassword] = mg?.password ?? defaults.password; | |
| values[ENV_KEYS.mongoSystemDatabase] = "system"; | |
| } | |
| function applyOllamaValues( | |
| values: Record<string, string | number | undefined>, | |
| choices: InitChoices, | |
| ): void { | |
| const usesOllama = | |
| choices.models.llmSmallProvider === "ollama" || | |
| choices.models.llmLargeProvider === "ollama" || | |
| choices.models.embeddingsProvider === "ollama"; | |
| if (usesOllama && choices.models.ollama) { | |
| const ol = choices.models.ollama; | |
| values[ENV_KEYS.ollamaHost] = ol.host; | |
| values[ENV_KEYS.ollamaPort] = ol.port; | |
| values[ENV_KEYS.ollamaSmallModel] = ol.smallModel; | |
| values[ENV_KEYS.ollamaLargeModel] = ol.largeModel; | |
| values[ENV_KEYS.embeddingsLocalModel] = ol.embeddingsLocalModel; | |
| return; | |
| } | |
| values[ENV_KEYS.ollamaHost] = OLLAMA_DEFAULT_HOST; | |
| values[ENV_KEYS.ollamaPort] = OLLAMA_DEFAULT_PORT; | |
| values[ENV_KEYS.ollamaSmallModel] = OLLAMA_DEFAULT_SMALL_MODEL; | |
| values[ENV_KEYS.ollamaLargeModel] = OLLAMA_DEFAULT_LARGE_MODEL; | |
| } | |
| function applyRemoteModelValues( | |
| values: Record<string, string | number | undefined>, | |
| choices: InitChoices, | |
| ): void { | |
| const gcp = choices.models.gcp; | |
| const azure = choices.models.azure; | |
| const openai = choices.models.openai; | |
| const anthropic = choices.models.anthropic; | |
| const deepseek = choices.models.deepseek; | |
| const bedrock = choices.models.bedrock; | |
| values[ENV_KEYS.llmSmallProvider] = choices.models.llmSmallProvider; | |
| values[ENV_KEYS.llmLargeProvider] = choices.models.llmLargeProvider; | |
| values[ENV_KEYS.embeddingsProvider] = choices.models.embeddingsProvider; | |
| values[ENV_KEYS.gcpCredentialsPath] = gcp?.credentialsPath ?? "gcp_credentials.json"; | |
| values[ENV_KEYS.gcpProjectId] = gcp?.projectId ?? "your-project-id"; | |
| values[ENV_KEYS.gcpSmallLlmModel] = gcp?.smallLlmModel ?? GCP_DEFAULT_SMALL_MODEL; | |
| values[ENV_KEYS.gcpExtraSmallLlmModel] = | |
| gcp?.smallLlmModel ?? GCP_DEFAULT_SMALL_MODEL; | |
| values[ENV_KEYS.gcpLargeLlmModel] = gcp?.largeLlmModel ?? GCP_DEFAULT_LARGE_MODEL; | |
| values[ENV_KEYS.gcpEmbeddingModel] = gcp?.embeddingModel ?? GCP_DEFAULT_EMBEDDING_MODEL; | |
| values[ENV_KEYS.azureSmallLlmModel] = | |
| azure?.smallLlmModel ?? AZURE_DEFAULT_SMALL_MODEL; | |
| values[ENV_KEYS.azureLargeLlmModel] = | |
| azure?.largeLlmModel ?? AZURE_DEFAULT_LARGE_MODEL; | |
| values[ENV_KEYS.azureLlmApiVersion] = | |
| azure?.llmApiVersion ?? AZURE_DEFAULT_LARGE_API_VERSION; | |
| values[ENV_KEYS.azureLlmEndpoint] = | |
| azure?.llmEndpoint ?? AZURE_ENDPOINT_PLACEHOLDER; | |
| values[ENV_KEYS.azureLlmSubscriptionKey] = | |
| azure?.llmSubscriptionKey ?? AZURE_KEY_PLACEHOLDER; | |
| values[ENV_KEYS.azureLargeLlmApiVersion] = | |
| azure?.llmApiVersion ?? AZURE_DEFAULT_LARGE_API_VERSION; | |
| values[ENV_KEYS.azureLargeLlmEndpoint] = | |
| azure?.llmEndpoint ?? AZURE_ENDPOINT_PLACEHOLDER; | |
| values[ENV_KEYS.azureLargeLlmSubscriptionKey] = | |
| azure?.llmSubscriptionKey ?? AZURE_KEY_PLACEHOLDER; | |
| values[ENV_KEYS.azureEmbeddingModel] = | |
| azure?.embeddingModel ?? AZURE_DEFAULT_EMBEDDING_MODEL; | |
| values[ENV_KEYS.azureEmbeddingFullEndpoint] = | |
| azure?.embeddingEndpoint ?? AZURE_EMBEDDING_PLACEHOLDER; | |
| values[ENV_KEYS.azureEmbeddingKey] = | |
| azure?.embeddingKey ?? AZURE_KEY_PLACEHOLDER; | |
| values[ENV_KEYS.openaiApiKey] = openai?.apiKey ?? OPENAI_KEY_PLACEHOLDER; | |
| values[ENV_KEYS.openaiSmallLlmModel] = | |
| openai?.smallLlmModel ?? OPENAI_DEFAULT_SMALL_MODEL; | |
| values[ENV_KEYS.openaiLargeLlmModel] = | |
| openai?.largeLlmModel ?? OPENAI_DEFAULT_LARGE_MODEL; | |
| values[ENV_KEYS.openaiEmbeddingModel] = | |
| openai?.embeddingModel ?? OPENAI_DEFAULT_EMBEDDING_MODEL; | |
| if (openai?.baseUrl) { | |
| values[ENV_KEYS.openaiBaseUrl] = openai.baseUrl; | |
| } else { | |
| values[ENV_KEYS.openaiBaseUrl] = ""; | |
| } | |
| values[ENV_KEYS.anthropicApiKey] = anthropic?.apiKey ?? ANTHROPIC_KEY_PLACEHOLDER; | |
| values[ENV_KEYS.anthropicSmallLlmModel] = | |
| anthropic?.smallLlmModel ?? ANTHROPIC_DEFAULT_SMALL_MODEL; | |
| values[ENV_KEYS.anthropicLargeLlmModel] = | |
| anthropic?.largeLlmModel ?? ANTHROPIC_DEFAULT_LARGE_MODEL; | |
| values[ENV_KEYS.deepseekApiKey] = deepseek?.apiKey ?? DEEPSEEK_KEY_PLACEHOLDER; | |
| values[ENV_KEYS.deepseekSmallLlmModel] = | |
| deepseek?.smallLlmModel ?? DEEPSEEK_DEFAULT_SMALL_MODEL; | |
| values[ENV_KEYS.deepseekLargeLlmModel] = | |
| deepseek?.largeLlmModel ?? DEEPSEEK_DEFAULT_LARGE_MODEL; | |
| values[ENV_KEYS.bedrockAccessKeyId] = bedrock?.accessKeyId ?? "your-access-key-id"; | |
| values[ENV_KEYS.bedrockSecretAccessKey] = | |
| bedrock?.secretAccessKey ?? "your-secret-access-key"; | |
| values[ENV_KEYS.bedrockSessionToken] = bedrock?.sessionToken ?? ""; | |
| values[ENV_KEYS.bedrockSmallLlmModel] = | |
| bedrock?.smallLlmModel ?? BEDROCK_DEFAULT_SMALL_MODEL; | |
| values[ENV_KEYS.bedrockLargeLlmModel] = | |
| bedrock?.largeLlmModel ?? BEDROCK_DEFAULT_LARGE_MODEL; | |
| values[ENV_KEYS.bedrockEmbeddingModel] = | |
| bedrock?.embeddingModel ?? BEDROCK_DEFAULT_EMBEDDING_MODEL; | |
| } | |
| function fallbackEmbeddingDimensions(choices: InitChoices) { | |
| const embeddingModel = resolveActiveEmbeddingModel(choices.models); | |
| const known = lookupEmbeddingDimensions(embeddingModel); | |
| if (known !== undefined) { | |
| return { | |
| nodesDimension: known, | |
| tripletsDimension: known, | |
| observationsDimension: known, | |
| dataDimension: known, | |
| relationshipsDimension: known, | |
| }; | |
| } | |
| if (choices.models.embeddingsProvider === "ollama") { | |
| return OLLAMA_EMBEDDING_DIMENSIONS; | |
| } | |
| return EMBEDDINGS_DEFAULTS; | |
| } | |
| function applyEmbeddingDefaults( | |
| values: Record<string, string | number | undefined>, | |
| choices: InitChoices, | |
| ): void { | |
| values[ENV_KEYS.embeddingsSmallModel] = EMBEDDINGS_DEFAULTS.smallModel; | |
| values[ENV_KEYS.embeddingsLocalModel] = | |
| values[ENV_KEYS.embeddingsLocalModel] ?? EMBEDDINGS_DEFAULTS.localModel; | |
| const embeddingDims = choices.models.embeddingDimensions ?? fallbackEmbeddingDimensions(choices); | |
| values[ENV_KEYS.embeddingNodesDimension] = embeddingDims.nodesDimension; | |
| values[ENV_KEYS.embeddingTripletsDimension] = embeddingDims.tripletsDimension; | |
| values[ENV_KEYS.embeddingObservationsDimension] = | |
| embeddingDims.observationsDimension; | |
| values[ENV_KEYS.embeddingDataDimension] = embeddingDims.dataDimension; | |
| values[ENV_KEYS.embeddingRelationshipsDimension] = | |
| embeddingDims.relationshipsDimension; | |
| if (choices.models.embeddingsProvider === "azure" && choices.models.azure) { | |
| values[ENV_KEYS.embeddingsFullEndpoint] = | |
| choices.models.azure.embeddingEndpoint; | |
| values[ENV_KEYS.embeddingsKey] = choices.models.azure.embeddingKey; | |
| } else { | |
| values[ENV_KEYS.embeddingsFullEndpoint] = | |
| values[ENV_KEYS.embeddingsFullEndpoint] ?? AZURE_EMBEDDING_PLACEHOLDER; | |
| values[ENV_KEYS.embeddingsKey] = | |
| values[ENV_KEYS.embeddingsKey] ?? AZURE_KEY_PLACEHOLDER; | |
| } | |
| } | |