Skip to content

Commit 893f2df

Browse files
tmattmat
authored
Implement support for Semantic Search agent and tool (#77784)
1 parent fae3e15 commit 893f2df

23 files changed

+561
-193
lines changed

‎src/Features/CSharpTest/SemanticSearch/Mocks/MockSemanticSearchResultsObserver.cs‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -43,11 +43,5 @@ public ValueTask OnUserCodeExceptionAsync(UserCodeExceptionInfo exception, Cance
4343
OnUserCodeExceptionImpl?.Invoke(exception);
4444
return ValueTaskFactory.CompletedTask;
4545
}
46-
47-
public ValueTask OnCompilationFailureAsync(ImmutableArray<QueryCompilationError> errors, CancellationToken cancellationToken)
48-
{
49-
OnCompilationFailureImpl?.Invoke(errors);
50-
return ValueTaskFactory.CompletedTask;
51-
}
5246
}
5347

‎src/Features/Core/Portable/SemanticSearch/AbstractSemanticSearchService.cs‎

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -120,13 +120,11 @@ public async Task<ExecuteQueryResult> ExecuteQueryAsync(
120120
}
121121
}
122122

123-
await observer.OnCompilationFailureAsync(
124-
emitResult.Diagnostics.SelectAsArray(
123+
var errors = emitResult.Diagnostics.SelectAsArray(
125124
d => d.Severity == DiagnosticSeverity.Error,
126-
d => new QueryCompilationError(d.Id, d.GetMessage(), (d.Location.SourceTree == queryTree) ? d.Location.SourceSpan : default)),
127-
cancellationToken).ConfigureAwait(false);
125+
d => new QueryCompilationError(d.Id, d.GetMessage(), (d.Location.SourceTree == queryTree) ? d.Location.SourceSpan : default));
128126

129-
return CreateResult(FeaturesResources.Semantic_search_query_failed_to_compile);
127+
return CreateResult(errors, FeaturesResources.Semantic_search_query_failed_to_compile);
130128
}
131129

132130
peStream.Position = 0;
@@ -145,7 +143,7 @@ await observer.OnCompilationFailureAsync(
145143
if (!TryGetFindMethod(queryAssembly, out var findMethod, out var queryKind, out var errorMessage, out var errorMessageArgs))
146144
{
147145
traceSource.TraceInformation($"Semantic search failed: {errorMessage}");
148-
return CreateResult(errorMessage, errorMessageArgs);
146+
return CreateResult(compilationErrors: [], errorMessage, errorMessageArgs);
149147
}
150148

151149
var invocationContext = new QueryExecutionContext(queryText, findMethod, observer, classificationOptions, traceSource);
@@ -155,7 +153,7 @@ await observer.OnCompilationFailureAsync(
155153

156154
if (invocationContext.TerminatedWithException)
157155
{
158-
return CreateResult(FeaturesResources.Semantic_search_query_terminated_with_exception);
156+
return CreateResult(compilationErrors: [], FeaturesResources.Semantic_search_query_terminated_with_exception);
159157
}
160158
}
161159
finally
@@ -175,10 +173,10 @@ await observer.OnCompilationFailureAsync(
175173
}
176174
}
177175

178-
return CreateResult(errorMessage: null);
176+
return CreateResult(compilationErrors: [], errorMessage: null);
179177

