Skip to content

Commit df3e5de

Browse files
intitniintitni
committed
Add new ActiveDocumentChatContextCollector
1 parent 430d693 commit df3e5de

20 files changed

+1451
-587
lines changed

‎Core/Package.swift‎

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -352,20 +352,25 @@ let package = Package(
352352
],
353353
path: "Sources/ChatContextCollectors/WebChatContextCollector"
354354
),
355-
356-
.target(
357-
name: "ActiveDocumentChatContextCollector",
358-
dependencies: [
359-
"ChatContextCollector",
360-
.product(name: "LangChain", package: "Tool"),
361-
.product(name: "OpenAIService", package: "Tool"),
362-
.product(name: "Preferences", package: "Tool"),
363-
.product(name: "ASTParser", package: "Tool"),
364-
.product(name: "SwiftSyntax", package: "swift-syntax"),
365-
.product(name: "SwiftParser", package: "swift-syntax"),
366-
],
367-
path: "Sources/ChatContextCollectors/ActiveDocumentChatContextCollector"
368-
),
355+
356+
.target(
357+
name: "ActiveDocumentChatContextCollector",
358+
dependencies: [
359+
"ChatContextCollector",
360+
.product(name: "LangChain", package: "Tool"),
361+
.product(name: "OpenAIService", package: "Tool"),
362+
.product(name: "Preferences", package: "Tool"),
363+
.product(name: "ASTParser", package: "Tool"),
364+
.product(name: "SwiftSyntax", package: "swift-syntax"),
365+
.product(name: "SwiftParser", package: "swift-syntax"),
366+
],
367+
path: "Sources/ChatContextCollectors/ActiveDocumentChatContextCollector"
368+
),
369+
370+
.testTarget(
371+
name: "ActiveDocumentChatContextCollectorTests",
372+
dependencies: ["ActiveDocumentChatContextCollector"]
373+
),
369374
]
370375
)
371376

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

