Skip to content

Commit 07630c4

Browse files
intitniintitni
committed
Fix that system prompt is not updated on function call
1 parent a0885cc commit 07630c4

File tree

12 files changed

+186
-83
lines changed

12 files changed

+186
-83
lines changed

‎Core/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/ActiveDocumentChatContextCollector.swift‎

Lines changed: 55 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,8 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
6262
functions.append(MoveToCodeAroundLineFunction(contextCollector: self))
6363
}
6464

65+
print(extractSystemPrompt(context))
66+
6567
return .init(
6668
systemPrompt: extractSystemPrompt(context),
6769
functions: functions
@@ -87,55 +89,84 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
8789

8890
func extractSystemPrompt(_ context: ActiveDocumentContext) -> String {
8991
let start = """
90-
User Editing Document Context:###
91-
(The context may change during the conversation, so it may not match our conversation.)
92+
## File and Code Scope
93+
94+
You can use the following context to answer user's questions about the editing document or code. The context shows only a part of the code in the editing document, and will change during the conversation, so it may not match our conversation.
95+
96+
User Editing Document Context: ###
9297
"""
9398
let end = "###"
9499
let relativePath = "Document Relative Path: \(context.relativePath)"
95100
let language = "Language: \(context.language.rawValue)"
96101

97102
if let focusedContext = context.focusedContext {
98-
let codeContext = "Focused Context:\(focusedContext.contextRange) \(focusedContext.context)"
103+
let codeContext = focusedContext.context.isEmpty
104+
? ""
105+
: """
106+
Focused Context:
107+
```
108+
\(focusedContext.context.joined(separator: "\n"))
109+
```
110+
"""
111+
99112
let codeRange = "Focused Range [line, character]: \(focusedContext.codeRange)"
113+
100114
let code = """
101-
Focused Code (start from line \(focusedContext.codeRange.start.line), call `expandFocusRange` if you need more context):
115+
Focused Code (start from line \(
116+
focusedContext.codeRange.start
117+
.line
118+
)):
102119
```\(context.language.rawValue)
103120
\(focusedContext.code)
104121
```
105122
"""
123+
106124
let fileAnnotations = focusedContext.otherLineAnnotations.isEmpty
107125
? ""
108126
: """
109-
File Annotations:
110-
\(focusedContext.otherLineAnnotations.map { " - \($0)" }.joined(separator: "\n"))
111-
(call `moveToCodeAroundLine` if you context about the line annotations here)
127+
Other Annotations:\"""
128+
(They are not inside the focused code. You don't known how to handle them until you get the code at the line)
129+
\(
130+
focusedContext.otherLineAnnotations
131+
.map(convertAnnotationToText)
132+
.joined(separator: "\n")
133+
)
134+
\"""
112135
"""
136+
113137
let codeAnnotations = focusedContext.lineAnnotations.isEmpty
114138
? ""
115139
: """
116-
Code Annotations:
117-
\(focusedContext.lineAnnotations.map { " - \($0)" }.joined(separator: "\n"))
140+
Annotations Inside Focused Range:\"""
141+
\(
142+
focusedContext.lineAnnotations
143+
.map(convertAnnotationToText)
144+
.joined(separator: "\n")
145+
)
146+
\"""
118147
"""
148+
119149
return [
120150
start,
121151
relativePath,
122152
language,
123-
fileAnnotations,
124153
codeContext,
125154
codeRange,
126-
codeAnnotations,
127155
code,
156+
codeAnnotations,
157+
fileAnnotations,
128158
end,
129159
]
130160
.filter { !$0.isEmpty }
131-
.joined(separator: "\n")
161+
.joined(separator: "\n\n")
132162
} else {
133163
let selectionRange = "Selection Range [line, character]: \(context.selectionRange)"
134164
let lineAnnotations = context.lineAnnotations.isEmpty
135165
? ""
136166
: """
137-
Line Annotations:
138-
\(context.lineAnnotations.map { " - \($0)" }.joined(separator: "\n"))
167+
Line Annotations:\"""
168+
\(context.lineAnnotations.map(convertAnnotationToText).joined(separator: "\n"))
169+
\"""
139170
"""
140171

141172
return [
@@ -150,6 +181,10 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
150181
.joined(separator: "\n")
151182
}
152183
}
184+
185+
func convertAnnotationToText(_ annotation: EditorInformation.LineAnnotation) -> String {
186+
return "- Line \(annotation.line), \(annotation.type): \(annotation.message)"
187+
}
153188
}
154189

155190
struct ActiveDocumentContext {
@@ -164,7 +199,7 @@ struct ActiveDocumentContext {
164199
var imports: [String]
165200

166201
struct FocusedContext {
167-
var context: String
202+
var context: [String]
168203
var contextRange: CursorRange
169204
var codeRange: CursorRange
170205
var code: String
@@ -206,30 +241,21 @@ struct ActiveDocumentContext {
206241
)
207242

208243
imports = codeContext.imports
209-
244+
210245
let startLine = codeContext.focusedRange.start.line
211246
let endLine = codeContext.focusedRange.end.line
212247
var matchedAnnotations = [EditorInformation.LineAnnotation]()
213248
var otherAnnotations = [EditorInformation.LineAnnotation]()
214249
for annotation in lineAnnotations {
215-
if annotation.line >= startLine && annotation.line <= endLine {
250+
if annotation.line >= startLine, annotation.line <= endLine {
216251
matchedAnnotations.append(annotation)
217252
} else {
218253
otherAnnotations.append(annotation)
219254
}
220255
}
221-
256+
222257
focusedContext = .init(
223-
context: {
224-
switch codeContext.scope {
225-
case .file:
226-
return "File"
227-
case .top:
228-
return "Top level of the file"
229-
case let .scope(signature):
230-
return signature
231-
}
232-
}(),
258+
context: codeContext.scopeSignatures,
233259
contextRange: codeContext.contextRange,
234260
codeRange: codeContext.focusedRange,
235261
code: codeContext.focusedCode,
@@ -258,7 +284,7 @@ struct ActiveDocumentContext {
258284
selectionRange = info.editorContent?.selections.first ?? .zero
259285
lineAnnotations = info.editorContent?.lineAnnotations ?? []
260286
imports = []
261-
287+
262288
if changed {
263289
moveToFocusedCode()
264290
}

‎Core/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/FocusedCodeFinder/FocusedCodeFinder.swift‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,19 @@ struct CodeContext: Equatable {
55
enum Scope: Equatable {
66
case file
77
case top
8-
case scope(signature: String)
8+
case scope(signature: [String])
99
}
1010

11+
var scopeSignatures: [String] {
12+
switch scope {
13+
case .file:
14+
return []
15+
case .top:
16+
return ["Top level of the file"]
17+
case .scope(let signature):
18+
return signature
19+
}
20+
}
1121
var scope: Scope
1222
var contextRange: CursorRange
1323
var focusedRange: CursorRange

‎Core/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/FocusedCodeFinder/SwiftFocusedCodeFinder.swift‎

Lines changed: 32 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ struct SwiftFocusedCodeFinder: FocusedCodeFinder {
1010
activeDocumentContext: ActiveDocumentContext
1111
) -> CodeContext {
1212
let source = activeDocumentContext.fileContent
13+
#warning("TODO: cache the tree")
1314
let tree = Parser.parse(source: source)
1415

1516
let locationConverter = SourceLocationConverter(
@@ -58,7 +59,6 @@ struct SwiftFocusedCodeFinder: FocusedCodeFinder {
5859
result.imports = visitor.imports
5960
return result
6061
}
61-
nodes.removeFirst()
6262
codeRange = convertRange(focusedNode)
6363
} else {
6464
codeRange = range
@@ -91,8 +91,7 @@ struct SwiftFocusedCodeFinder: FocusedCodeFinder {
9191
}
9292

9393
return .init(
94-
scope: signature
95-
.isEmpty ? .file : .scope(signature: signature.joined(separator: " > ")),
94+
scope: signature.isEmpty ? .file : .scope(signature: signature),
9695
contextRange: contextRange,
9796
focusedRange: codeRange,
9897
focusedCode: code,
@@ -130,7 +129,8 @@ extension SwiftFocusedCodeFinder {
130129
return (.init(
131130
signature: "\(type) \(name)"
132131
.prefixedModifiers(node.modifierAndAttributeText(extractText))
133-
.suffixedInheritance(node.inheritanceClauseTexts(extractText)),
132+
.suffixedInheritance(node.inheritanceClauseTexts(extractText))
133+
.replacingOccurrences(of: "\n", with: " "),
134134
contextRange: convertRange(node)
135135
), false)
136136

@@ -140,7 +140,8 @@ extension SwiftFocusedCodeFinder {
140140
return (.init(
141141
signature: "\(type) \(name)"
142142
.prefixedModifiers(node.modifierAndAttributeText(extractText))
143-
.suffixedInheritance(node.inheritanceClauseTexts(extractText)),
143+
.suffixedInheritance(node.inheritanceClauseTexts(extractText))
144+
.replacingOccurrences(of: "\n", with: " "),
144145
contextRange: convertRange(node)
145146
), false)
146147

@@ -150,7 +151,8 @@ extension SwiftFocusedCodeFinder {
150151
return (.init(
151152
signature: "\(type) \(name)"
152153
.prefixedModifiers(node.modifierAndAttributeText(extractText))
153-
.suffixedInheritance(node.inheritanceClauseTexts(extractText)),
154+
.suffixedInheritance(node.inheritanceClauseTexts(extractText))
155+
.replacingOccurrences(of: "\n", with: " "),
154156
contextRange: convertRange(node)
155157
), false)
156158

@@ -160,7 +162,8 @@ extension SwiftFocusedCodeFinder {
160162
return (.init(
161163
signature: "\(type) \(name)"
162164
.prefixedModifiers(node.modifierAndAttributeText(extractText))
163-
.suffixedInheritance(node.inheritanceClauseTexts(extractText)),
165+
.suffixedInheritance(node.inheritanceClauseTexts(extractText))
166+
.replacingOccurrences(of: "\n", with: ""),
164167
contextRange: convertRange(node)
165168
), false)
166169

@@ -169,7 +172,8 @@ extension SwiftFocusedCodeFinder {
169172
let name = node.identifier.text
170173
return (.init(
171174
signature: "\(type) \(name)"
172-
.prefixedModifiers(node.modifierAndAttributeText(extractText)),
175+
.prefixedModifiers(node.modifierAndAttributeText(extractText))
176+
.replacingOccurrences(of: "\n", with: " "),
173177
contextRange: convertRange(node)
174178
), false)
175179

@@ -179,7 +183,8 @@ extension SwiftFocusedCodeFinder {
179183
return (.init(
180184
signature: "\(type) \(name)"
181185
.prefixedModifiers(node.modifierAndAttributeText(extractText))
182-
.suffixedInheritance(node.inheritanceClauseTexts(extractText)),
186+
.suffixedInheritance(node.inheritanceClauseTexts(extractText))
187+
.replacingOccurrences(of: "\n", with: " "),
183188
contextRange: convertRange(node)
184189
), false)
185190

@@ -189,14 +194,18 @@ extension SwiftFocusedCodeFinder {
189194
return (.init(
190195
signature: "\(type) \(name)"
191196
.prefixedModifiers(node.modifierAndAttributeText(extractText))
192-
.suffixedInheritance(node.inheritanceClauseTexts(extractText)),
197+
.suffixedInheritance(node.inheritanceClauseTexts(extractText))
198+
.replacingOccurrences(of: "\n", with: " "),
193199
contextRange: convertRange(node)
194200
), false)
195201

196202
case let node as FunctionDeclSyntax:
197203
let type = node.funcKeyword.text
198204
let name = node.identifier.text
199205
let signature = node.signature.trimmedDescription
206+
.split(separator: "\n")
207+
.map { $0.trimmingCharacters(in: .whitespacesAndNewlines) }
208+
.joined(separator: " ")
200209

201210
return (.init(
202211
signature: "\(type) \(name)\(signature)"
@@ -211,7 +220,8 @@ extension SwiftFocusedCodeFinder {
211220

212221
return (.init(
213222
signature: "\(type) \(name)\(signature.isEmpty ? "" : ": \(signature)")"
214-
.prefixedModifiers(node.modifierAndAttributeText(extractText)),
223+
.prefixedModifiers(node.modifierAndAttributeText(extractText))
224+
.replacingOccurrences(of: "\n", with: " "),
215225
contextRange: convertRange(node)
216226
), true)
217227

@@ -221,7 +231,8 @@ extension SwiftFocusedCodeFinder {
221231

222232
return (.init(
223233
signature: signature
224-
.prefixedModifiers(node.modifierAndAttributeText(extractText)),
234+
.prefixedModifiers(node.modifierAndAttributeText(extractText))
235+
.replacingOccurrences(of: "\n", with: " "),
225236
contextRange: convertRange(node)
226237
), true)
227238

@@ -233,7 +244,8 @@ extension SwiftFocusedCodeFinder {
233244

234245
return (.init(
235246
signature: signature
236-
.prefixedModifiers(node.modifierAndAttributeText(extractText)),
247+
.prefixedModifiers(node.modifierAndAttributeText(extractText))
248+
.replacingOccurrences(of: "\n", with: " "),
237249
contextRange: convertRange(node)
238250
), true)
239251

@@ -242,7 +254,8 @@ extension SwiftFocusedCodeFinder {
242254

243255
return (.init(
244256
signature: "\(signature)"
245-
.prefixedModifiers(node.modifierAndAttributeText(extractText)),
257+
.prefixedModifiers(node.modifierAndAttributeText(extractText))
258+
.replacingOccurrences(of: "\n", with: " "),
246259

247260
contextRange: convertRange(node)
248261
), true)
@@ -252,7 +265,8 @@ extension SwiftFocusedCodeFinder {
252265

253266
return (.init(
254267
signature: signature
255-
.prefixedModifiers(node.modifierAndAttributeText(extractText)),
268+
.prefixedModifiers(node.modifierAndAttributeText(extractText))
269+
.replacingOccurrences(of: "\n", with: " "),
256270

257271
contextRange: convertRange(node)
258272
), true)
@@ -261,22 +275,22 @@ extension SwiftFocusedCodeFinder {
261275
let signature = "closure"
262276

263277
return (.init(
264-
signature: signature,
278+
signature: signature.replacingOccurrences(of: "\n", with: " "),
265279
contextRange: convertRange(node)
266280
), true)
267281

268282
case let node as FunctionCallExprSyntax:
269283
let signature = "function call"
270284

271285
return (.init(
272-
signature: signature,
286+
signature: signature.replacingOccurrences(of: "\n", with: " "),
273287
contextRange: convertRange(node),
274288
canBeUsedAsCodeRange: false
275289
), true)
276290

277291
case let node as SwitchCaseSyntax:
278292
return (.init(
279-
signature: node.trimmedDescription,
293+
signature: node.trimmedDescription.replacingOccurrences(of: "\n", with: " "),
280294
contextRange: convertRange(node)
281295
), true)
282296

0 commit comments

Comments
 (0)