wip

sourcegraph · Feb 1, 2025 · c82e8cc · c82e8cc
1 parent cea74cc
commit c82e8cc
Show file tree

Hide file tree

Showing 16 changed files with 152 additions and 2 deletions.
diff --git a/lib/shared/src/context/openctx/api.ts b/lib/shared/src/context/openctx/api.ts
@@ -32,3 +32,4 @@ export const REMOTE_DIRECTORY_PROVIDER_URI = 'internal-remote-directory-search'
 export const WEB_PROVIDER_URI = 'internal-web-provider'
 export const GIT_OPENCTX_PROVIDER_URI = 'internal-git-openctx-provider'
 export const CODE_SEARCH_PROVIDER_URI = 'internal-code-search-provider'
+export const RULES_PROVIDER_URI = 'internal-rules-provider'
diff --git a/lib/shared/src/context/openctx/context.ts b/lib/shared/src/context/openctx/context.ts
@@ -51,7 +51,7 @@ export const getContextForChatMessage = async (
                         providerUri: item.providerUri,
                         content: item.ai?.content || '',
                         provider: 'openctx',
-                    }) as ContextItemOpenCtx
+                    }) satisfies ContextItemOpenCtx
             )
     } catch {
         return []

diff --git a/lib/shared/src/index.ts b/lib/shared/src/index.ts
@@ -360,6 +360,7 @@ export {
     WEB_PROVIDER_URI,
     GIT_OPENCTX_PROVIDER_URI,
     CODE_SEARCH_PROVIDER_URI,
+    RULES_PROVIDER_URI,
 } from './context/openctx/api'
 export * from './context/openctx/context'
 export * from './lexicalEditor/editorState'

diff --git a/vscode/src/autoedits/autoedits-config.ts b/vscode/src/autoedits/autoedits-config.ts
@@ -36,6 +36,7 @@ const defaultTokenLimit = {
         [RetrieverIdentifier.RecentCopyRetriever]: 500,
         [RetrieverIdentifier.DiagnosticsRetriever]: 500,
         [RetrieverIdentifier.RecentViewPortRetriever]: 2500,
+        [RetrieverIdentifier.RulesRetriever]: 1000,
     },
 } as const satisfies AutoEditsTokenLimit
 

diff --git a/vscode/src/autoedits/prompt/constants.ts b/vscode/src/autoedits/prompt/constants.ts
@@ -19,6 +19,10 @@ export const AREA_FOR_CODE_MARKER_OPEN = ps`<area_around_code_to_rewrite>`
 export const AREA_FOR_CODE_MARKER_CLOSE = ps`</area_around_code_to_rewrite>`
 export const CODE_TO_REWRITE_TAG_CLOSE = ps`</code_to_rewrite>`
 export const CODE_TO_REWRITE_TAG_OPEN = ps`<code_to_rewrite>`
+export const RULES_TAG_OPEN = ps`<rules>`
+export const RULES_TAG_CLOSE = ps`</rules>`
+export const RULE_TAG_OPEN = ps`<rule>`
+export const RULE_TAG_CLOSE = ps`</rule>`
 
 // Some common prompt instructions
 export const SYSTEM_PROMPT = ps`You are an intelligent programmer named CodyBot. You are an expert at coding. Your goal is to help your colleague finish a code change.`
