fixed citations

microsoft · Nov 19, 2024 · 8ada3b3 · 8ada3b3
1 parent 9bcd35f
commit 8ada3b3
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 35 deletions.
diff --git a/.../packages/Microsoft.TeamsAI/Microsoft.TeamsAI.Tests/Application/StreamingResponseTests.cs b/.../packages/Microsoft.TeamsAI/Microsoft.TeamsAI.Tests/Application/StreamingResponseTests.cs
@@ -220,7 +220,8 @@ void CaptureSend(Activity[] arg)
             StreamingResponse streamer = new(turnContext);
             List<Citation> citations = new List<Citation>();
             citations.Add(new Citation(content: "test-content", title: "test", url: "https://example.com"));
-            streamer.QueueTextChunk("first", citations);
+            streamer.SetCitations(citations);
+            streamer.QueueTextChunk("first");
             await streamer.WaitForQueue();
             streamer.QueueTextChunk("second");
             await streamer.WaitForQueue();
@@ -229,6 +230,10 @@ void CaptureSend(Activity[] arg)
             streamer.SensitivityLabel = new SensitivityUsageInfo() { Name= "Sensitivity"};
             await streamer.EndStream();
             Assert.Equal(2, streamer.UpdatesSent());
+            if (streamer.Citations != null)
+            {
+                Assert.Single(streamer.Citations);
+            }
         }
 
         [Fact]

diff --git a/dotnet/packages/Microsoft.TeamsAI/Microsoft.TeamsAI/AI/Clients/LLMClient.cs b/dotnet/packages/Microsoft.TeamsAI/Microsoft.TeamsAI/AI/Clients/LLMClient.cs
@@ -199,10 +199,15 @@ public async Task<PromptResponse> CompletePromptAsync(
                 string text = args.Chunk.delta?.GetContent<string>() ?? "";
                 IList<Citation>? citations = args.Chunk.delta?.Context?.Citations ?? null;
 
+                if (citations != null)
+                {
+                    streamer.SetCitations(citations);
+                }
+
 
                 if (text.Length > 0)
                 {
-                    streamer.QueueTextChunk(text, citations);
+                    streamer.QueueTextChunk(text);
                 }
             });
 

diff --git a/dotnet/packages/Microsoft.TeamsAI/Microsoft.TeamsAI/Application/StreamingResponse.cs b/dotnet/packages/Microsoft.TeamsAI/Microsoft.TeamsAI/Application/StreamingResponse.cs
@@ -92,6 +92,40 @@ public Task WaitForQueue()
             return this._queueSync != null ? this._queueSync : Task.CompletedTask;
         }
 
+        /// <summary>
+        ///  Sets the citations for the full message.
+        /// </summary>
+        /// <param name="citations">Citations to be included in the message.</param>
+        public void SetCitations(IList<Citation> citations)
+        {
+            if (citations.Count > 0)
+            {
+                if (this.Citations == null)
+                {
+                    this.Citations = new List<ClientCitation>();
+                }
+
+                int currPos = this.Citations.Count;
+
+                foreach (Citation citation in citations)
+                {
+                    string abs = CitationUtils.Snippet(citation.Content, 480);
+
+                    this.Citations.Add(new ClientCitation()
+                    {
+                        Position = $"{currPos}",
+                        Appearance = new ClientCitationAppearance()
+                        {
+                            Name = citation.Title,
+                            Abstract = abs
+                        }
+                    });
+                    currPos++;
+                }
+
+            }
+        }
+
         /// <summary>
         /// Queues an informative update to be sent to the client.
         /// </summary>
@@ -131,38 +165,8 @@ public void QueueTextChunk(string text, IList<Citation>? citations = null)
 
             Message += text;
 
-            if (citations != null && citations.Count > 0)
-            {
-                if (this.Citations == null)
-                {
-                    this.Citations = new List<ClientCitation>();
-                }
-
-                int currPos = this.Citations.Count;
-
-                foreach (Citation citation in citations)
-                {
-                    string abs = CitationUtils.Snippet(citation.Content, 480);
-
-                    this.Citations.Add(new ClientCitation()
-                    {
-                        Position = $"{currPos}",
-                        Appearance = new ClientCitationAppearance()
-                        {
-                            Name = citation.Title,
-                            Abstract = abs
-                        }
-                    });
-                    currPos++;
-                }
-
-                // If there are citations, modify the content so that the sources are numbers instead of [doc1], [doc2], etc.
-                this.Message = this.Citations.Count == 0 ? this.Message : CitationUtils.FormatCitationsResponse(this.Message);
-
-                // If there are citations, filter out the citations unused in content.
-                this.Citations = this.Citations.Count > 0 ? CitationUtils.GetUsedCitations(this.Message, this.Citations) : new List<ClientCitation>();
-
-            }
+            // If there are citations, modify the content so that the sources are numbers instead of [doc1], [doc2], etc.
+            this.Message = CitationUtils.FormatCitationsResponse(this.Message);
 
             QueueNextChunk();
         }
@@ -320,6 +324,19 @@ private async Task SendActivity(Activity activity)
                 }
             };
 
+            if (this.Citations != null && this.Citations.Count > 0 && !this._ended)
+            {
+                // If there are citations, filter out the citations unused in content.
+                List<ClientCitation>? currCitations = CitationUtils.GetUsedCitations(this.Message, this.Citations);
+                AIEntity entity = new AIEntity();
+                if (currCitations != null && currCitations.Count > 0)
+                {
+                    entity.Citation = currCitations;
+                }
+
+                activity.Entities.Add(entity);
+            }
+
             // Add in Powered by AI feature flags
             if (this._ended)
             {

diff --git a/dotnet/packages/Microsoft.TeamsAI/Microsoft.TeamsAI/TypingTimer.cs b/dotnet/packages/Microsoft.TeamsAI/Microsoft.TeamsAI/TypingTimer.cs
@@ -1,6 +1,7 @@
 using Microsoft.Teams.AI.Utilities;
 using Microsoft.Bot.Schema;
 using Microsoft.Bot.Builder;
+using Microsoft.Teams.AI.Application;
 
 namespace Microsoft.Teams.AI
 {
@@ -121,7 +122,7 @@ private Task<ResourceResponse[]> StopTimerWhenSendMessageActivityHandlerAsync(IT
             {
                 foreach (Activity activity in activities)
                 {
-                    if (activity.Type == ActivityTypes.Message)
+                    if (activity.Type == ActivityTypes.Message || activity.GetChannelData<StreamingChannelData>()?.StreamType != null)
                     {
                         Dispose();
                         break;