From 765daa854435f2dedac9c7f719698e82dc1112a3 Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Wed, 18 Jun 2025 14:51:26 +0200 Subject: [PATCH] Refactor citation and token usage handling Updated comments for clarity and streamlined logic for managing tokenUsageResponse. Removed explicit null checks in favor of a null-coalescing assignment. Ensured citations are always extracted and returned at the end of the streaming process. --- .../Services/VectorSearchService.cs | 20 +++++-------------- 1 file changed, 5 insertions(+), 15 deletions(-) diff --git a/SqlDatabaseVectorSearch/Services/VectorSearchService.cs b/SqlDatabaseVectorSearch/Services/VectorSearchService.cs index 8133b79..4d1b108 100644 --- a/SqlDatabaseVectorSearch/Services/VectorSearchService.cs +++ b/SqlDatabaseVectorSearch/Services/VectorSearchService.cs @@ -115,7 +115,7 @@ public partial class VectorSearchService(IServiceProvider serviceProvider, Appli if (token?.Contains('【') == true) { - // Citations are started when the first token contains a 【 character. + // Citations are started when we encounter a token containing a 【 character. // We need to track it because we don't want to return the citations in the actual response. areCitationsStarted = true; } @@ -126,22 +126,12 @@ public partial class VectorSearchService(IServiceProvider serviceProvider, Appli } // Token usage is expected in the last message. - tokenUsageResponse = tokenUsage is not null ? new(tokenUsage) : null; - if (tokenUsageResponse is not null) - { - // Response is complete, we can return the citations. - var (_, citations) = ExtractCitations(fullAnswer.ToString()); - yield return new(null, StreamState.End, tokenUsageResponse, citations); - } + tokenUsageResponse ??= tokenUsage is not null ? new(tokenUsage) : null; } - // If the token usage has not been returned in the last message, we must explicitly tell that the stream is ended. - if (tokenUsageResponse is null) - { - // Extract citations at the end of streaming. - var (_, citations) = ExtractCitations(fullAnswer.ToString()); - yield return new(null, StreamState.End, null, citations); - } + // Extract citations at the end of streaming. + var (_, citations) = ExtractCitations(fullAnswer.ToString()); + yield return new(null, StreamState.End, tokenUsageResponse, citations); } private async Task<(ChatResponse ReformulatedQuestion, int EmbeddingTokenCount, IEnumerable Chunks)> CreateContextAsync(Question question, bool reformulate, CancellationToken cancellationToken)