okengine 0.9.0 → 0.10.0

package/src/cli/ai-setup/apply.ts

@@ -1,9 +1,10 @@
 /**
- * Write AI driver config, env, and `src/ai.ts` for `oke ai setup`.
+ * Write AI driver config, env, and `src/core/ai.ts` for `oke ai setup`.
  */
 
 import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
 import { dirname, join } from "node:path";
+import { OLLAMA_IMAGE } from "../../docker/recipes/index.ts";
 
 /** Choices applied to the project. */
 export type AiSetupApplyInput = {
@@ -13,6 +14,10 @@ export type AiSetupApplyInput = {
   readonly visionModel?: string | null;
   readonly embedModel?: string | null;
   readonly apiKeyEnv?: string;
+  /** When set with {@link apiKeyEnv}, writes the token into `.env.local`. */
+  readonly apiKey?: string;
+  /** Optional `images.ai` pin (llama.cpp / Ollama / vLLM / SGLang). */
+  readonly image?: string;
 };
 
 /** Options for {@link applyAiSetup}. */
@@ -49,8 +54,10 @@ export function applyAiSetup(
   if (updateDrivers) {
     let config = readFileSync(configPath, "utf8");
     config = upsertAiDrivers(config, input.driver);
-    if (input.driver === "ollama") {
-      config = upsertImage(config, "ai", "ollama/ollama:latest");
+    if (input.image) {
+      config = upsertImage(config, "ai", input.image);
+    } else if (input.driver === "ollama") {
+      config = upsertImage(config, "ai", OLLAMA_IMAGE);
     }
     writeFileSync(configPath, config, "utf8");
   }
@@ -62,20 +69,29 @@ export function applyAiSetup(
     : existsSync(envExample)
       ? readFileSync(envExample, "utf8")
       : "";
-  env = upsertEnv(env, "OKE_AI_DRIVER", input.driver);
-  if (input.baseUrl) env = upsertEnv(env, "OKE_AI_URL", input.baseUrl);
+  // Driver / URL stay commented so docker mode can hydrate `OKE_AI_URL` from
+  // `docker/.env.docker` without being shadowed. The selected chat model is
+  // written active — `oke dev` seeds it into stack controls / `LLAMA_ARG_DOCKER_REPO`
+  // (commented-only left the compose default `smollm2` forever).
+  env = upsertEnv(env, "OKE_AI_DRIVER", input.driver, { comment: true });
+  if (input.baseUrl) env = upsertEnv(env, "OKE_AI_URL", input.baseUrl, { comment: true });
   if (input.chatModel) env = upsertEnv(env, "OKE_AI_MODEL", input.chatModel);
-  if (input.visionModel) env = upsertEnv(env, "OKE_AI_VISION_MODEL", input.visionModel);
-  if (input.embedModel) env = upsertEnv(env, "OKE_AI_EMBED_MODEL", input.embedModel);
+  if (input.visionModel) {
+    env = upsertEnv(env, "OKE_AI_VISION_MODEL", input.visionModel, { comment: true });
+  }
+  if (input.embedModel) {
+    env = upsertEnv(env, "OKE_AI_EMBED_MODEL", input.embedModel, { comment: true });
+  }
   if (input.apiKeyEnv) {
-    // Ensure the key name is present (value left blank for the user)
-    if (!new RegExp(`^${input.apiKeyEnv}=`, "m").test(env)) {
-      env = `${env.trimEnd()}\n${input.apiKeyEnv}=\n`;
+    if (input.apiKey !== undefined && input.apiKey.length > 0) {
+      env = upsertEnv(env, input.apiKeyEnv, input.apiKey);
+    } else if (!new RegExp(`^#?\\s*${input.apiKeyEnv}=`, "m").test(env)) {
+      env = `${env.trimEnd()}\n# ${input.apiKeyEnv}=\n`;
     }
   }
   writeFileSync(envPath, env.endsWith("\n") ? env : `${env}\n`, "utf8");
 
-  const aiTsPath = join(cwd, "src", "ai.ts");
+  const aiTsPath = join(cwd, "src", "core", "ai.ts");
   mkdirSync(dirname(aiTsPath), { recursive: true });
   writeFileSync(aiTsPath, renderAiTs(input), "utf8");
   ensureAiImported(cwd);
@@ -89,11 +105,11 @@ export function applyAiSetup(
  */
 export function upsertAiDrivers(source: string, driver: string): string {
   const block = `{
-    local: "${driver}",
-    docker: "${driver === "mock" ? "mock" : driver}",
-    test: "mock",
-    prod: "${driver === "mock" ? "mock" : driver}",
-  }`;
+      local: "${driver}",
+      docker: "${driver === "mock" ? "mock" : driver}",
+      test: "mock",
+      prod: "${driver === "mock" ? "mock" : driver}",
+    }`;
   // Top-level drivers.ai only (4-space indent) — never nest under channel.email.
   const aiRe = /^( {4})ai:\s*\{[\s\S]*?\n\1\}/m;
   if (aiRe.test(source)) {
@@ -131,13 +147,30 @@ export function upsertImage(source: string, key: string, image: string): string
   return source.replace(imagesRe, `images: {${body}\n${line}\n  }`);
 }
 
+/** Options for {@link upsertEnv}. */
+export type UpsertEnvOptions = {
+  /**
+   * When true, write `# KEY=value` (opt-in override). Used for AI stack keys
+   * so they do not shadow `docker/.env.docker`.
+   */
+  readonly comment?: boolean;
+};
+
 /**
+ * Insert or replace a dotenv assignment (optionally commented).
+ *
  * @param env - dotenv text
  * @param key - Variable name
  * @param value - Value
+ * @param options - Comment vs active line
  */
-export function upsertEnv(env: string, key: string, value: string): string {
-  const line = `${key}=${value}`;
+export function upsertEnv(
+  env: string,
+  key: string,
+  value: string,
+  options: UpsertEnvOptions = {},
+): string {
+  const line = options.comment ? `# ${key}=${value}` : `${key}=${value}`;
   const re = new RegExp(`^#?\\s*${key}=.*$`, "m");
   if (re.test(env)) return env.replace(re, line);
   const aiSection = /(# ── AI[^\n]*\n)/;
@@ -197,18 +230,27 @@ export function renderAiTs(input: AiSetupApplyInput): string {
 }
 
 /**
- * Ensure `src/app.ts` (or `src/core.ts`) imports `./ai`.
+ * Ensure `src/app.ts` or `src/core/index.ts` imports AI wiring.
  *
  * @param cwd - Project root
  */
 function ensureAiImported(cwd: string): void {
-  for (const rel of ["src/app.ts", "src/core.ts"]) {
+  const candidates: ReadonlyArray<{ readonly rel: string; readonly importLine: string }> = [
+    { rel: "src/app.ts", importLine: `import "./core/ai";\n` },
+    { rel: "src/core/index.ts", importLine: `import "./ai";\n` },
+    { rel: "src/core.ts", importLine: `import "./ai";\n` },
+  ];
+  for (const { rel, importLine } of candidates) {
     const path = join(cwd, rel);
     if (!existsSync(path)) continue;
     const src = readFileSync(path, "utf8");
-    if (/from\s+["']\.\/ai["']/.test(src) || /import\s+["']\.\/ai["']/.test(src)) return;
-    const next = `import "./ai";\n${src}`;
-    writeFileSync(path, next, "utf8");
+    if (
+      /from\s+["']\.\/(?:core\/)?ai["']/.test(src) ||
+      /import\s+["']\.\/(?:core\/)?ai["']/.test(src)
+    ) {
+      return;
+    }
+    writeFileSync(path, `${importLine}${src}`, "utf8");
     return;
   }
 }