From 17c2ed35651c030719dccba15b06199b1774bb6b Mon Sep 17 00:00:00 2001
From: XiaoYun Zhang <xiaoyuz@microsoft.com>
Date: Thu, 18 Jul 2024 21:33:39 -0700
Subject: [PATCH] add streaming tool call example

---
 .../GettingStart/Streaming_Tool_Call.cs       | 56 +++++++++++++++++++
 .../Orchestrator/RolePlayOrchestratorTests.cs | 31 ++++++++++
 2 files changed, 87 insertions(+)
 create mode 100644 dotnet/sample/AutoGen.BasicSamples/GettingStart/Streaming_Tool_Call.cs

diff --git a/dotnet/sample/AutoGen.BasicSamples/GettingStart/Streaming_Tool_Call.cs b/dotnet/sample/AutoGen.BasicSamples/GettingStart/Streaming_Tool_Call.cs
new file mode 100644
index 000000000000..48ebd127b562
--- /dev/null
+++ b/dotnet/sample/AutoGen.BasicSamples/GettingStart/Streaming_Tool_Call.cs
@@ -0,0 +1,56 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Streaming_Tool_Call.cs
+
+using AutoGen.Core;
+using AutoGen.OpenAI;
+using AutoGen.OpenAI.Extension;
+using Azure.AI.OpenAI;
+using FluentAssertions;
+
+namespace AutoGen.BasicSample.GettingStart;
+
+internal class Streaming_Tool_Call
+{
+    public static async Task RunAsync()
+    {
+        #region Create_tools
+        var tools = new Tools();
+        #endregion Create_tools
+
+        #region Create_auto_invoke_middleware
+        var autoInvokeMiddleware = new FunctionCallMiddleware(
+            functions: [tools.GetWeatherFunctionContract],
+            functionMap: new Dictionary<string, Func<string, Task<string>>>()
+            {
+                { tools.GetWeatherFunctionContract.Name, tools.GetWeatherWrapper },
+            });
+        #endregion Create_auto_invoke_middleware
+
+        #region Create_Agent
+        var apiKey = Environment.GetEnvironmentVariable("OPENAI_API_KEY") ?? throw new Exception("Please set OPENAI_API_KEY environment variable.");
+        var model = "gpt-4o";
+        var openaiClient = new OpenAIClient(apiKey);
+        var agent = new OpenAIChatAgent(
+            openAIClient: openaiClient,
+            name: "agent",
+            modelName: model,
+            systemMessage: "You are a helpful AI assistant")
+            .RegisterMessageConnector()
+            .RegisterStreamingMiddleware(autoInvokeMiddleware)
+            .RegisterPrintMessage();
+        #endregion Create_Agent
+
+        IMessage finalReply = null;
+        var question = new TextMessage(Role.User, "What's the weather in Seattle");
+
+        // In streaming function call
+        // function can only be invoked untill all the chunks are collected
+        // therefore, only one ToolCallAggregateMessage chunk will be return here.
+        await foreach (var message in agent.GenerateStreamingReplyAsync([question]))
+        {
+            finalReply = message;
+        }
+
+        finalReply?.GetContent().Should().Be("The weather in Seattle is sunny.");
+    }
+}
diff --git a/dotnet/test/AutoGen.Tests/Orchestrator/RolePlayOrchestratorTests.cs b/dotnet/test/AutoGen.Tests/Orchestrator/RolePlayOrchestratorTests.cs
index eea0a90c1caf..5a2cebb66cff 100644
--- a/dotnet/test/AutoGen.Tests/Orchestrator/RolePlayOrchestratorTests.cs
+++ b/dotnet/test/AutoGen.Tests/Orchestrator/RolePlayOrchestratorTests.cs
@@ -228,6 +228,37 @@ public async Task GPT_3_5_CoderReviewerRunnerTestAsync()
         await CoderReviewerRunnerTestAsync(openAIChatAgent);
     }
 
+    [ApiKeyFact("OPENAI_API_KEY")]
+    public async Task GPT_4o_CoderReviewerRunnerTestAsync()
+    {
+        var apiKey = Environment.GetEnvironmentVariable("OPENAI_API_KEY");
+        var model = "gpt-4o";
+        var openaiClient = new OpenAIClient(apiKey);
+        var openAIChatAgent = new OpenAIChatAgent(
+            openAIClient: openaiClient,
+            name: "assistant",
+            modelName: model)
+            .RegisterMessageConnector();
+
+        await CoderReviewerRunnerTestAsync(openAIChatAgent);
+    }
+
+    [ApiKeyFact("OPENAI_API_KEY")]
+    public async Task GPT_4o_mini_CoderReviewerRunnerTestAsync()
+    {
+        var apiKey = Environment.GetEnvironmentVariable("OPENAI_API_KEY");
+        var model = "gpt-4o-mini";
+        var openaiClient = new OpenAIClient(apiKey);
+        var openAIChatAgent = new OpenAIChatAgent(
+            openAIClient: openaiClient,
+            name: "assistant",
+            modelName: model)
+            .RegisterMessageConnector();
+
+        await CoderReviewerRunnerTestAsync(openAIChatAgent);
+    }
+
+
     [ApiKeyFact("GOOGLE_GEMINI_API_KEY")]
     public async Task GoogleGemini_1_5_flash_001_CoderReviewerRunnerTestAsync()
     {