Skip to content

Commit 3d8544c

Browse files
charlesliencharleslien
committed
move request-files-prompt to agent-runtime
1 parent 8c070b6 commit 3d8544c

File tree

6 files changed

+30
-125
lines changed

6 files changed

+30
-125
lines changed

‎backend/src/__tests__/main-prompt.test.ts‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@ import {
2525

2626
// Mock imports
2727
import * as checkTerminalCommandModule from '../check-terminal-command'
28-
import * as requestFilesPrompt from '../find-files/request-files-prompt'
2928
import * as getDocumentationForQueryModule from '../get-documentation-for-query'
3029
import { mainPrompt } from '../main-prompt'
3130
import * as processFileBlockModule from '../process-file-block'
@@ -154,10 +153,6 @@ describe('mainPrompt', () => {
154153
}),
155154
)
156155

157-
spyOn(requestFilesPrompt, 'requestRelevantFiles').mockImplementation(
158-
async () => [],
159-
)
160-
161156
spyOn(
162157
checkTerminalCommandModule,
163158
'checkTerminalCommand',

‎backend/src/__tests__/read-docs-tool.test.ts‎

Lines changed: 1 addition & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,8 @@ import {
2020
test,
2121
} from 'bun:test'
2222

23-
import researcherAgent from '../../../.agents/researcher/researcher'
24-
import * as checkTerminalCommandModule from '../check-terminal-command'
2523
import { mockFileContext } from './test-utils'
26-
import * as requestFilesPrompt from '../find-files/request-files-prompt'
24+
import researcherAgent from '../../../.agents/researcher/researcher'
2725
import * as context7Api from '../llm-apis/context7-api'
2826
import { runAgentStep } from '../run-agent-step'
2927

@@ -106,25 +104,6 @@ describe('read_docs tool with researcher agent', () => {
106104
},
107105
],
108106
})
109-
110-
// Mock other required modules
111-
const requestRelevantFilesSpy = spyOn(
112-
requestFilesPrompt,
113-
'requestRelevantFiles',
114-
).mockImplementation(async () => [])
115-
mockedFunctions.push({
116-
name: 'requestFilesPrompt.requestRelevantFiles',
117-
spy: requestRelevantFilesSpy,
118-
})
119-
120-
const checkTerminalCommandSpy = spyOn(
121-
checkTerminalCommandModule,
122-
'checkTerminalCommand',
123-
).mockImplementation(async () => null)
124-
mockedFunctions.push({
125-
name: 'checkTerminalCommand',
126-
spy: checkTerminalCommandSpy,
127-
})
128107
})
129108

