Smart files context (#184)

2025-05-16 22:21:45 -07:00
parent 2455c554ee
commit f9f33596bd
8 changed files with 224 additions and 31 deletions
--- a/src/components/ProModeSelector.tsx
+++ b/src/components/ProModeSelector.tsx
@@ -27,6 +27,12 @@ export function ProModeSelector() {
    });
  };

+  const toggleSmartContext = () => {
+    updateSettings({
+      enableProSmartFilesContextMode: !settings?.enableProSmartFilesContextMode,
+    });
+  };
+
  if (!settings?.enableDyadPro) {
    return null;
  }
@@ -105,6 +111,30 @@ export function ProModeSelector() {
              onCheckedChange={toggleLazyEdits}
            />
          </div>
+          <div className="flex items-center justify-between">
+            <div className="space-y-2">
+              <Label htmlFor="smart-context">Smart Context</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">
+                    Improve efficiency and save credits working on large
+                    codebases.
+                  </TooltipContent>
+                </Tooltip>
+                <p className="text-xs text-muted-foreground max-w-55">
+                  Automatically detects the most relevant files for your chat
+                </p>
+              </div>
+            </div>
+            <Switch
+              id="smart-context"
+              checked={Boolean(settings?.enableProSmartFilesContextMode)}
+              onCheckedChange={toggleSmartContext}
+            />
+          </div>
        </div>
      </PopoverContent>
    </Popover>
--- a/src/components/chat/DyadCodebaseContext.tsx
+++ b/src/components/chat/DyadCodebaseContext.tsx
@@ -0,0 +1,117 @@
+import React, { useState, useEffect } from "react";
+import { ChevronUp, ChevronDown, Code2, FileText } from "lucide-react";
+import { CustomTagState } from "./stateTypes";
+
+interface DyadCodebaseContextProps {
+  children: React.ReactNode;
+  node?: {
+    properties?: {
+      files?: string;
+      state?: CustomTagState;
+    };
+  };
+}
+
+export const DyadCodebaseContext: React.FC<DyadCodebaseContextProps> = ({
+  node,
+}) => {
+  const state = node?.properties?.state as CustomTagState;
+  const inProgress = state === "pending";
+  const [isExpanded, setIsExpanded] = useState(inProgress);
+  const files = node?.properties?.files?.split(",") || [];
+
+  // Collapse when transitioning from in-progress to not-in-progress
+  useEffect(() => {
+    if (!inProgress && isExpanded) {
+      setIsExpanded(false);
+    }
+  }, [inProgress]);
+
+  return (
+    <div
+      className={`relative bg-(--background-lightest) dark:bg-zinc-900 hover:bg-(--background-lighter) rounded-lg px-4 py-2 border my-2 cursor-pointer ${
+        inProgress ? "border-blue-500" : "border-border"
+      }`}
+      onClick={() => setIsExpanded(!isExpanded)}
+      role="button"
+      aria-expanded={isExpanded}
+      tabIndex={0}
+      onKeyDown={(e) => {
+        if (e.key === "Enter" || e.key === " ") {
+          e.preventDefault();
+          setIsExpanded(!isExpanded);
+        }
+      }}
+    >
+      {/* Top-left label badge */}
+      <div
+        className="absolute top-2 left-2 flex items-center gap-1 px-2 py-0.5 rounded text-xs font-semibold text-blue-500 bg-white dark:bg-zinc-900"
+        style={{ zIndex: 1 }}
+      >
+        <Code2 size={16} className="text-blue-500" />
+        <span>Codebase Context</span>
+      </div>
+
+      {/* File count when collapsed */}
+      {files.length > 0 && (
+        <div className="absolute top-2 left-40 flex items-center">
+          <span className="px-1.5 py-0.5 bg-gray-100 dark:bg-zinc-800 text-xs rounded text-gray-600 dark:text-gray-300">
+            Using {files.length} file{files.length !== 1 ? "s" : ""}
+          </span>
+        </div>
+      )}
+
+      {/* Indicator icon */}
+      <div className="absolute top-2 right-2 p-1 text-gray-500">
+        {isExpanded ? <ChevronUp size={16} /> : <ChevronDown size={16} />}
+      </div>
+
+      {/* Main content with smooth transition */}
+      <div
+        className="pt-6 overflow-hidden transition-all duration-300 ease-in-out"
+        style={{
+          maxHeight: isExpanded ? "1000px" : "0px",
+          opacity: isExpanded ? 1 : 0,
+          marginBottom: isExpanded ? "0" : "-6px", // Compensate for padding
+        }}
+      >
+        {/* File list when expanded */}
+        {files.length > 0 && (
+          <div className="mb-3">
+            <div className="flex flex-wrap gap-2 mt-2">
+              {files.map((file, index) => {
+                const filePath = file.trim();
+                const fileName = filePath.split("/").pop() || filePath;
+                const pathPart =
+                  filePath.substring(0, filePath.length - fileName.length) ||
+                  "";
+
+                return (
+                  <div
+                    key={index}
+                    className="px-2 py-1     bg-gray-100 dark:bg-zinc-800 rounded-lg"
+                  >
+                    <div className="flex items-center gap-1.5">
+                      <FileText
+                        size={14}
+                        className="text-gray-500 dark:text-gray-400 flex-shrink-0"
+                      />
+                      <div className="text-sm font-medium text-gray-700 dark:text-gray-300">
+                        {fileName}
+                      </div>
+                    </div>
+                    {pathPart && (
+                      <div className="text-xs text-gray-500 dark:text-gray-400 ml-5">
+                        {pathPart}
+                      </div>
+                    )}
+                  </div>
+                );
+              })}
+            </div>
+          </div>
+        )}
+      </div>
+    </div>
+  );
+};
--- a/src/components/chat/DyadMarkdownParser.tsx
+++ b/src/components/chat/DyadMarkdownParser.tsx
@@ -8,6 +8,7 @@ import { DyadAddDependency } from "./DyadAddDependency";
 import { DyadExecuteSql } from "./DyadExecuteSql";
 import { DyadAddIntegration } from "./DyadAddIntegration";
 import { DyadEdit } from "./DyadEdit";
+import { DyadCodebaseContext } from "./DyadCodebaseContext";
 import { CodeHighlight } from "./CodeHighlight";
 import { useAtomValue } from "jotai";
 import { isStreamingAtom } from "@/atoms/chatAtoms";
@@ -117,6 +118,7 @@ function preprocessUnclosedTags(content: string): {
    "dyad-output",
    "dyad-chat-summary",
    "dyad-edit",
+    "dyad-codebase-context",
  ];

  let processedContent = content;
@@ -180,6 +182,7 @@ function parseCustomTags(content: string): ContentPiece[] {
    "dyad-output",
    "dyad-chat-summary",
    "dyad-edit",
+    "dyad-codebase-context",
  ];

  const tagPattern = new RegExp(
@@ -362,6 +365,20 @@ function renderCustomTag(
        </DyadEdit>
      );

+    case "dyad-codebase-context":
+      return (
+        <DyadCodebaseContext
+          node={{
+            properties: {
+              files: attributes.files || "",
+              state: getState({ isStreaming, inProgress }),
+            },
+          }}
+        >
+          {content}
+        </DyadCodebaseContext>
+      );
+
    case "dyad-output":
      return (
        <DyadOutput
--- a/src/ipc/handlers/chat_stream_handlers.ts
+++ b/src/ipc/handlers/chat_stream_handlers.ts
@@ -236,11 +236,8 @@ export function registerChatStreamHandlers() {
          "estimated tokens",
          codebaseInfo.length / 4,
        );
-        const { modelClient, backupModelClients } = await getModelClient(
-          settings.selectedModel,
-          settings,
-          files,
-        );
+        const { modelClient, backupModelClients, isEngineEnabled } =
+          await getModelClient(settings.selectedModel, settings, files);

        // Prepare message history for the AI
        const messageHistory = updatedChat.messages.map((message) => ({
@@ -328,15 +325,22 @@ This conversation includes one or more image attachments. When the user uploads
 `;
        }

+        const codebasePrefix = isEngineEnabled
+          ? // No codebase prefix if engine is set, we will take of it there.
+            []
+          : ([
+              {
+                role: "user",
+                content: "This is my codebase. " + codebaseInfo,
+              },
+              {
+                role: "assistant",
+                content: "OK, got it. I'm ready to help",
+              },
+            ] as const);
+
        let chatMessages: CoreMessage[] = [
-          {
-            role: "user",
-            content: "This is my codebase. " + codebaseInfo,
-          },
-          {
-            role: "assistant",
-            content: "OK, got it. I'm ready to help",
-          },
+          ...codebasePrefix,
          ...limitedMessageHistory.map((msg) => ({
            role: msg.role as "user" | "assistant" | "system",
            content: msg.content,
--- a/src/ipc/utils/get_model_client.ts
+++ b/src/ipc/utils/get_model_client.ts
@@ -50,6 +50,7 @@ export async function getModelClient(
 ): Promise<{
  modelClient: ModelClient;
  backupModelClients: ModelClient[];
+  isEngineEnabled?: boolean;
 }> {
  const allProviders = await getLanguageModelProviders();

@@ -103,9 +104,12 @@ export async function getModelClient(
    // so we do a nullish and not a truthy check here.
    if (providerConfig.gatewayPrefix != null || dyadLocalEngine) {
      const languageModel = await findLanguageModel(model);
+      const engineProMode =
+        settings.enableProSmartFilesContextMode ||
+        settings.enableProLazyEditsMode;
      // Currently engine is only used for turbo edits.
      const isEngineEnabled = Boolean(
-        settings.enableProLazyEditsMode &&
+        engineProMode &&
          languageModel?.type === "cloud" &&
          languageModel?.supportsTurboEdits,
      );
@@ -113,6 +117,10 @@ export async function getModelClient(
        ? createDyadEngine({
            apiKey: dyadApiKey,
            baseURL: dyadLocalEngine ?? "https://engine.dyad.sh/v1",
+            dyadOptions: {
+              enableLazyEdits: settings.enableProLazyEditsMode,
+              enableSmartFilesContext: settings.enableProSmartFilesContextMode,
+            },
          })
        : createOpenAICompatible({
            name: "dyad-gateway",
@@ -126,7 +134,7 @@ export async function getModelClient(
      const autoModelClient = {
        model: provider(
          `${providerConfig.gatewayPrefix || ""}${modelName}`,
-          settings.enableProLazyEditsMode
+          engineProMode
            ? {
                files,
              }
@@ -161,6 +169,7 @@ export async function getModelClient(
            providerConfig,
          ).modelClient,
          backupModelClients: [autoModelClient],
+          isEngineEnabled,
        };
      } else {
        return {
--- a/src/ipc/utils/llm_engine_provider.ts
+++ b/src/ipc/utils/llm_engine_provider.ts
@@ -41,6 +41,11 @@ Custom fetch implementation. You can use it as a middleware to intercept request
 or to provide a custom fetch implementation for e.g. testing.
 */
  fetch?: FetchFunction;
+
+  dyadOptions: {
+    enableLazyEdits?: boolean;
+    enableSmartFilesContext?: boolean;
+  };
 }

 export interface DyadEngineProvider {
@@ -62,7 +67,7 @@ Creates a chat model for text generation.
 }

 export function createDyadEngine(
-  options: ExampleProviderSettings = {},
+  options: ExampleProviderSettings,
 ): DyadEngineProvider {
  const baseURL = withoutTrailingSlash(
    options.baseURL ?? "https://api.example.com/v1",
@@ -124,7 +129,9 @@ export function createDyadEngine(
              if (files?.length) {
                parsedBody.dyad_options = {
                  files,
-                  enable_lazy_edits: true,
+                  enable_lazy_edits: options.dyadOptions.enableLazyEdits,
+                  enable_smart_files_context:
+                    options.dyadOptions.enableSmartFilesContext,
                };
              }

--- a/src/lib/schemas.ts
+++ b/src/lib/schemas.ts
@@ -120,6 +120,7 @@ export const UserSettingsSchema = z.object({
  maxChatTurnsInContext: z.number().optional(),
  enableProSaverMode: z.boolean().optional(),
  enableProLazyEditsMode: z.boolean().optional(),
+  enableProSmartFilesContextMode: z.boolean().optional(),
  // DEPRECATED.
  runtimeMode: RuntimeModeSchema.optional(),
 });
--- a/src/utils/codebase.ts
+++ b/src/utils/codebase.ts
@@ -223,6 +223,21 @@ async function collectFiles(dir: string, baseDir: string): Promise<string[]> {
  return files;
 }

+// Skip large configuration files or generated code (just include the path)
+function isOmittedFile(relativePath: string): boolean {
+  return (
+    relativePath.includes(path.join("src", "components", "ui")) ||
+    relativePath.includes("eslint.config") ||
+    relativePath.includes("tsconfig.json") ||
+    relativePath.includes("package-lock.json") ||
+    // These should already be excluded based on file type, but
+    // just in case, we'll redact the contents here.
+    relativePath.includes(".env")
+  );
+}
+
+const OMITTED_FILE_CONTENT = "// Contents omitted for brevity";
+
 /**
 * Format a file for inclusion in the codebase extract
 */
@@ -230,18 +245,9 @@ async function formatFile(filePath: string, baseDir: string): Promise<string> {
  try {
    const relativePath = path.relative(baseDir, filePath);

-    // Skip large configuration files or generated code (just include the path)
-    if (
-      relativePath.includes(path.join("src", "components", "ui")) ||
-      relativePath.includes("eslint.config") ||
-      relativePath.includes("tsconfig.json") ||
-      relativePath.includes("package-lock.json") ||
-      // These should already be excluded based on file type, but
-      // just in case, we'll redact the contents here.
-      relativePath.includes(".env")
-    ) {
+    if (isOmittedFile(relativePath)) {
      return `<dyad-file path="${relativePath}">
-// Contents omitted for brevity
+${OMITTED_FILE_CONTENT}
 </dyad-file>

 `;
@@ -305,11 +311,13 @@ export async function extractCodebase(appPath: string): Promise<{

    // Get raw content for the files array
    const relativePath = path.relative(appPath, file);
-    const rawContent = await readFileWithCache(file);
-    if (rawContent !== null) {
+    const fileContent = isOmittedFile(relativePath)
+      ? OMITTED_FILE_CONTENT
+      : await readFileWithCache(file);
+    if (fileContent !== null) {
      filesArray.push({
        path: relativePath,
-        content: rawContent,
+        content: fileContent,
      });
    }