Lines changed: 234 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -9,45 +9,252 @@ import XcodeInspector
99
public final class ActiveDocumentChatContextCollector: ChatContextCollector {
1010
public init() {}
1111

12+
var activeDocumentContext: ActiveDocumentContext?
13+
1214
public func generateContext(
1315
history: [ChatMessage],
1416
scopes: Set<String>,
1517
content: String
1618
) -> ChatContext? {
17-
guard scopes.contains("file") else { return nil }
18-
let info = getEditorInformation()
19+
guard let info = getEditorInformation() else { return nil }
20+
let context = getActiveDocumentContext(info)
21+
activeDocumentContext = context
22+
23+
guard scopes.contains("code") || scopes.contains("c") else {
24+
if scopes.contains("file") || scopes.contains("f") {
25+
var removedCode = context
26+
removedCode.focusedContext = nil
27+
return .init(
28+
systemPrompt: extractSystemPrompt(removedCode),
29+
functions: []
30+
)
31+
}
32+
return nil
33+
}
34+
35+
var functions = [any ChatGPTFunction]()
36+
37+
// When the bot is already focusing on a piece of code, it can expand the range.
38+
39+
if context.focusedContext != nil {
40+
functions.append(ExpandFocusRangeFunction(contextCollector: self))
41+
}
42+
43+
// When the bot is not focusing on any code, or the focusing area is not the user's
44+
// selection, it can move the focus back to the user's selection.
45+
46+
if context.focusedContext == nil ||
47+
!(context.focusedContext?.codeRange.contains(context.selectionRange) ?? false)
48+
{
49+
functions.append(MoveToFocusedCodeFunction(contextCollector: self))
50+
}
51+
52+
// When there is a line annotation not in the focused area, the bot can move the focus area
53+
// to the code covering the line of the annotation.
54+
55+
if let focusedContext = context.focusedContext,
56+
!focusedContext.otherLineAnnotations.isEmpty
57+
{
58+
functions.append(MoveToCodeAroundLineFunction(contextCollector: self))
59+
}
60+
61+
if context.focusedContext == nil, !context.lineAnnotations.isEmpty {
62+
functions.append(MoveToCodeAroundLineFunction(contextCollector: self))
63+
}
1964

2065
return .init(
21-
systemPrompt: extractSystemPrompt(info),
22-
functions: []
66+
systemPrompt: extractSystemPrompt(context),
67+
functions: functions
2368
)
2469
}
25-
26-
func extractSystemPrompt(_ info: EditorInformation) -> String {
27-
let relativePath = info.documentURL.path
28-
.replacingOccurrences(of: info.projectURL.path, with: "")
29-
let selectionRange = info.editorContent?.selections.first ?? .outOfScope
30-
let lineAnnotations = info.editorContent?.lineAnnotations ?? []
31-
32-
var result = """
33-
Active Document Context:###
34-
Document Relative Path: \(relativePath)
35-
Language: \(info.language.rawValue)
36-
Selection Range [line, character]: \
37-
[\(selectionRange.start.line), \(selectionRange.start.character)] - \
38-
[\(selectionRange.end.line), \(selectionRange.end.character)]
39-
###
40-
"""
41-
42-
if !lineAnnotations.isEmpty {
43-
result += """
44-
Line Annotations:
45-
\(lineAnnotations.map { " - \($0)" }.joined(separator: "\n"))
70+
71+
func getActiveDocumentContext(_ info: EditorInformation) -> ActiveDocumentContext {
72+
var activeDocumentContext = activeDocumentContext ?? .init(
73+
relativePath: "",
74+
language: .builtIn(.swift),
75+
fileContent: "",
76+
lines: [],
77+
selectedCode: "",
78+
selectionRange: .outOfScope,
79+
lineAnnotations: [],
80+
imports: []
81+
)
82+
83+
activeDocumentContext.update(info)
84+
return activeDocumentContext
85+
}
86+
87+
func extractSystemPrompt(_ context: ActiveDocumentContext) -> String {
88+
let start = "User Editing Document Context:###"
89+
let end = "###"
90+
let relativePath = "Document Relative Path: \(context.relativePath)"
91+
let language = "Language: \(context.language)"
92+
93+
if let focusedContext = context.focusedContext {
94+
let codeContext = "\(focusedContext.contextRange) \(focusedContext.context)"
95+
let codeRange = "Focused Range [line, character]: \(focusedContext.codeRange)"
96+
let code = """
97+
Focused Code (start from line \(focusedContext.codeRange.start.line)):
98+
```\(context.language.rawValue)
99+
\(focusedContext.code)
100+
```
46101
"""
102+
let fileAnnotations = focusedContext.otherLineAnnotations.isEmpty
103+
? ""
104+
: """
105+
File Annotations:
106+
\(focusedContext.otherLineAnnotations.map { " - \($0)" }.joined(separator: "\n"))
107+
"""
108+
let codeAnnotations = focusedContext.lineAnnotations.isEmpty
109+
? ""
110+
: """
111+
Code Annotations:
112+
\(focusedContext.lineAnnotations.map { " - \($0)" }.joined(separator: "\n"))
113+
"""
114+
return [
115+
start,
116+
relativePath,
117+
language,
118+
fileAnnotations,
119+
codeContext,
120+
codeRange,
121+
codeAnnotations,
122+
code,
123+
end,
124+
]
125+
.filter { !$0.isEmpty }
126+
.joined(separator: "\n")
127+
} else {
128+
let selectionRange = "Selection Range [line, character]: \(context.selectionRange)"
129+
let lineAnnotations = context.lineAnnotations.isEmpty
130+
? ""
131+
: """
132+
Line Annotations:
133+
\(context.lineAnnotations.map { " - \($0)" }.joined(separator: "\n"))
134+
"""
135+
136+
return [
137+
start,
138+
relativePath,
139+
language,
140+
lineAnnotations,
141+
selectionRange,
142+
end,
143+
]
144+
.filter { !$0.isEmpty }
145+
.joined(separator: "\n")
47146
}
48-
49-
return result
50147
}
51148
}
52149

150+
struct ActiveDocumentContext {
151+
var relativePath: String
152+
var language: CodeLanguage
153+
var fileContent: String
154+
var lines: [String]
155+
var selectedCode: String
156+
var selectionRange: CursorRange
157+
var lineAnnotations: [EditorInformation.LineAnnotation]
158+
var imports: [String]
159+
160+
struct FocusedContext {
161+
var context: String
162+
var contextRange: CursorRange
163+
var codeRange: CursorRange
164+
var code: String
165+
var lineAnnotations: [EditorInformation.LineAnnotation]
166+
var otherLineAnnotations: [EditorInformation.LineAnnotation]
167+
}
168+
169+
var focusedContext: FocusedContext?
170+
171+
mutating func moveToFocusedCode() {
172+
moveToCodeContainingRange(selectionRange)
173+
}
174+
175+
mutating func moveToCodeAroundLine(_ line: Int) {
176+
moveToCodeContainingRange(.init(
177+
start: .init(line: line, character: 0),
178+
end: .init(line: line, character: 0)
179+
))
180+
}
181+
182+
mutating func expandFocusedRangeToContextRange() {
183+
guard let focusedContext else { return }
184+
moveToCodeContainingRange(focusedContext.contextRange)
185+
}
186+
187+
mutating func moveToCodeContainingRange(_ range: CursorRange) {
188+
let finder: FocusedCodeFinder = {
189+
switch language {
190+
case .builtIn(.swift):
191+
return SwiftFocusedCodeFinder()
192+
default:
193+
return UnknownLanguageFocusedCodeFinder()
194+
}
195+
}()
196+
197+
let codeContext = finder.findFocusedCode(
198+
containingRange: range,
199+
activeDocumentContext: self
200+
)
201+
202+
imports = codeContext.imports
203+
204+
let startLine = codeContext.focusedRange.start.line
205+
let endLine = codeContext.focusedRange.end.line
206+
var matchedAnnotations = [EditorInformation.LineAnnotation]()
207+
var otherAnnotations = [EditorInformation.LineAnnotation]()
208+
for annotation in lineAnnotations {
209+
if annotation.line >= startLine && annotation.line <= endLine {
210+
matchedAnnotations.append(annotation)
211+
} else {
212+
otherAnnotations.append(annotation)
213+
}
214+
}
215+
216+
focusedContext = .init(
217+
context: {
218+
switch codeContext.scope {
219+
case .file:
220+
return "File"
221+
case .top:
222+
return "Top level of the file"
223+
case let .scope(signature):
224+
return signature
225+
}
226+
}(),
227+
contextRange: codeContext.contextRange,
228+
codeRange: codeContext.focusedRange,
229+
code: codeContext.focusedCode,
230+
lineAnnotations: matchedAnnotations,
231+
otherLineAnnotations: otherAnnotations
232+
)
233+
}
234+
235+
mutating func update(_ info: EditorInformation) {
236+
/// Whenever the file content, relative path, or selection range changes,
237+
/// we should reset the context.
238+
let changed: Bool = {
239+
if info.relativePath != relativePath { return true }
240+
if info.editorContent?.content != fileContent { return true }
241+
if let range = info.editorContent?.selections.first,
242+
range != selectionRange { return true }
243+
return false
244+
}()
245+
246+
relativePath = info.relativePath
247+
language = info.language
248+
fileContent = info.editorContent?.content ?? ""
249+
lines = info.editorContent?.lines ?? []
250+
selectedCode = info.selectedContent
251+
selectionRange = info.editorContent?.selections.first ?? .zero
252+
lineAnnotations = info.editorContent?.lineAnnotations ?? []
253+
imports = []
254+
255+
if changed {
256+
moveToFocusedCode()
257+
}
258+
}
259+
}
53260

0 commit comments

Comments
 (0)