130109
afterEach(() => {

‎backend/src/__tests__/web-search-tool.test.ts‎

Lines changed: 1 addition & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,8 @@ import {
2323
test,
2424
} from 'bun:test'
2525

26-
import researcherAgent from '../../../.agents/researcher/researcher'
27-
import * as checkTerminalCommandModule from '../check-terminal-command'
2826
import { mockFileContext } from './test-utils'
29-
import * as requestFilesPrompt from '../find-files/request-files-prompt'
27+
import researcherAgent from '../../../.agents/researcher/researcher'
3028
import * as linkupApi from '../llm-apis/linkup-api'
3129
import { runAgentStep } from '../run-agent-step'
3230

@@ -82,15 +80,6 @@ describe('web_search tool with researcher agent', () => {
8280
agentRuntimeImpl.promptAiSdk = async function () {
8381
return 'Test response'
8482
}
85-
86-
// Mock other required modules
87-
spyOn(requestFilesPrompt, 'requestRelevantFiles').mockImplementation(
88-
async () => [],
89-
)
90-
spyOn(
91-
checkTerminalCommandModule,
92-
'checkTerminalCommand',
93-
).mockImplementation(async () => null)
9483
})
9584

9685
afterEach(() => {

‎backend/src/tools/handlers/tool/find-files.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
1+
import {
2+
requestRelevantFiles,
3+
requestRelevantFilesForTraining,
4+
} from '@codebuff/agent-runtime/find-files/request-files-prompt'
15
import {
26
countTokens,
37
countTokensJson,
48
} from '@codebuff/agent-runtime/util/token-counter'
59
import { insertTrace } from '@codebuff/bigquery'
610

7-
import {
8-
requestRelevantFiles,
9-
requestRelevantFilesForTraining,
10-
} from '../../../find-files/request-files-prompt'
1111
import { getFileReadingUpdates } from '../../../get-file-reading-updates'
1212
import { getSearchSystemPrompt } from '../../../system-prompt/search-system-prompt'
1313
import { renderReadFilesResult } from '../../../util/parse-tool-call-xml'

backend/src/__tests__/request-files-prompt.test.ts renamed to packages/agent-runtime/src/find-files/__tests__/request-files-prompt.test.ts

Lines changed: 15 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,15 @@
1-
import * as geminiWithFallbacksModule from '@codebuff/agent-runtime/llm-api/gemini-with-fallbacks'
21
import { finetunedVertexModels } from '@codebuff/common/old-constants'
32
import {
43
TEST_AGENT_RUNTIME_IMPL,
54
TEST_AGENT_RUNTIME_SCOPED_IMPL,
65
} from '@codebuff/common/testing/impl/agent-runtime'
7-
import {
8-
clearMockedModules,
9-
mockModule,
10-
} from '@codebuff/common/testing/mock-modules'
11-
import {
12-
afterAll,
13-
beforeAll,
14-
beforeEach,
15-
mock as bunMockFn,
16-
describe,
17-
expect,
18-
it,
19-
} from 'bun:test'
6+
import { beforeEach, describe, expect, it, mock } from 'bun:test'
207

21-
import * as OriginalRequestFilesPromptModule from '../find-files/request-files-prompt'
8+
import { requestRelevantFiles } from '../request-files-prompt'
229

23-
import type { CostMode } from '@codebuff/common/old-constants'
2410
import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
2511
import type { Message } from '@codebuff/common/types/messages/codebuff-message'
2612
import type { ProjectFileContext } from '@codebuff/common/util/file'
27-
import type { Mock } from 'bun:test'
2813

2914
let agentRuntimeImpl: AgentRuntimeDeps
3015

@@ -62,54 +47,17 @@ describe('requestRelevantFiles', () => {
6247
const mockFingerprintId = 'fingerprint1'
6348
const mockUserInputId = 'input1'
6449
const mockUserId = 'user1'
65-
const mockCostMode: CostMode = 'normal'
6650
const mockRepoId = 'owner/repo'
6751

68-
beforeAll(() => {
69-
mockModule('@codebuff/backend/llm-apis/gemini-with-fallbacks', () => ({
70-
promptFlashWithFallbacks: bunMockFn(() =>
71-
Promise.resolve('file1.ts\nfile2.ts'),
72-
),
73-
}))
74-
75-
mockModule('@codebuff/backend/websockets/request-context', () => ({
76-
getRequestContext: bunMockFn(() => ({
77-
approvedOrgIdForRepo: 'org123',
78-
isRepoApprovedForUserInOrg: true,
79-
})),
80-
}))
81-
82-
mockModule('@codebuff/common/db', () => ({
83-
default: {
84-
insert: bunMockFn(() => ({
85-
values: bunMockFn(() => ({
86-
onConflictDoNothing: bunMockFn(() => Promise.resolve()),
87-
})),
88-
})),
89-
},
90-
}))
91-
mockModule('@codebuff/bigquery', () => ({
92-
insertTrace: bunMockFn(() => Promise.resolve()),
93-
}))
94-
})
95-
96-
afterAll(() => {
97-
clearMockedModules()
98-
})
99-
10052
beforeEach(() => {
101-
agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
102-
103-
const promptFlashWithFallbacksMock =
104-
geminiWithFallbacksModule.promptFlashWithFallbacks as Mock<
105-
typeof geminiWithFallbacksModule.promptFlashWithFallbacks
106-
>
107-
promptFlashWithFallbacksMock.mockResolvedValue('file1.ts\nfile2.ts')
108-
promptFlashWithFallbacksMock.mockClear()
53+
agentRuntimeImpl = {
54+
...TEST_AGENT_RUNTIME_IMPL,
55+
promptAiSdk: mock(() => Promise.resolve('file1.ts\nfile2.ts')),
56+
}
10957
})
11058

11159
it('should use default file counts and maxFiles when no custom config', async () => {
112-
await OriginalRequestFilesPromptModule.requestRelevantFiles({
60+
await requestRelevantFiles({
11361
...agentRuntimeImpl,
11462
...TEST_AGENT_RUNTIME_SCOPED_IMPL,
11563
messages: mockMessages,
@@ -123,9 +71,7 @@ describe('requestRelevantFiles', () => {
12371
userId: mockUserId,
12472
repoId: mockRepoId,
12573
})
126-
expect(
127-
geminiWithFallbacksModule.promptFlashWithFallbacks,
128-
).toHaveBeenCalled()
74+
expect(agentRuntimeImpl.promptAiSdk).toHaveBeenCalled()
12975
})
13076

13177
it('should use custom file counts from config', async () => {
@@ -135,7 +81,7 @@ describe('requestRelevantFiles', () => {
13581
maxFilesPerRequest: 10,
13682
}
13783

138-
await OriginalRequestFilesPromptModule.requestRelevantFiles({
84+
await requestRelevantFiles({
13985
...agentRuntimeImpl,
14086
...TEST_AGENT_RUNTIME_SCOPED_IMPL,
14187
messages: mockMessages,
@@ -149,9 +95,7 @@ describe('requestRelevantFiles', () => {
14995
userId: mockUserId,
15096
repoId: mockRepoId,
15197
})
152-
expect(
153-
geminiWithFallbacksModule.promptFlashWithFallbacks,
154-
).toHaveBeenCalled()
98+
expect(agentRuntimeImpl.promptAiSdk).toHaveBeenCalled()
15599
})
156100

157101
it('should use custom maxFilesPerRequest from config', async () => {
@@ -160,7 +104,7 @@ describe('requestRelevantFiles', () => {
160104
maxFilesPerRequest: 3,
161105
}
162106

163-
const result = await OriginalRequestFilesPromptModule.requestRelevantFiles({
107+
const result = await requestRelevantFiles({
164108
...agentRuntimeImpl,
165109
...TEST_AGENT_RUNTIME_SCOPED_IMPL,
166110
messages: mockMessages,
@@ -185,7 +129,7 @@ describe('requestRelevantFiles', () => {
185129
modelName: 'ft_filepicker_010',
186130
}
187131

188-
await OriginalRequestFilesPromptModule.requestRelevantFiles({
132+
await requestRelevantFiles({
189133
...agentRuntimeImpl,
190134
...TEST_AGENT_RUNTIME_SCOPED_IMPL,
191135
messages: mockMessages,
@@ -199,9 +143,7 @@ describe('requestRelevantFiles', () => {
199143
userId: mockUserId,
200144
repoId: mockRepoId,
201145
})
202-
expect(
203-
geminiWithFallbacksModule.promptFlashWithFallbacks,
204-
).toHaveBeenCalledWith(
146+
expect(agentRuntimeImpl.promptAiSdk).toHaveBeenCalledWith(
205147
expect.objectContaining({
206148
useFinetunedModel: finetunedVertexModels.ft_filepicker_010,
207149
}),
@@ -213,7 +155,7 @@ describe('requestRelevantFiles', () => {
213155
modelName: 'invalid-model-name',
214156
}
215157

216-
await OriginalRequestFilesPromptModule.requestRelevantFiles({
158+
await requestRelevantFiles({
217159
...agentRuntimeImpl,
218160
...TEST_AGENT_RUNTIME_SCOPED_IMPL,
219161
messages: mockMessages,
@@ -228,9 +170,7 @@ describe('requestRelevantFiles', () => {
228170
repoId: mockRepoId,
229171
})
230172
const expectedModel = finetunedVertexModels.ft_filepicker_010
231-
expect(
232-
geminiWithFallbacksModule.promptFlashWithFallbacks,
233-
).toHaveBeenCalledWith(
173+
expect(agentRuntimeImpl.promptAiSdk).toHaveBeenCalledWith(
234174
expect.objectContaining({
235175
useFinetunedModel: expectedModel,
236176
}),

backend/src/find-files/request-files-prompt.ts renamed to packages/agent-runtime/src/find-files/request-files-prompt.ts

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,5 @@
11
import { dirname, isAbsolute, normalize } from 'path'
22

3-
import { promptFlashWithFallbacks } from '@codebuff/agent-runtime/llm-api/gemini-with-fallbacks'
4-
import {
5-
castAssistantMessage,
6-
messagesWithSystem,
7-
getMessagesSubset,
8-
} from '@codebuff/agent-runtime/util/messages'
93
import { insertTrace } from '@codebuff/bigquery'
104
import {
115
finetunedVertexModels,
@@ -15,7 +9,15 @@ import {
159
import { getAllFilePaths } from '@codebuff/common/project-file-tree'
1610
import { range, shuffle, uniq } from 'lodash'
1711

18-
import type { TextBlock } from '@codebuff/agent-runtime/llm-api/claude'
12+
import {
13+
castAssistantMessage,
14+
messagesWithSystem,
15+
getMessagesSubset,
16+
} from '@codebuff/agent-runtime/util/messages'
17+
18+
import { promptFlashWithFallbacks } from '../llm-api/gemini-with-fallbacks'
19+
20+
import type { TextBlock } from '../llm-api/claude'
1921
import type {
2022
GetExpandedFileContextForTrainingTrace,
2123
GetRelevantFilesTrace,

0 commit comments

Comments
 (0)