@@ -31,3 +35,4 @@ export const LINT_ERRORS_INSTRUCTION = ps`Here are some linter errors from the c
 export const RECENT_COPY_INSTRUCTION = ps`Here is some recent code I copied from the editor.`
 export const CURRENT_FILE_INSTRUCTION = ps`Here is the file that I am looking at `
 export const SHORT_TERM_SNIPPET_VIEWS_INSTRUCTION = ps`Here are some snippets of code I just looked at:`
+export const RULES_INSTRUCTION = ps`Your response MUST comply with the following rules: `
diff --git a/vscode/src/autoedits/prompt/default-prompt-strategy.test.ts b/vscode/src/autoedits/prompt/default-prompt-strategy.test.ts
@@ -62,6 +62,7 @@ describe('DefaultUserPromptStrategy', () => {
                 [RetrieverIdentifier.RecentCopyRetriever]: 100,
                 [RetrieverIdentifier.JaccardSimilarityRetriever]: 100,
                 [RetrieverIdentifier.DiagnosticsRetriever]: 100,
+                [RetrieverIdentifier.RulesRetriever]: 100,
             },
         }
 

diff --git a/vscode/src/autoedits/prompt/default-prompt-strategy.ts b/vscode/src/autoedits/prompt/default-prompt-strategy.ts
@@ -16,6 +16,7 @@ import {
     getRecentCopyPrompt,
     getRecentEditsPrompt,
     getRecentlyViewedSnippetsPrompt,
+    getRulesPrompt,
     joinPromptsWithNewlineSeparator,
 } from './prompt-utils'
 
@@ -55,6 +56,12 @@ export class DefaultUserPromptStrategy extends AutoeditsUserPromptStrategy {
             getJaccardSimilarityPrompt
         )
 
+        const rulesPrompt = getPromptForTheContextSource(
+            contextItemMapping.get(RetrieverIdentifier.RulesRetriever) || [],
+            constants.RULES_INSTRUCTION,
+            getRulesPrompt
+        )
+
         const { fileWithMarkerPrompt, areaPrompt } = getCurrentFilePromptComponents(
             document,
             codeToReplaceData
@@ -71,6 +78,7 @@ export class DefaultUserPromptStrategy extends AutoeditsUserPromptStrategy {
             getPromptWithNewline(lintErrorsPrompt),
             getPromptWithNewline(recentCopyPrompt),
             getPromptWithNewline(areaPrompt),
+            getPromptWithNewline(rulesPrompt),
             constants.FINAL_USER_PROMPT
         )
 

diff --git a/vscode/src/autoedits/prompt/prompt-utils.ts b/vscode/src/autoedits/prompt/prompt-utils.ts
@@ -389,6 +389,24 @@ export function getJaccardSimilarityPrompt(contextItems: AutocompleteContextSnip
     )
 }
 
+export function getRulesPrompt(contextItems: AutocompleteContextSnippet[]): PromptString {
+    if (contextItems.length === 0) {
+        return ps``
+    }
+    const rulesPrompts = contextItems.map(item =>
+        joinPromptsWithNewlineSeparator(
+            constants.RULE_TAG_OPEN,
+            PromptString.fromAutocompleteContextSnippet(item).content,
+            constants.RULE_TAG_CLOSE
+        )
+    )
+    return joinPromptsWithNewlineSeparator(
+        constants.RULES_TAG_OPEN,
+        PromptString.join(rulesPrompts, ps`\n`),
+        constants.RULES_TAG_CLOSE
+    )
+}
+
 //  Helper functions
 export function getContextItemMappingWithTokenLimit(
     contextItems: AutocompleteContextSnippet[],

diff --git a/vscode/src/autoedits/prompt/short-term-diff-prompt-strategy.test.ts b/vscode/src/autoedits/prompt/short-term-diff-prompt-strategy.test.ts
@@ -67,6 +67,7 @@ describe('ShortTermPromptStrategy', () => {
                     [RetrieverIdentifier.RecentCopyRetriever]: 100,
                     [RetrieverIdentifier.JaccardSimilarityRetriever]: 100,
                     [RetrieverIdentifier.DiagnosticsRetriever]: 100,
+                    [RetrieverIdentifier.RulesRetriever]: 1000,
                 },
             }
             const codeToReplaceData = getCodeToReplaceData({

diff --git a/vscode/src/autoedits/prompt/short-term-diff-prompt-strategy.ts b/vscode/src/autoedits/prompt/short-term-diff-prompt-strategy.ts
@@ -18,6 +18,7 @@ import {
     getRecentCopyPrompt,
     getRecentEditsPrompt,
     getRecentlyViewedSnippetsPrompt,
+    getRulesPrompt,
     joinPromptsWithNewlineSeparator,
 } from './prompt-utils'
 
@@ -53,6 +54,12 @@ export class ShortTermPromptStrategy extends AutoeditsUserPromptStrategy {
             getJaccardSimilarityPrompt
         )
 
+        const rulesPrompt = getPromptForTheContextSource(
+            contextItemMapping.get(RetrieverIdentifier.RulesRetriever) || [],
+            constants.RULES_INSTRUCTION,
+            getRulesPrompt
+        )
+
         const { fileWithMarkerPrompt, areaPrompt } = getCurrentFilePromptComponents(
             document,
             codeToReplaceData
@@ -71,6 +78,7 @@ export class ShortTermPromptStrategy extends AutoeditsUserPromptStrategy {
             getPromptWithNewline(recentCopyPrompt),
             getPromptWithNewline(areaPrompt),
             getPromptWithNewline(shortTermEditsPrompt),
+            getPromptWithNewline(rulesPrompt),
             constants.FINAL_USER_PROMPT
         )
 

diff --git a/vscode/src/completions/context/context-strategy.ts b/vscode/src/completions/context/context-strategy.ts
@@ -15,6 +15,7 @@ import { LineLevelDiffStrategy } from './retrievers/recent-user-actions/recent-e
 import { UnifiedDiffStrategy } from './retrievers/recent-user-actions/recent-edits-diff-helpers/unified-diff'
 import { RecentEditsRetriever } from './retrievers/recent-user-actions/recent-edits-retriever'
 import { RecentViewPortRetriever } from './retrievers/recent-user-actions/recent-view-port'
+import { RulesRetriever } from './retrievers/rules-retriever'
 import { loadTscRetriever } from './retrievers/tsc/load-tsc-retriever'
 
 export type ContextStrategy =
@@ -151,6 +152,7 @@ export class DefaultContextStrategyFactory implements ContextStrategyFactory {
                                     maxTrackedViewPorts: 50,
                                     maxRetrievedViewPorts: 10,
                                 }),
+                                new RulesRetriever(),
                             ]
                             break
                         case 'jaccard-similarity':

diff --git a/vscode/src/completions/context/retrievers/rules-retriever.ts b/vscode/src/completions/context/retrievers/rules-retriever.ts
@@ -0,0 +1,39 @@
+import {
+    type AutocompleteContextSnippet,
+    type RuleService,
+    firstValueFrom,
+} from '@sourcegraph/cody-shared'
+import type { Disposable } from 'vscode'
+import { URI } from 'vscode-uri'
+import { createFileSystemRuleService } from '../../../rules/fs-rules'
+import type { ContextRetriever, ContextRetrieverOptions } from '../../types'
+import { RetrieverIdentifier } from '../utils'
+
+export class RulesRetriever implements Disposable, ContextRetriever {
+    public identifier = RetrieverIdentifier.RulesRetriever
+
+    private ruleService: RuleService
+    public constructor() {
+        this.ruleService = createFileSystemRuleService()
+    }
+
+    public async retrieve({ document }: ContextRetrieverOptions): Promise<AutocompleteContextSnippet[]> {
+        const rules = await firstValueFrom(this.ruleService.rulesForPaths([document.uri]))
+
+        return rules.map(
+            rule =>
+                ({
+                    type: 'base',
+                    identifier: this.identifier,
+                    content: rule.instruction,
+                    uri: URI.parse(rule.uri),
+                }) satisfies AutocompleteContextSnippet
+        )
+    }
+
+    public isSupportedForLanguageId(): boolean {
+        return true
+    }
+
+    public dispose() {}
+}
diff --git a/vscode/src/completions/context/utils.ts b/vscode/src/completions/context/utils.ts
@@ -9,6 +9,7 @@ export enum RetrieverIdentifier {
     RecentCopyRetriever = 'recent-copy',
     DiagnosticsRetriever = 'diagnostics',
     RecentViewPortRetriever = 'recent-view-port',
+    RulesRetriever = 'rules',
 }
 
 export interface ShouldUseContextParams {

diff --git a/vscode/src/context/openctx.ts b/vscode/src/context/openctx.ts
@@ -6,6 +6,7 @@ import {
     type CodyClientConfig,
     FeatureFlag,
     GIT_OPENCTX_PROVIDER_URI,
+    RULES_PROVIDER_URI,
     WEB_PROVIDER_URI,
     authStatus,
     clientCapabilities,
@@ -41,6 +42,7 @@ import LinearIssuesProvider from './openctx/linear-issues'
 import RemoteDirectoryProvider, { createRemoteDirectoryProvider } from './openctx/remoteDirectorySearch'
 import RemoteFileProvider, { createRemoteFileProvider } from './openctx/remoteFileSearch'
 import RemoteRepositorySearch, { createRemoteRepositoryProvider } from './openctx/remoteRepositorySearch'
+import { createRulesProvider } from './openctx/rules'
 import { createWebProvider } from './openctx/web'
 
 export function exposeOpenCtxClient(
@@ -146,6 +148,11 @@ export function getOpenCtxProviders(
                     provider: createWebProvider(false),
                     providerUri: WEB_PROVIDER_URI,
                 },
+                {
+                    settings: true,
+                    provider: createRulesProvider(),
+                    providerUri: RULES_PROVIDER_URI,
+                },
             ]
 
             if (!isDotCom(authStatus)) {

diff --git a/vscode/src/context/openctx/rules.ts b/vscode/src/context/openctx/rules.ts
@@ -0,0 +1,55 @@
+import {
+    RULES_PROVIDER_URI,
+    type RulePromptPart,
+    firstValueFrom,
+    pluralize,
+} from '@sourcegraph/cody-shared'
+import { URI } from 'vscode-uri'
+import { createFileSystemRuleService } from '../../rules/fs-rules'
+import type { OpenCtxProvider } from './types'
+
+/**
+ * An OpenCtx provider that attaches rules for the current file.
+ */
+export function createRulesProvider(): OpenCtxProvider {
+    const ruleService = createFileSystemRuleService()
+    return {
+        providerUri: RULES_PROVIDER_URI,
+
+        meta() {
+            return {
+                name: 'Rules',
+                mentions: { autoInclude: true },
+            }
+        },
+
+        async mentions({ autoInclude, uri }) {
+            if (!uri || !autoInclude) {
+                return []
+            }
+
+            const rules = await firstValueFrom(ruleService.rulesForPaths([URI.parse(uri)]))
+            return rules.length === 0
+                ? []
+                : [
+                      {
+                          title: `${rules.length} ${pluralize('rule', rules.length)}`,
+                          uri: 'rules+openctx://rules', // dummy URI
+                          data: { rules },
+                      },
+                  ]
+        },
+
+        async items(params) {
+            const rules = params.mention?.data?.rules as RulePromptPart[] | undefined
+            return (
+                rules?.map(rule => ({
+                    title: rule.title ?? rule.display_name,
+                    uri: rule.uri,
+                    url: rule.uri,
+                    ai: { content: rule.instruction },
+                })) ?? []
+            )
+        },
+    }
+}
diff --git a/vscode/src/rules/fs-rules.ts b/vscode/src/rules/fs-rules.ts
@@ -98,7 +98,9 @@ export function createFileSystemRuleService(): RuleService {
                                 } catch (error) {
                                     if (
                                         !(
-                                            error instanceof vscode.FileSystemError &&
+                                            error &&
+                                            typeof error === 'object' &&
+                                            'code' in error &&
                                             error.code === 'FileNotFound'
                                         )
                                     ) {