180-
ExecuteQueryResult CreateResult(string? errorMessage, params string[]? args)
181-
=> new(errorMessage, args, emitTime, executionTime);
178+
ExecuteQueryResult CreateResult(ImmutableArray<QueryCompilationError> compilationErrors, string? errorMessage, params string[]? args)
179+
=> new(compilationErrors, errorMessage, args, emitTime, executionTime);
182180
}
183181
catch (Exception e) when (FatalError.ReportAndPropagateUnlessCanceled(e, cancellationToken, ErrorSeverity.Critical))
184182
{

‎src/Features/Core/Portable/SemanticSearch/ExecuteQueryResult.cs‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,15 @@
33
// See the LICENSE file in the project root for more information.
44

55
using System;
6+
using System.Collections.Immutable;
67
using System.Runtime.Serialization;
78

89
namespace Microsoft.CodeAnalysis.SemanticSearch;
910

1011
/// <summary>
1112
/// The result of Semantic Search query execution.
1213
/// </summary>
14+
/// <param name="compilationErrors">Compilation errors.</param>
1315
/// <param name="ErrorMessage">An error message if the execution failed.</param>
1416
/// <param name="ErrorMessageArgs">
1517
/// Arguments to be substituted to <paramref name="ErrorMessage"/>.
@@ -20,7 +22,8 @@ namespace Microsoft.CodeAnalysis.SemanticSearch;
2022
/// <param name="ExecutionTime">Time it took to execute the query.</param>
2123
[DataContract]
2224
internal readonly record struct ExecuteQueryResult(
23-
[property: DataMember(Order = 0)] string? ErrorMessage,
24-
[property: DataMember(Order = 1)] string[]? ErrorMessageArgs = null,
25-
[property: DataMember(Order = 2)] TimeSpan EmitTime = default,
26-
[property: DataMember(Order = 3)] TimeSpan ExecutionTime = default);
25+
[property: DataMember(Order = 0)] ImmutableArray<QueryCompilationError> compilationErrors,
26+
[property: DataMember(Order = 1)] string? ErrorMessage,
27+
[property: DataMember(Order = 2)] string[]? ErrorMessageArgs = null,
28+
[property: DataMember(Order = 3)] TimeSpan EmitTime = default,
29+
[property: DataMember(Order = 4)] TimeSpan ExecutionTime = default);

‎src/Features/Core/Portable/SemanticSearch/IRemoteSemanticSearchService.cs‎

