Support turbo edits (pro) (#166)

2025-05-14 23:35:50 -07:00
parent d545babb63
commit 35b459d82d
12 changed files with 400 additions and 26 deletions
--- a/src/components/ProModeSelector.tsx
+++ b/src/components/ProModeSelector.tsx
@@ -20,6 +20,13 @@ export function ProModeSelector() {
  const toggleSaverMode = () => {
    updateSettings({ enableProSaverMode: !settings?.enableProSaverMode });
  };
+
+  const toggleLazyEdits = () => {
+    updateSettings({
+      enableProLazyEditsMode: !settings?.enableProLazyEditsMode,
+    });
+  };
+
  if (!settings?.enableDyadPro) {
    return null;
  }
@@ -75,6 +82,29 @@ export function ProModeSelector() {
              onCheckedChange={toggleSaverMode}
            />
          </div>
+          <div className="flex items-center justify-between">
+            <div className="space-y-2">
+              <Label htmlFor="lazy-edits">Turbo Edits</Label>
+              <div className="flex items-center gap-1">
+                <Tooltip>
+                  <TooltipTrigger asChild>
+                    <Info className="h-4 w-4 text-muted-foreground cursor-help" />
+                  </TooltipTrigger>
+                  <TooltipContent side="right" className="max-w-72">
+                    Edits files faster.
+                  </TooltipContent>
+                </Tooltip>
+                <p className="text-xs text-muted-foreground max-w-55">
+                  Makes editing files faster and cheaper.
+                </p>
+              </div>
+            </div>
+            <Switch
+              id="lazy-edits"
+              checked={Boolean(settings?.enableProLazyEditsMode)}
+              onCheckedChange={toggleLazyEdits}
+            />
+          </div>
        </div>
      </PopoverContent>
    </Popover>
--- a/src/components/chat/DyadEdit.tsx
+++ b/src/components/chat/DyadEdit.tsx
@@ -0,0 +1,110 @@
+import type React from "react";
+import type { ReactNode } from "react";
+import { useState } from "react";
+import {
+  ChevronsDownUp,
+  ChevronsUpDown,
+  Loader,
+  CircleX,
+  Rabbit,
+} from "lucide-react";
+import { CodeHighlight } from "./CodeHighlight";
+import { CustomTagState } from "./stateTypes";
+
+interface DyadEditProps {
+  children?: ReactNode;
+  node?: any;
+  path?: string;
+  description?: string;
+}
+
+export const DyadEdit: React.FC<DyadEditProps> = ({
+  children,
+  node,
+  path: pathProp,
+  description: descriptionProp,
+}) => {
+  const [isContentVisible, setIsContentVisible] = useState(false);
+
+  // Use props directly if provided, otherwise extract from node
+  const path = pathProp || node?.properties?.path || "";
+  const description = descriptionProp || node?.properties?.description || "";
+  const state = node?.properties?.state as CustomTagState;
+  const inProgress = state === "pending";
+  const aborted = state === "aborted";
+
+  // Extract filename from path
+  const fileName = path ? path.split("/").pop() : "";
+
+  return (
+    <div
+      className={`bg-(--background-lightest) hover:bg-(--background-lighter) rounded-lg px-4 py-2 border my-2 cursor-pointer ${
+        inProgress
+          ? "border-amber-500"
+          : aborted
+            ? "border-red-500"
+            : "border-border"
+      }`}
+      onClick={() => setIsContentVisible(!isContentVisible)}
+    >
+      <div className="flex items-center justify-between">
+        <div className="flex items-center gap-2">
+          <div className="flex items-center">
+            <Rabbit size={16} />
+            <span className="bg-blue-500 text-white text-xs px-1.5 py-0.5 rounded ml-1 font-medium">
+              Turbo Edit
+            </span>
+          </div>
+          {fileName && (
+            <span className="text-gray-700 dark:text-gray-300 font-medium text-sm">
+              {fileName}
+            </span>
+          )}
+          {inProgress && (
+            <div className="flex items-center text-amber-600 text-xs">
+              <Loader size={14} className="mr-1 animate-spin" />
+              <span>Editing...</span>
+            </div>
+          )}
+          {aborted && (
+            <div className="flex items-center text-red-600 text-xs">
+              <CircleX size={14} className="mr-1" />
+              <span>Did not finish</span>
+            </div>
+          )}
+        </div>
+        <div className="flex items-center">
+          {isContentVisible ? (
+            <ChevronsDownUp
+              size={20}
+              className="text-gray-500 hover:text-gray-700 dark:text-gray-400 dark:hover:text-gray-200"
+            />
+          ) : (
+            <ChevronsUpDown
+              size={20}
+              className="text-gray-500 hover:text-gray-700 dark:text-gray-400 dark:hover:text-gray-200"
+            />
+          )}
+        </div>
+      </div>
+      {path && (
+        <div className="text-xs text-gray-500 dark:text-gray-400 font-medium mb-1">
+          {path}
+        </div>
+      )}
+      {description && (
+        <div className="text-sm text-gray-600 dark:text-gray-300">
+          <span className="font-medium">Summary: </span>
+          {description}
+        </div>
+      )}
+      {isContentVisible && (
+        <div className="text-xs">
+          <CodeHighlight className="language-typescript">
+            {children}
+          </CodeHighlight>
+        </div>
+      )}
+    </div>
+  );
+};
--- a/src/components/chat/DyadMarkdownParser.tsx
+++ b/src/components/chat/DyadMarkdownParser.tsx
@@ -7,6 +7,7 @@ import { DyadDelete } from "./DyadDelete";
 import { DyadAddDependency } from "./DyadAddDependency";
 import { DyadExecuteSql } from "./DyadExecuteSql";
 import { DyadAddIntegration } from "./DyadAddIntegration";
+import { DyadEdit } from "./DyadEdit";
 import { CodeHighlight } from "./CodeHighlight";
 import { useAtomValue } from "jotai";
 import { isStreamingAtom } from "@/atoms/chatAtoms";
@@ -115,6 +116,7 @@ function preprocessUnclosedTags(content: string): {
    "dyad-add-integration",
    "dyad-output",
    "dyad-chat-summary",
+    "dyad-edit",
  ];

  let processedContent = content;
@@ -177,6 +179,7 @@ function parseCustomTags(content: string): ContentPiece[] {
    "dyad-add-integration",
    "dyad-output",
    "dyad-chat-summary",
+    "dyad-edit",
  ];

  const tagPattern = new RegExp(
@@ -344,6 +347,21 @@ function renderCustomTag(
        </DyadAddIntegration>
      );

+    case "dyad-edit":
+      return (
+        <DyadEdit
+          node={{
+            properties: {
+              path: attributes.path || "",
+              description: attributes.description || "",
+              state: getState({ isStreaming, inProgress }),
+            },
+          }}
+        >
+          {content}
+        </DyadEdit>
+      );
+
    case "dyad-output":
      return (
        <DyadOutput
--- a/src/ipc/handlers/chat_stream_handlers.ts
+++ b/src/ipc/handlers/chat_stream_handlers.ts
@@ -215,17 +215,16 @@ export function registerChatStreamHandlers() {
      } else {
        // Normal AI processing for non-test prompts
        const settings = readSettings();
-        const { modelClient, backupModelClients } = await getModelClient(
-          settings.selectedModel,
-          settings,
-        );

        // Extract codebase information if app is associated with the chat
        let codebaseInfo = "";
+        let files: { path: string; content: string }[] = [];
        if (updatedChat.app) {
          const appPath = getDyadAppPath(updatedChat.app.path);
          try {
-            codebaseInfo = await extractCodebase(appPath);
+            const out = await extractCodebase(appPath);
+            codebaseInfo = out.formattedOutput;
+            files = out.files;
            logger.log(`Extracted codebase information from ${appPath}`);
          } catch (error) {
            logger.error("Error extracting codebase:", error);
@@ -237,6 +236,11 @@ export function registerChatStreamHandlers() {
          "estimated tokens",
          codebaseInfo.length / 4,
        );
+        const { modelClient, backupModelClients } = await getModelClient(
+          settings.selectedModel,
+          settings,
+          files,
+        );

        // Prepare message history for the AI
        const messageHistory = updatedChat.messages.map((message) => ({
--- a/src/ipc/handlers/debug_handlers.ts
+++ b/src/ipc/handlers/debug_handlers.ts
@@ -140,7 +140,7 @@ export function registerDebugHandlers() {

        // Extract codebase
        const appPath = getDyadAppPath(app.path);
-        const codebase = await extractCodebase(appPath);
+        const codebase = (await extractCodebase(appPath)).formattedOutput;

        return {
          debugInfo,
--- a/src/ipc/handlers/proposal_handlers.ts
+++ b/src/ipc/handlers/proposal_handlers.ts
@@ -92,7 +92,8 @@ async function getCodebaseTokenCount(

  // Calculate and cache the token count
  logger.log(`Calculating codebase token count for chatId: ${chatId}`);
-  const codebase = await extractCodebase(getDyadAppPath(appPath));
+  const codebase = (await extractCodebase(getDyadAppPath(appPath)))
+    .formattedOutput;
  const tokenCount = estimateTokens(codebase);

  // Store in cache
--- a/src/ipc/handlers/token_count_handlers.ts
+++ b/src/ipc/handlers/token_count_handlers.ts
@@ -68,7 +68,7 @@ export function registerTokenCountHandlers() {

      if (chat.app) {
        const appPath = getDyadAppPath(chat.app.path);
-        codebaseInfo = await extractCodebase(appPath);
+        codebaseInfo = (await extractCodebase(appPath)).formattedOutput;
        codebaseTokens = estimateTokens(codebaseInfo);
        logger.log(
          `Extracted codebase information from ${appPath}, tokens: ${codebaseTokens}`,
--- a/src/ipc/utils/get_model_client.ts
+++ b/src/ipc/utils/get_model_client.ts
@@ -11,6 +11,9 @@ import log from "electron-log";
 import { getLanguageModelProviders } from "../shared/language_model_helpers";
 import { LanguageModelProvider } from "../ipc_types";
 import { llmErrorStore } from "@/main/llm_error_store";
+import { createDyadEngine } from "./llm_engine_provider";
+
+const dyadLocalEngine = process.env.DYAD_LOCAL_ENGINE;

 const AUTO_MODELS = [
  {
@@ -32,10 +35,16 @@ export interface ModelClient {
  builtinProviderId?: string;
 }

+interface File {
+  path: string;
+  content: string;
+}
+
 const logger = log.scope("getModelClient");
 export async function getModelClient(
  model: LargeLanguageModel,
  settings: UserSettings,
+  files?: File[],
 ): Promise<{
  modelClient: ModelClient;
  backupModelClients: ModelClient[];
@@ -65,8 +74,9 @@ export async function getModelClient(
          {
            provider: autoModel.provider,
            name: autoModel.name,
-          } as LargeLanguageModel,
+          },
          settings,
+          files,
        );
      }
    }
@@ -85,17 +95,33 @@ export async function getModelClient(

  // Handle Dyad Pro override
  if (dyadApiKey && settings.enableDyadPro) {
-    // Check if the selected provider supports Dyad Pro (has a gateway prefix)
-    if (providerConfig.gatewayPrefix) {
-      const provider = createOpenAI({
-        apiKey: dyadApiKey,
-        baseURL: "https://llm-gateway.dyad.sh/v1",
-      });
-      logger.info("Using Dyad Pro API key via Gateway");
+    // Check if the selected provider supports Dyad Pro (has a gateway prefix) OR
+    // we're using local engine.
+    if (providerConfig.gatewayPrefix || dyadLocalEngine) {
+      const provider = settings.enableProLazyEditsMode
+        ? createDyadEngine({
+            apiKey: dyadApiKey,
+            baseURL: dyadLocalEngine ?? "https://engine.dyad.sh/v1",
+          })
+        : createOpenAI({
+            apiKey: dyadApiKey,
+            baseURL: "https://llm-gateway.dyad.sh/v1",
+          });
+
+      logger.info(
+        `Using Dyad Pro API key. engine_enabled=${settings.enableProLazyEditsMode}`,
+      );
      // Do not use free variant (for openrouter).
      const modelName = model.name.split(":free")[0];
      const autoModelClient = {
-        model: provider(`${providerConfig.gatewayPrefix}${modelName}`),
+        model: provider(
+          `${providerConfig.gatewayPrefix || ""}${modelName}`,
+          settings.enableProLazyEditsMode
+            ? {
+                files,
+              }
+            : undefined,
+        ),
        builtinProviderId: "auto",
      };
      const googleSettings = settings.providerSettings?.google;
@@ -235,7 +261,7 @@ function getRegularModelClient(
        const provider = createOpenAICompatible({
          name: providerConfig.id,
          baseURL: providerConfig.apiBaseUrl,
-          apiKey: apiKey,
+          apiKey,
        });
        return {
          modelClient: {
--- a/src/ipc/utils/llm_engine_provider.ts
+++ b/src/ipc/utils/llm_engine_provider.ts
@@ -0,0 +1,159 @@
+import {
+  LanguageModelV1,
+  LanguageModelV1ObjectGenerationMode,
+} from "@ai-sdk/provider";
+import { OpenAICompatibleChatLanguageModel } from "@ai-sdk/openai-compatible";
+import {
+  FetchFunction,
+  loadApiKey,
+  withoutTrailingSlash,
+} from "@ai-sdk/provider-utils";
+
+import { OpenAICompatibleChatSettings } from "@ai-sdk/openai-compatible";
+import log from "electron-log";
+
+const logger = log.scope("llm_engine_provider");
+
+export type ExampleChatModelId = string & {};
+
+export interface ExampleChatSettings extends OpenAICompatibleChatSettings {
+  files?: { path: string; content: string }[];
+}
+export interface ExampleProviderSettings {
+  /**
+Example API key.
+*/
+  apiKey?: string;
+  /**
+Base URL for the API calls.
+*/
+  baseURL?: string;
+  /**
+Custom headers to include in the requests.
+*/
+  headers?: Record<string, string>;
+  /**
+Optional custom url query parameters to include in request urls.
+*/
+  queryParams?: Record<string, string>;
+  /**
+Custom fetch implementation. You can use it as a middleware to intercept requests,
+or to provide a custom fetch implementation for e.g. testing.
+*/
+  fetch?: FetchFunction;
+}
+
+export interface DyadEngineProvider {
+  /**
+Creates a model for text generation.
+*/
+  (
+    modelId: ExampleChatModelId,
+    settings?: ExampleChatSettings,
+  ): LanguageModelV1;
+
+  /**
+Creates a chat model for text generation.
+*/
+  chatModel(
+    modelId: ExampleChatModelId,
+    settings?: ExampleChatSettings,
+  ): LanguageModelV1;
+}
+
+export function createDyadEngine(
+  options: ExampleProviderSettings = {},
+): DyadEngineProvider {
+  const baseURL = withoutTrailingSlash(
+    options.baseURL ?? "https://api.example.com/v1",
+  );
+  const getHeaders = () => ({
+    Authorization: `Bearer ${loadApiKey({
+      apiKey: options.apiKey,
+      environmentVariableName: "DYAD_PRO_API_KEY",
+      description: "Example API key",
+    })}`,
+    ...options.headers,
+  });
+
+  interface CommonModelConfig {
+    provider: string;
+    url: ({ path }: { path: string }) => string;
+    headers: () => Record<string, string>;
+    fetch?: FetchFunction;
+  }
+
+  const getCommonModelConfig = (modelType: string): CommonModelConfig => ({
+    provider: `example.${modelType}`,
+    url: ({ path }) => {
+      const url = new URL(`${baseURL}${path}`);
+      if (options.queryParams) {
+        url.search = new URLSearchParams(options.queryParams).toString();
+      }
+      return url.toString();
+    },
+    headers: getHeaders,
+    fetch: options.fetch,
+  });
+
+  const createChatModel = (
+    modelId: ExampleChatModelId,
+    settings: ExampleChatSettings = {},
+  ) => {
+    // Extract files from settings to process them appropriately
+    const { files, ...restSettings } = settings;
+
+    // Create configuration with file handling
+    const config = {
+      ...getCommonModelConfig("chat"),
+      defaultObjectGenerationMode:
+        "tool" as LanguageModelV1ObjectGenerationMode,
+      // Custom fetch implementation that adds files to the request
+      fetch: files?.length
+        ? (input: RequestInfo | URL, init?: RequestInit) => {
+            // Use default fetch if no init or body
+            if (!init || !init.body || typeof init.body !== "string") {
+              return (options.fetch || fetch)(input, init);
+            }
+
+            try {
+              // Parse the request body to manipulate it
+              const parsedBody = JSON.parse(init.body);
+
+              // Add files to the request if they exist
+              if (files?.length) {
+                parsedBody.dyad_options = {
+                  files,
+                  enable_lazy_edits: true,
+                };
+              }
+
+              // Return modified request with files included
+              const modifiedInit = {
+                ...init,
+                body: JSON.stringify(parsedBody),
+              };
+
+              // Use the provided fetch or default fetch
+              return (options.fetch || fetch)(input, modifiedInit);
+            } catch (e) {
+              logger.error("Error parsing request body", e);
+              // If parsing fails, use original request
+              return (options.fetch || fetch)(input, init);
+            }
+          }
+        : options.fetch,
+    };
+
+    return new OpenAICompatibleChatLanguageModel(modelId, restSettings, config);
+  };
+
+  const provider = (
+    modelId: ExampleChatModelId,
+    settings?: ExampleChatSettings,
+  ) => createChatModel(modelId, settings);
+
+  provider.chatModel = createChatModel;
+
+  return provider;
+}
--- a/src/lib/schemas.ts
+++ b/src/lib/schemas.ts
@@ -119,6 +119,7 @@ export const UserSettingsSchema = z.object({
  lastShownReleaseNotesVersion: z.string().optional(),
  maxChatTurnsInContext: z.number().optional(),
  enableProSaverMode: z.boolean().optional(),
+  enableProLazyEditsMode: z.boolean().optional(),
  // DEPRECATED.
  runtimeMode: RuntimeModeSchema.optional(),
 });
--- a/src/utils/codebase.ts
+++ b/src/utils/codebase.ts
@@ -275,13 +275,19 @@ ${content}
 /**
 * Extract and format codebase files as a string to be included in prompts
 * @param appPath - Path to the codebase to extract
- * @returns A string containing formatted file contents
+ * @returns Object containing formatted output and individual files
 */
-export async function extractCodebase(appPath: string): Promise<string> {
+export async function extractCodebase(appPath: string): Promise<{
+  formattedOutput: string;
+  files: { path: string; content: string }[];
+}> {
  try {
    await fsAsync.access(appPath);
  } catch {
-    return `# Error: Directory ${appPath} does not exist or is not accessible`;
+    return {
+      formattedOutput: `# Error: Directory ${appPath} does not exist or is not accessible`,
+      files: [],
+    };
  }
  const startTime = Date.now();

@@ -292,15 +298,33 @@ export async function extractCodebase(appPath: string): Promise<string> {
  // This is important for cache-ability.
  const sortedFiles = await sortFilesByModificationTime(files);

-  // Format files
-  let output = "";
-  const formatPromises = sortedFiles.map((file) => formatFile(file, appPath));
+  // Format files and collect individual file contents
+  const filesArray: { path: string; content: string }[] = [];
+  const formatPromises = sortedFiles.map(async (file) => {
+    const formattedContent = await formatFile(file, appPath);
+
+    // Get raw content for the files array
+    const relativePath = path.relative(appPath, file);
+    const rawContent = await readFileWithCache(file);
+    if (rawContent !== null) {
+      filesArray.push({
+        path: relativePath,
+        content: rawContent,
+      });
+    }
+
+    return formattedContent;
+  });
+
  const formattedFiles = await Promise.all(formatPromises);
-  output = formattedFiles.join("");
+  const formattedOutput = formattedFiles.join("");

  const endTime = Date.now();
  logger.log("extractCodebase: time taken", endTime - startTime);
-  return output;
+  return {
+    formattedOutput,
+    files: filesArray,
+  };
 }

 /**