Skip to content

Commit

Permalink
wip - apply feedback 2
Browse files Browse the repository at this point in the history
  • Loading branch information
@abeatrix
abeatrix committed Jan 7, 2025
1 parent 848120d commit 2e5d77e
Show file tree
Hide file tree
Showing 11 changed files with 156 additions and 118 deletions.
14 changes: 11 additions & 3 deletions lib/shared/src/configuration.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -466,12 +466,20 @@ export interface FireworksCodeCompletionParams {

export interface AgentToolboxSettings {
/**
* The agent ID that user has currently enabled.
* The agent that user has currently enabled.
*/
agent?: string
agent?: {
/**
* The name of the agent that user has currently enabled. E.g. "deep-cody"
*/
name?: string
}
/**
* Whether the user has enabled terminal context.
* Defaulted to undefined if shell context is not enabled by site admin via feature flag.
*/
shell?: boolean
shell?: {
enabled: boolean
error?: string
}
}
71 changes: 38 additions & 33 deletions lib/shared/src/models/sync.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -205,47 +205,52 @@ export function syncModels({
featureFlagProvider.evaluatedFeatureFlag(
FeatureFlag.CodyEarlyAccess
),
featureFlagProvider.evaluatedFeatureFlag(FeatureFlag.DeepCody),
featureFlagProvider.evaluatedFeatureFlag(
FeatureFlag.CodyChatDefaultToClaude35Haiku
),
enableToolCody
).pipe(
switchMap(([hasEarlyAccess, defaultToHaiku]) => {
// TODO(sqs): remove waitlist from localStorage when user has access
const isOnWaitlist = config.clientState.waitlist_o1
if (isDotComUser && (hasEarlyAccess || isOnWaitlist)) {
data.primaryModels = data.primaryModels.map(model => {
if (model.tags.includes(ModelTag.Waitlist)) {
const newTags = model.tags.filter(
tag => tag !== ModelTag.Waitlist
)
newTags.push(
hasEarlyAccess
? ModelTag.EarlyAccess
: ModelTag.OnWaitlist
)
return { ...model, tags: newTags }
switchMap(
([hasEarlyAccess, isDeepCodyEnabled, defaultToHaiku]) => {
// TODO(sqs): remove waitlist from localStorage when user has access
const isOnWaitlist = config.clientState.waitlist_o1
if (isDotComUser && (hasEarlyAccess || isOnWaitlist)) {
data.primaryModels = data.primaryModels.map(
model => {
if (model.tags.includes(ModelTag.Waitlist)) {
const newTags = model.tags.filter(
tag => tag !== ModelTag.Waitlist
)
newTags.push(
hasEarlyAccess
? ModelTag.EarlyAccess
: ModelTag.OnWaitlist
)
return { ...model, tags: newTags }
}
return model
}
)
}
if (isDeepCodyEnabled && enableToolCody) {
data.primaryModels.push(
createModelFromServerModel(TOOL_CODY_MODEL)
)
}
// set the default model to Haiku for free users
if (isDotComUser && isCodyFreeUser && defaultToHaiku) {
const haikuModel = data.primaryModels.find(m =>
m.id.includes('claude-3-5-haiku')
)
if (haikuModel) {
data.preferences!.defaults.chat = haikuModel.id
}
return model
})
}
if (enableToolCody) {
data.primaryModels.push(
createModelFromServerModel(TOOL_CODY_MODEL)
)
}
// set the default model to Haiku for free users
if (isDotComUser && isCodyFreeUser && defaultToHaiku) {
const haikuModel = data.primaryModels.find(m =>
m.id.includes('claude-3-5-haiku')
)
if (haikuModel) {
data.preferences!.defaults.chat = haikuModel.id
}
}

return Observable.of(data)
})
return Observable.of(data)
}
)
)
})
)
Expand Down
4 changes: 2 additions & 2 deletions vscode/src/chat/agentic/CodyTool.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -209,7 +209,7 @@ describe('CodyTool', () => {
})

it('should register all default tools based on toolbox settings', () => {
const mockedToolboxSettings = { agent: 'mock-agent', shell: true }
const mockedToolboxSettings = { agent: { name: 'deep-cody' }, shell: { enabled: true } }
vi.spyOn(toolboxManager, 'getSettings').mockReturnValue(mockedToolboxSettings)
const localStorageData: { [key: string]: unknown } = {}
mockLocalStorage({
Expand All @@ -226,7 +226,7 @@ describe('CodyTool', () => {
expect(tools.some(t => t.config.title.includes('Terminal'))).toBeTruthy()

// Disable shell and check if terminal tool is removed.
mockedToolboxSettings.shell = false
mockedToolboxSettings.shell.enabled = false
vi.spyOn(toolboxManager, 'getSettings').mockReturnValue(mockedToolboxSettings)
const newTools = CodyToolProvider.getTools()
expect(newTools.some(t => t.config.title.includes('Terminal'))).toBeFalsy()
Expand Down
16 changes: 8 additions & 8 deletions vscode/src/chat/agentic/CodyTool.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,17 +69,17 @@ export abstract class CodyTool {
*/
protected parse(): string[] {
const { subTag } = this.config.tags
const regex = new RegExp(`<${subTag}>(.+?)</?${subTag}>`, 's')
const parsed = (this.unprocessedText.match(new RegExp(regex, 'g')) || [])
.map(match => regex.exec(match)?.[1].trim())
.filter(Boolean) as string[]
const filtered = parsed.filter(q => !this.performedQueries.has(q))
// Store the search query to avoid running the same query again.
for (const query of filtered) {
const regex = new RegExp(`<${subTag}>(.+?)</?${subTag}>`, 'gs')
// Use matchAll for more efficient iteration and destructuring
const newQueries = [...this.unprocessedText.matchAll(regex)]
.map(([, group]) => group?.trim())
.filter(query => query && !this.performedQueries.has(query))
// Add all new queries to the set at once
for (const query of newQueries) {
this.performedQueries.add(query)
}
this.reset()
return filtered
return newQueries
}
/**
* The raw text input stream.
Expand Down
10 changes: 8 additions & 2 deletions vscode/src/chat/agentic/CodyToolProvider.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,13 +89,19 @@ describe('CodyToolProvider', () => {
})

it('should not include CLI tool if shell is disabled', () => {
vi.spyOn(toolboxManager, 'getSettings').mockReturnValue({ agent: 'deep-cody', shell: false })
vi.spyOn(toolboxManager, 'getSettings').mockReturnValue({
agent: { name: 'deep-cody' },
shell: { enabled: false },
})
const tools = CodyToolProvider.getTools()
expect(tools.some(tool => tool.config.title === 'Terminal')).toBe(false)
})

it('should include CLI tool if shell is enabled', () => {
vi.spyOn(toolboxManager, 'getSettings').mockReturnValue({ agent: 'deep-cody', shell: true })
vi.spyOn(toolboxManager, 'getSettings').mockReturnValue({
agent: { name: 'deep-cody' },
shell: { enabled: true },
})
const tools = CodyToolProvider.getTools()
expect(tools.some(tool => tool.config.title === 'Terminal')).toBe(true)
})
Expand Down
50 changes: 30 additions & 20 deletions vscode/src/chat/agentic/CodyToolProvider.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@ import { type CodyTool, type CodyToolConfig, OpenCtxTool, TOOL_CONFIGS } from '.
import { toolboxManager } from './ToolboxManager'
import { OPENCTX_TOOL_CONFIG } from './config'

type Retriever = Pick<ContextRetriever, 'retrieveContext'>

/**
* Interface for tool execution status callbacks.
* Used to track and report tool execution progress.
Expand All @@ -29,7 +31,7 @@ export interface ToolStatusCallback {
*/
export interface ToolConfiguration extends CodyToolConfig {
name: string
createInstance: (config: CodyToolConfig, ...args: any[]) => CodyTool
createInstance: (config: CodyToolConfig, retriever?: Retriever) => CodyTool
}

/**
Expand All @@ -44,30 +46,30 @@ export interface ToolConfiguration extends CodyToolConfig {
export class ToolFactory {
private tools: Map<string, ToolConfiguration> = new Map()

constructor(private contextRetriever: Pick<ContextRetriever, 'retrieveContext'>) {}
constructor(private contextRetriever: Retriever) {}

public register(toolConfig: ToolConfiguration): void {
this.tools.set(toolConfig.name, toolConfig)
}

public createTool(name: string, ...args: any[]): CodyTool | undefined {
public createTool(name: string, retriever?: Retriever): CodyTool | undefined {
const config = this.tools.get(name)
if (!config) {
return undefined
}
const instance = config.createInstance(config, ...args)
const instance = config.createInstance(config, retriever)
return instance
}

public getInstances(): CodyTool[] {
// Create fresh instances of all registered tools
return Array.from(this.tools.entries())
.filter(([name]) => name !== 'CliTool' || toolboxManager.getSettings()?.shell)
.filter(([name]) => name !== 'CliTool' || toolboxManager.getSettings()?.shell?.enabled)
.map(([_, config]) => config.createInstance(config, this.contextRetriever))
.filter(isDefined) as CodyTool[]
}

public createDefaultTools(contextRetriever?: Pick<ContextRetriever, 'retrieveContext'>): CodyTool[] {
public createDefaultTools(contextRetriever?: Retriever): CodyTool[] {
return Object.entries(TOOL_CONFIGS)
.map(([name]) => this.createTool(name, contextRetriever))
.filter(isDefined) as CodyTool[]
Expand All @@ -90,8 +92,16 @@ export class ToolFactory {

private generateToolName(provider: ContextMentionProviderMetadata): string {
const suffix = provider.id.includes('modelcontextprotocol') ? 'MCP' : ''
const title = provider.title.replace(' ', '').split('/').at(-1)
return 'TOOL' + title?.toUpperCase().replace(/[^a-zA-Z0-9]/g, '') + suffix
return (
'TOOL' +
provider.title
.split('/')
.pop()
?.replace(/\s+/g, '')
?.toUpperCase()
?.replace(/[^A-Z0-9]/g, '') +
suffix
)
}

private getToolConfig(provider: ContextMentionProviderMetadata): CodyToolConfig {
Expand Down Expand Up @@ -140,7 +150,7 @@ export namespace CodyToolProvider {
export let factory: ToolFactory
let openCtxSubscription: Unsubscribable | undefined

export function initialize(contextRetriever: Pick<ContextRetriever, 'retrieveContext'>): void {
export function initialize(contextRetriever: Retriever): void {
factory = new ToolFactory(contextRetriever)
initializeRegistry()
}
Expand All @@ -151,17 +161,12 @@ export namespace CodyToolProvider {
}

export function setupOpenCtxProviderListener(): void {
if (!openCtx.controller) {
console.error('OpenCtx controller not available')
if (!openCtxSubscription && factory && openCtx.controller) {
openCtxSubscription = openCtx.controller
.metaChanges({}, {})
.pipe(map(providers => providers.filter(p => !!p.mentions).map(openCtxProviderMetadata)))
.subscribe(providerMeta => factory?.createOpenCtxTools(providerMeta))
}
if (openCtxSubscription || !openCtx.controller) {
return
}

openCtxSubscription = openCtx.controller
.metaChanges({}, {})
.pipe(map(providers => providers.filter(p => !!p.mentions).map(openCtxProviderMetadata)))
.subscribe(providerMeta => factory.createOpenCtxTools(providerMeta))
}

function initializeRegistry(): void {
Expand All @@ -170,7 +175,12 @@ export namespace CodyToolProvider {
name,
...tool.prototype.config,
createInstance: useContextRetriever
? (_, contextRetriever) => new tool(contextRetriever)
? (_, contextRetriever) => {
if (!contextRetriever) {
throw new Error(`Context retriever required for ${name}`)
}
return new tool(contextRetriever)
}
: () => new tool(),
})
}
Expand Down
29 changes: 20 additions & 9 deletions vscode/src/chat/agentic/ToolboxManager.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ type StoredToolboxSettings = {
* NOTE: This is a Singleton class.
*/
class ToolboxManager {
private static readonly STORAGE_KEY = 'CODY_CHATAGENTS_TOOLBOX_SETTINGS'
private static readonly STORAGE_KEY = 'CODYAGENT_TOOLBOX_SETTINGS'
private static instance?: ToolboxManager

private constructor() {
Expand All @@ -51,8 +51,8 @@ class ToolboxManager {
private getStoredUserSettings(): StoredToolboxSettings {
return (
localStorage.get<StoredToolboxSettings>(ToolboxManager.STORAGE_KEY) ?? {
agent: undefined,
shell: this.shellConfig.user,
agent: this.isEnabled ? 'deep-cody' : undefined,
shell: false,
}
)
}
Expand All @@ -62,20 +62,31 @@ class ToolboxManager {
return null
}
const { agent, shell } = this.getStoredUserSettings()
const isShellEnabled = this.shellConfig.instance && this.shellConfig.client ? shell : undefined
return {
agent,
agent: { name: agent },
// Only show shell option if it's supported by instance and client.
shell: this.shellConfig.instance && this.shellConfig.client ? shell : undefined,
shell: { enabled: isShellEnabled ?? false },
}
}

public async updateToolboxSettings(settings: AgentToolboxSettings): Promise<void> {
public async updateSettings(settings: AgentToolboxSettings): Promise<void> {
logDebug('ToolboxManager', 'Updating toolbox settings', { verbose: settings })
await localStorage.set(ToolboxManager.STORAGE_KEY, settings)
await localStorage.set(ToolboxManager.STORAGE_KEY, {
agent: settings.agent?.name,
shell: settings.shell?.enabled ?? false,
})
this.changeNotifications.next()
}

public readonly settings: Observable<AgentToolboxSettings | null> = combineLatest(
/**
* Returns a real-time Observable stream of toolbox settings that updates when any of the following changes:
* - Feature flags
* - User subscription
* - Available models
* - Manual settings updates
* Use this when you need to react to settings changes over time.
*/
public readonly observable: Observable<AgentToolboxSettings | null> = combineLatest(
featureFlagProvider.evaluatedFeatureFlag(FeatureFlag.DeepCody),
featureFlagProvider.evaluatedFeatureFlag(FeatureFlag.ContextAgentDefaultChatModel),
featureFlagProvider.evaluatedFeatureFlag(FeatureFlag.DeepCodyShellContext),
Expand Down
6 changes: 3 additions & 3 deletions vscode/src/chat/chat-view/ChatController.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -663,7 +663,7 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
editorState,
intent: detectedIntent,
manuallySelectedIntent: manuallySelectedIntent ? detectedIntent : undefined,
agent: toolboxManager.getSettings()?.agent,
agent: toolboxManager.getSettings()?.agent?.name,
})
this.postViewTranscript({ speaker: 'assistant' })

Expand Down Expand Up @@ -1544,10 +1544,10 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
userProductSubscription.pipe(
map(value => (value === pendingOperation ? null : value))
),
toolboxSettings: () => toolboxManager.settings,
toolboxSettings: () => toolboxManager.observable,
updateToolboxSettings: settings => {
return promiseFactoryToObservable(async () => {
await toolboxManager.updateToolboxSettings(settings)
await toolboxManager.updateSettings(settings)
})
},
}
Expand Down
3 changes: 0 additions & 3 deletions vscode/src/chat/chat-view/handlers/DeepCodyHandler.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ import {
featureFlagProvider,
storeLastValue,
} from '@sourcegraph/cody-shared'
import { CodyToolProvider } from '../../agentic/CodyToolProvider'
import { DeepCodyAgent } from '../../agentic/DeepCody'
import { DeepCodyRateLimiter } from '../../agentic/DeepCodyRateLimiter'
import type { ChatBuilder } from '../ChatBuilder'
Expand Down Expand Up @@ -34,8 +33,6 @@ export class DeepCodyHandler extends ChatHandler implements AgentHandler {
error?: Error
abort?: boolean
}> {
CodyToolProvider.initialize(this.contextRetriever)

// NOTE: Skip query rewrite for deep-cody as the agent will reviewed and rewrite the query.
const skipQueryRewrite = true
const baseContextResult = await super.computeContext(
Expand Down
Loading

0 comments on commit 2e5d77e

Please sign in to comment.