Lines changed: 1 addition & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ internal interface ICallback
2121
{
2222
ValueTask OnDefinitionFoundAsync(RemoteServiceCallbackId callbackId, SerializableDefinitionItem definition, CancellationToken cancellationToken);
2323
ValueTask OnUserCodeExceptionAsync(RemoteServiceCallbackId callbackId, UserCodeExceptionInfo exception, CancellationToken cancellationToken);
24-
ValueTask OnCompilationFailureAsync(RemoteServiceCallbackId callbackId, ImmutableArray<QueryCompilationError> errors, CancellationToken cancellationToken);
2524
ValueTask<ClassificationOptions> GetClassificationOptionsAsync(RemoteServiceCallbackId callbackId, string language, CancellationToken cancellationToken);
2625
ValueTask AddItemsAsync(RemoteServiceCallbackId callbackId, int itemCount, CancellationToken cancellationToken);
2726
ValueTask ItemsCompletedAsync(RemoteServiceCallbackId callbackId, int itemCount, CancellationToken cancellationToken);
@@ -43,9 +42,6 @@ public ValueTask OnDefinitionFoundAsync(RemoteServiceCallbackId callbackId, Seri
4342
public ValueTask OnUserCodeExceptionAsync(RemoteServiceCallbackId callbackId, UserCodeExceptionInfo exception, CancellationToken cancellationToken)
4443
=> ((ServerCallback)GetCallback(callbackId)).OnUserCodeExceptionAsync(exception, cancellationToken);
4544

46-
public ValueTask OnCompilationFailureAsync(RemoteServiceCallbackId callbackId, ImmutableArray<QueryCompilationError> errors, CancellationToken cancellationToken)
47-
=> ((ServerCallback)GetCallback(callbackId)).OnCompilationFailureAsync(errors, cancellationToken);
48-
4945
public ValueTask AddItemsAsync(RemoteServiceCallbackId callbackId, int itemCount, CancellationToken cancellationToken)
5046
=> ((ServerCallback)GetCallback(callbackId)).AddItemsAsync(itemCount, cancellationToken);
5147

@@ -81,17 +77,6 @@ public async ValueTask OnUserCodeExceptionAsync(UserCodeExceptionInfo exception,
8177
}
8278
}
8379

84-
public async ValueTask OnCompilationFailureAsync(ImmutableArray<QueryCompilationError> errors, CancellationToken cancellationToken)
85-
{
86-
try
87-
{
88-
await observer.OnCompilationFailureAsync(errors, cancellationToken).ConfigureAwait(false);
89-
}
90-
catch (Exception e) when (FatalError.ReportAndCatchUnlessCanceled(e, cancellationToken))
91-
{
92-
}
93-
}
94-
9580
public async ValueTask AddItemsAsync(int itemCount, CancellationToken cancellationToken)
9681
{
9782
try
@@ -132,7 +117,7 @@ public static async ValueTask<ExecuteQueryResult> ExecuteQueryAsync(Solution sol
132117
var client = await RemoteHostClient.TryGetClientAsync(solution.Services, cancellationToken).ConfigureAwait(false);
133118
if (client == null)
134119
{
135-
return new ExecuteQueryResult(FeaturesResources.Semantic_search_only_supported_on_net_core);
120+
return new ExecuteQueryResult(compilationErrors: [], FeaturesResources.Semantic_search_only_supported_on_net_core);
136121
}
137122

138123
var serverCallback = new ServerCallback(solution, results, classificationOptions);

‎src/Features/Core/Portable/SemanticSearch/ISemanticSearchResultsObserver.cs‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ namespace Microsoft.CodeAnalysis.SemanticSearch;
1414
internal interface ISemanticSearchResultsObserver
1515
{
1616
ValueTask OnUserCodeExceptionAsync(UserCodeExceptionInfo exception, CancellationToken cancellationToken);
17-
ValueTask OnCompilationFailureAsync(ImmutableArray<QueryCompilationError> errors, CancellationToken cancellationToken);
1817
ValueTask OnDefinitionFoundAsync(DefinitionItem definition, CancellationToken cancellationToken);
1918
ValueTask AddItemsAsync(int itemCount, CancellationToken cancellationToken);
2019
ValueTask ItemsCompletedAsync(int itemCount, CancellationToken cancellationToken);

‎src/Features/Core/Portable/SemanticSearch/SearchCompilationFailureDefinitionItem.cs‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
namespace Microsoft.CodeAnalysis.SemanticSearch;
1212

13-
internal sealed class SearchCompilationFailureDefinitionItem(QueryCompilationError error, Document queryDocument)
13+
internal sealed class SearchCompilationFailureDefinitionItem(QueryCompilationError error, Document? queryDocument)
1414
: DefinitionItem(
1515
tags:
1616
[
@@ -24,10 +24,7 @@ internal sealed class SearchCompilationFailureDefinitionItem(QueryCompilationErr
2424
new TaggedText(TextTags.Text, error.Message)
2525
],
2626
nameDisplayParts: [],
27-
sourceSpans:
28-
[
29-
new DocumentSpan(queryDocument, error.Span)
30-
],
27+
sourceSpans: queryDocument != null ? [new DocumentSpan(queryDocument, error.Span)] : [],
3128
classifiedSpans: [],
3229
metadataLocations: [],
3330
properties: null,

‎src/Features/Core/Portable/SemanticSearch/SearchExceptionDefinitionItem.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ internal sealed class SearchExceptionDefinitionItem(string message, ImmutableArr
2828
.. stackTrace
2929
],
3030
nameDisplayParts: exceptionTypeName,
31-
sourceSpans: [documentSpan],
31+
sourceSpans: documentSpan.SourceSpan.IsEmpty ? [] : [documentSpan],
3232
classifiedSpans: [],
3333
metadataLocations: [],
3434
properties: null,

‎src/Features/ExternalAccess/Copilot/Internal/SemanticSearch/CopilotSemanticSearchQueryExecutor.cs‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the MIT license.
3+
// See the LICENSE file in the project root for more information.
4+
5+
using System;
6+
using System.Collections.Immutable;
7+
using System.Composition;
8+
using System.Threading;
9+
using System.Threading.Tasks;
10+
using Microsoft.CodeAnalysis.Classification;
11+
using Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch;
12+
using Microsoft.CodeAnalysis.FindUsages;
13+
using Microsoft.CodeAnalysis.Host;
14+
using Microsoft.CodeAnalysis.Host.Mef;
15+
using Microsoft.CodeAnalysis.SemanticSearch;
16+
using Roslyn.Utilities;
17+
18+
namespace Microsoft.CodeAnalysis.ExternalAccess.Copilot.Internal.SemanticSearch;
19+
20+
[Export(typeof(ICopilotSemanticSearchQueryExecutor)), Shared]
21+
[method: ImportingConstructor]
22+
[method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
23+
internal sealed class CopilotSemanticSearchQueryExecutor(IHostWorkspaceProvider workspaceProvider) : ICopilotSemanticSearchQueryExecutor
24+
{
25+
private sealed class ResultsObserver(CancellationTokenSource cancellationSource, int resultCountLimit) : ISemanticSearchResultsObserver
26+
{
27+
private ImmutableList<string> _results = [];
28+
public string? RuntimeException { get; private set; }
29+
public bool LimitReached { get; private set; }
30+
31+
public ImmutableList<string> Results => _results;
32+
33+
public ValueTask AddItemsAsync(int itemCount, CancellationToken cancellationToken)
34+
=> ValueTaskFactory.CompletedTask;
35+
36+
public ValueTask ItemsCompletedAsync(int itemCount, CancellationToken cancellationToken)
37+
=> ValueTaskFactory.CompletedTask;
38+
39+
public ValueTask OnUserCodeExceptionAsync(UserCodeExceptionInfo exception, CancellationToken cancellationToken)
40+
{
41+
RuntimeException ??= $"{exception.TypeName.ToVisibleDisplayString(includeLeftToRightMarker: false)}: {exception.Message}{Environment.NewLine}{exception.StackTrace.ToVisibleDisplayString(includeLeftToRightMarker: false)}";
42+
cancellationSource.Cancel();
43+
return ValueTaskFactory.CompletedTask;
44+
}
45+
46+
public ValueTask OnDefinitionFoundAsync(DefinitionItem definition, CancellationToken cancellationToken)
47+
{
48+
if (!ImmutableInterlocked.Update(ref _results,
49+
list => list.Count == resultCountLimit ? list : list.Add(definition.NameDisplayParts.ToVisibleDisplayString(includeLeftToRightMarker: false))))
50+
{
51+
LimitReached = true;
52+
cancellationSource.Cancel();
53+
}
54+
55+
return ValueTaskFactory.CompletedTask;
56+
}
57+
}
58+
59+
/// <summary>
60+
/// We only use symbol display names, classification is not relevant.
61+
/// </summary>
62+
private sealed class DefaultClassificationOptionsProvider : OptionsProvider<ClassificationOptions>
63+
{
64+
public static readonly DefaultClassificationOptionsProvider Instance = new();
65+
66+
public ValueTask<ClassificationOptions> GetOptionsAsync(LanguageServices languageServices, CancellationToken cancellationToken)
67+
=> new(ClassificationOptions.Default);
68+
}
69+
70+
private readonly Workspace _workspace = workspaceProvider.Workspace;
71+
72+
public async Task<CopilotSemanticSearchQueryResults> ExecuteAsync(string query, int resultCountLimit, CancellationToken cancellationToken)
73+
{
74+
Contract.ThrowIfFalse(resultCountLimit > 0);
75+
76+
using var cancellationSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
77+
var observer = new ResultsObserver(cancellationSource, resultCountLimit);
78+
79+
try
80+
{
81+
var result = await RemoteSemanticSearchServiceProxy.ExecuteQueryAsync(
82+
_workspace.CurrentSolution,
83+
LanguageNames.CSharp,
84+
query,
85+
SemanticSearchUtilities.ReferenceAssembliesDirectory,
86+
observer,
87+
DefaultClassificationOptionsProvider.Instance,
88+
cancellationSource.Token).ConfigureAwait(false);
89+
90+
return new CopilotSemanticSearchQueryResults()
91+
{
92+
Symbols = observer.Results,
93+
CompilationErrors = result.compilationErrors.SelectAsArray(e => (e.Id, e.Message)),
94+
Error = (result.ErrorMessage != null) ? string.Format(result.ErrorMessage, result.ErrorMessageArgs ?? []) : null,
95+
LimitReached = false,
96+
};
97+
}
98+
catch (OperationCanceledException) when (cancellationSource.IsCancellationRequested)
99+
{
100+
return new CopilotSemanticSearchQueryResults()
101+
{
102+
Symbols = observer.Results,
103+
CompilationErrors = [],
104+
Error = observer.RuntimeException != null ? $"The query failed with an exception: {observer.RuntimeException}" : null,
105+
LimitReached = observer.LimitReached,
106+
};
107+
}
108+
}
109+
}

‎src/Features/ExternalAccess/Copilot/InternalAPI.Unshipped.txt‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,18 @@ Microsoft.CodeAnalysis.ExternalAccess.Copilot.IExternalCSharpOnTheFlyDocsService
5656
Microsoft.CodeAnalysis.ExternalAccess.Copilot.IExternalCSharpCopilotGenerateImplementationService
5757
Microsoft.CodeAnalysis.ExternalAccess.Copilot.RelatedDocuments.ICopilotRelatedDocumentsService
5858
Microsoft.CodeAnalysis.ExternalAccess.Copilot.RelatedDocuments.ICopilotRelatedDocumentsService.GetRelatedDocumentIdsAsync(Microsoft.CodeAnalysis.Document! document, int position, System.Func<System.Collections.Immutable.ImmutableArray<Microsoft.CodeAnalysis.DocumentId!>, System.Threading.CancellationToken, System.Threading.Tasks.ValueTask>! callbackAsync, System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.ValueTask
59+
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.CopilotSemanticSearchQueryResults
60+
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.CopilotSemanticSearchQueryResults.CompilationErrors.get -> System.Collections.Generic.IReadOnlyList<(string! id, string! message)>!
61+
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.CopilotSemanticSearchQueryResults.CompilationErrors.init -> void
62+
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.CopilotSemanticSearchQueryResults.CopilotSemanticSearchQueryResults() -> void
63+
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.CopilotSemanticSearchQueryResults.Error.get -> string?
64+
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.CopilotSemanticSearchQueryResults.Error.init -> void
65+
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.CopilotSemanticSearchQueryResults.LimitReached.get -> bool
66+
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.CopilotSemanticSearchQueryResults.LimitReached.init -> void
67+
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.CopilotSemanticSearchQueryResults.Symbols.get -> System.Collections.Generic.IReadOnlyList<string!>!
68+
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.CopilotSemanticSearchQueryResults.Symbols.init -> void
69+
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.ICopilotSemanticSearchQueryExecutor
70+
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.ICopilotSemanticSearchQueryExecutor.ExecuteAsync(string! query, int resultCountLimit, System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task<Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.CopilotSemanticSearchQueryResults>!
5971
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.ISemanticSearchCopilotServiceImpl
6072
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.ISemanticSearchCopilotServiceImpl.TryGetQueryAsync(string! text, Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.SemanticSearchCopilotContextImpl! context, System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.ValueTask<Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.SemanticSearchCopilotGeneratedQueryImpl>
6173
Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch.SemanticSearchCopilotContextImpl

‎src/Features/ExternalAccess/Copilot/SemanticSearch/ICopilotSemanticSearchQueryExecutor.cs‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the MIT license.
3+
// See the LICENSE file in the project root for more information.
4+
5+
using System.Collections.Generic;
6+
using System.Threading;
7+
using System.Threading.Tasks;
8+
9+
namespace Microsoft.CodeAnalysis.ExternalAccess.Copilot.SemanticSearch;
10+
11+
internal interface ICopilotSemanticSearchQueryExecutor
12+
{
13+
Task<CopilotSemanticSearchQueryResults> ExecuteAsync(string query, int resultCountLimit, CancellationToken cancellationToken);
14+
}
15+
16+
internal readonly struct CopilotSemanticSearchQueryResults
17+
{
18+
public required IReadOnlyList<string> Symbols { get; init; }
19+
public required IReadOnlyList<(string id, string message)> CompilationErrors { get; init; }
20+
public string? Error { get; init; }
21+
public required bool LimitReached { get; init; }
22+
}

0 commit comments

Comments
 (0)