feat(compiler): support custom prompts (#1002)

Author: mathio

.changeset/smooth-houses-sit.md

@@ -0,0 +1,6 @@
+---
+"@lingo.dev/_compiler": patch
+"lingo.dev": patch
+---
+
+support custom prompts in compiler

packages/compiler/src/_base.ts

@@ -79,6 +79,15 @@ export type CompilerParams = {
    * @default {}
    */
   models: "lingo.dev" | Record<string, string>;
+  /**
+   * Custom system prompt for the translation engine. If set, this prompt will override the default system prompt defined in Compiler.
+   * Only works with custom models, not with Lingo.dev Engine.
+   *
+   * Example: "You are a helpful assistant that translates {SOURCE_LOCALE} to {TARGET_LOCALE}."
+   *
+   * @default null
+   */
+  prompt?: string | null;
 };
 export type CompilerInput = {
   relativeFilePath: string;
@@ -146,4 +155,5 @@ export const defaultParams: CompilerParams = {
   useDirective: false,
   debug: false,
   models: {},
+  prompt: null,
 };

packages/compiler/src/lib/lcp/api.spec.ts

@@ -25,9 +25,27 @@ describe("LCPAPI", () => {
 
       expect(chunkSpy).toHaveBeenCalledWith(0);
       expect(translateSpy).toHaveBeenCalledTimes(3);
-      expect(translateSpy).toHaveBeenCalledWith(modelsMock, 1, "en", "es");
-      expect(translateSpy).toHaveBeenCalledWith(modelsMock, 2, "en", "es");
-      expect(translateSpy).toHaveBeenCalledWith(modelsMock, 3, "en", "es");
+      expect(translateSpy).toHaveBeenCalledWith(
+        modelsMock,
+        1,
+        "en",
+        "es",
+        undefined,
+      );
+      expect(translateSpy).toHaveBeenCalledWith(
+        modelsMock,
+        2,
+        "en",
+        "es",
+        undefined,
+      );
+      expect(translateSpy).toHaveBeenCalledWith(
+        modelsMock,
+        3,
+        "en",
+        "es",
+        undefined,
+      );
       expect(mergeSpy).toHaveBeenCalledWith([10, 20, 30]);
       expect(result).toEqual(100);
     });

packages/compiler/src/lib/lcp/api/index.ts

@@ -34,6 +34,7 @@ export class LCPAPI {
     sourceDictionary: DictionarySchema,
     sourceLocale: string,
     targetLocale: string,
+    prompt?: string | null,
   ): Promise<DictionarySchema> {
     const timeLabel = `LCPAPI.translate: ${targetLocale}`;
     console.time(timeLabel);
@@ -45,6 +46,7 @@ export class LCPAPI {
         chunk,
         sourceLocale,
         targetLocale,
+        prompt,
       );
       translatedChunks.push(translatedChunk);
     }
@@ -161,6 +163,7 @@ export class LCPAPI {
     sourceDictionary: DictionarySchema,
     sourceLocale: string,
     targetLocale: string,
+    prompt?: string | null,
   ): Promise<DictionarySchema> {
     if (models === "lingo.dev") {
       try {
@@ -231,7 +234,11 @@ export class LCPAPI {
           messages: [
             {
               role: "system",
-              content: getSystemPrompt({ sourceLocale, targetLocale }),
+              content: getSystemPrompt({
+                sourceLocale,
+                targetLocale,
+                prompt: prompt ?? undefined,
+              }),
             },
             ...shots.flatMap((shotsTuple) => [
               {

packages/compiler/src/lib/lcp/api/prompt.spec.ts

@@ -0,0 +1,57 @@
+import prompt from "./prompt";
+import { describe, it, expect, vi } from "vitest";
+
+const baseArgs = {
+  sourceLocale: "en",
+  targetLocale: "es",
+};
+
+describe("prompt", () => {
+  it("returns user-defined prompt with replacements", () => {
+    const args = {
+      ...baseArgs,
+      prompt: "Translate from {SOURCE_LOCALE} to {TARGET_LOCALE}.",
+    };
+    const result = prompt(args);
+    expect(result).toBe("Translate from en to es.");
+  });
+
+  it("trims and replaces variables in user prompt", () => {
+    const args = {
+      ...baseArgs,
+      prompt: "  {SOURCE_LOCALE} => {TARGET_LOCALE}  ",
+    };
+    const result = prompt(args);
+    expect(result).toBe("en => es");
+  });
+
+  it("falls back to built-in prompt if no user prompt", () => {
+    const args = { ...baseArgs };
+    const result = prompt(args);
+    expect(result).toContain("You are an advanced AI localization engine");
+    expect(result).toContain("Source language (locale code): en");
+    expect(result).toContain("Target language (locale code): es");
+  });
+
+  it("logs when using user-defined prompt", () => {
+    const spy = vi.spyOn(console, "log");
+    const args = {
+      ...baseArgs,
+      prompt: "Prompt {SOURCE_LOCALE} {TARGET_LOCALE}",
+    };
+    prompt(args);
+    expect(spy).toHaveBeenCalledWith(
+      "✨ Compiler is using user-defined prompt.",
+    );
+    spy.mockRestore();
+  });
+
+  it("returns built-in prompt if user prompt is empty or whitespace", () => {
+    const args = {
+      ...baseArgs,
+      prompt: "   ",
+    };
+    const result = prompt(args);
+    expect(result).toContain("You are an advanced AI localization engine");
+  });
+});

packages/compiler/src/lib/lcp/server.spec.ts

@@ -405,6 +405,7 @@ describe("LCPServer", () => {
         },
         "en",
         "fr",
+        undefined,
       );
 
       // Verify final result combines cached and newly translated content

packages/compiler/src/lib/lcp/server.ts

@@ -8,22 +8,21 @@ import _ from "lodash";
 import { LCPCache } from "./cache";
 import { LCPAPI } from "./api";
 
-export type LCPServerParams = {
+type LCPServerBaseParams = {
   lcp: LCPSchema;
   sourceLocale: string;
-  targetLocales: string[];
   sourceRoot: string;
   lingoDir: string;
   models: "lingo.dev" | Record<string, string>;
+  prompt?: string | null;
 };
 
-export type LCPServerParamsForLocale = {
-  lcp: LCPSchema;
-  sourceLocale: string;
+export type LCPServerParams = LCPServerBaseParams & {
+  targetLocales: string[];
+};
+
+export type LCPServerParamsForLocale = LCPServerBaseParams & {
   targetLocale: string;
-  sourceRoot: string;
-  lingoDir: string;
-  models: "lingo.dev" | Record<string, string>;
 };
 
 export class LCPServer {
@@ -119,6 +118,7 @@ export class LCPServer {
         uncachedSourceDictionary,
         params.sourceLocale,
         params.targetLocale,
+        params.prompt,
       );
 
       // we merge new translations with cache, so that we can cache empty strings