agent-framework/dotnet/tests/OpenAIAssistant.IntegrationTests/OpenAIAssistantClientExtensionsTests.cs

// Copyright (c) Microsoft. All rights reserved.

#pragma warning disable CS0618 // Type or member is obsolete - Testing deprecated OpenAI Assistants API extension methods

using System;
using System.Diagnostics;
using System.IO;
using System.Threading.Tasks;
using AgentConformance.IntegrationTests.Support;
using Microsoft.Agents.AI;
using Microsoft.Extensions.AI;
using OpenAI;
using OpenAI.Assistants;
using OpenAI.Files;
using OpenAI.VectorStores;
using Shared.IntegrationTests;

namespace OpenAIAssistant.IntegrationTests;

public class OpenAIAssistantClientExtensionsTests
{
    private const string SkipCodeInterpreterReason = "OpenAI Assistant Code Interpreter intermittently fails in CI";

    private readonly AssistantClient _assistantClient = new OpenAIClient(TestConfiguration.GetRequiredValue(TestSettings.OpenAIApiKey)).GetAssistantClient();
    private readonly OpenAIFileClient _fileClient = new OpenAIClient(TestConfiguration.GetRequiredValue(TestSettings.OpenAIApiKey)).GetOpenAIFileClient();

    [Theory]
    [InlineData("CreateWithChatClientAgentOptionsAsync")]
    [InlineData("CreateWithChatClientAgentOptionsSync")]
    [InlineData("CreateWithParamsAsync")]
    public async Task CreateAIAgentAsync_WithAIFunctionTool_InvokesFunctionAsync(string createMechanism)
    {
        // Arrange
        const string AgentInstructions = "You are a helpful weather assistant. Always call the GetWeather function to answer questions about weather.";

        static string GetWeather(string location) => $"The weather in {location} is sunny with a high of 23C.";
        var weatherFunction = AIFunctionFactory.Create(GetWeather, nameof(GetWeather));

        // Act
        var agent = createMechanism switch
        {
            "CreateWithChatClientAgentOptionsAsync" => await this._assistantClient.CreateAIAgentAsync(
                model: TestConfiguration.GetRequiredValue(TestSettings.OpenAIChatModelName),
                options: new ChatClientAgentOptions()
                {
                    ChatOptions = new()
                    {
                        Instructions = AgentInstructions,
                        Tools = [weatherFunction]
                    }
                }),
            "CreateWithChatClientAgentOptionsSync" => await this._assistantClient.CreateAIAgentAsync(
                model: TestConfiguration.GetRequiredValue(TestSettings.OpenAIChatModelName),
                options: new ChatClientAgentOptions()
                {
                    ChatOptions = new()
                    {
                        Instructions = AgentInstructions,
                        Tools = [weatherFunction]
                    }
                }),
            "CreateWithParamsAsync" => await this._assistantClient.CreateAIAgentAsync(
                model: TestConfiguration.GetRequiredValue(TestSettings.OpenAIChatModelName),
                instructions: AgentInstructions,
                tools: [weatherFunction]),
            _ => throw new InvalidOperationException($"Unknown create mechanism: {createMechanism}")
        };

        try
        {
            // Trigger function call.
            var response = await agent.RunAsync("What is the weather like in Amsterdam?");
            var text = response.Text;

            // Assert
            Assert.Contains("Amsterdam", text, StringComparison.OrdinalIgnoreCase);
            Assert.Contains("sunny", text, StringComparison.OrdinalIgnoreCase);
            Assert.Contains("23", text, StringComparison.OrdinalIgnoreCase);
        }
        finally
        {
            await this._assistantClient.DeleteAssistantAsync(agent.Id);
        }
    }

    [Theory(Skip = SkipCodeInterpreterReason)]
    [InlineData("CreateWithChatClientAgentOptionsAsync")]
    [InlineData("CreateWithChatClientAgentOptionsSync")]
    [InlineData("CreateWithParamsAsync")]
    public async Task CreateAIAgentAsync_WithHostedCodeInterpreter_RunsCodeAsync(string createMechanism)
    {
        // Arrange
        const string Instructions = "Use the Code Interpreter Tool to run the uploaded python file and respond only with the secret number.";

        // Create a python file that prints a known value.
        var codeFilePath = Path.GetTempFileName() + "openai_secret_number.py";
        File.WriteAllText(
            path: codeFilePath,
            contents: "print(\"OPENAI_SECRET=13579\")" // Deterministic output we will look for.
        );

        // Upload file to OpenAI Assistants file store for use with the Code Interpreter.
        var uploadResult = await this._fileClient.UploadFileAsync(codeFilePath, FileUploadPurpose.Assistants);
        string uploadedFileId = uploadResult.Value.Id;
        var codeInterpreterTool = new HostedCodeInterpreterTool() { Inputs = [new HostedFileContent(uploadedFileId)] };

        var agent = createMechanism switch
        {
            "CreateWithChatClientAgentOptionsAsync" => await this._assistantClient.CreateAIAgentAsync(
                model: TestConfiguration.GetRequiredValue(TestSettings.OpenAIChatModelName),
                options: new ChatClientAgentOptions()
                {
                    ChatOptions = new()
                    {
                        Instructions = Instructions,
                        Tools = [codeInterpreterTool]
                    }
                }),
            "CreateWithChatClientAgentOptionsSync" => await this._assistantClient.CreateAIAgentAsync(
                model: TestConfiguration.GetRequiredValue(TestSettings.OpenAIChatModelName),
                options: new ChatClientAgentOptions()
                {
                    ChatOptions = new()
                    {
                        Instructions = Instructions,
                        Tools = [codeInterpreterTool]
                    }
                }),
            "CreateWithParamsAsync" => await this._assistantClient.CreateAIAgentAsync(
                model: TestConfiguration.GetRequiredValue(TestSettings.OpenAIChatModelName),
                instructions: Instructions,
                tools: [codeInterpreterTool]),
            _ => throw new InvalidOperationException($"Unknown create mechanism: {createMechanism}")
        };

        try
        {
            var response = await agent.RunAsync("What is the OPENAI_SECRET number?");
            var text = response.ToString();
            Assert.Contains("13579", text);
        }
        finally
        {
            await this._assistantClient.DeleteAssistantAsync(agent.Id);
            await this._fileClient.DeleteFileAsync(uploadedFileId);
            File.Delete(codeFilePath);
        }
    }

    [Theory(Skip = "For manual testing only")]
    [InlineData("CreateWithChatClientAgentOptionsAsync")]
    [InlineData("CreateWithChatClientAgentOptionsSync")]
    [InlineData("CreateWithParamsAsync")]
    public async Task CreateAIAgentAsync_WithHostedFileSearchTool_SearchesFilesAsync(string createMechanism)
    {
        // Arrange.
        const string Instructions = """
            You are a helpful agent that can help fetch data from files you know about.
            Use the File Search Tool to look up codes for words.
            Do not answer a question unless you can find the answer using the File Search Tool.
            """;

        // Create a local file with deterministic content and upload it.
        var searchFilePath = Path.GetTempFileName() + "wordcodelookup.txt";
        File.WriteAllText(
            path: searchFilePath,
            contents: "The word 'apple' uses the code 442345, while the word 'banana' uses the code 673457.");
        var uploadResult = await this._fileClient.UploadFileAsync(searchFilePath, FileUploadPurpose.Assistants);
        string uploadedFileId = uploadResult.Value.Id;

        // Create a vector store backing the file search (HostedFileSearchTool requires a vector store id).
        var vectorStoreClient = new OpenAIClient(TestConfiguration.GetRequiredValue(TestSettings.OpenAIApiKey)).GetVectorStoreClient();
        var vectorStoreCreate = await vectorStoreClient.CreateVectorStoreAsync(options: new VectorStoreCreationOptions()
        {
            Name = "WordCodeLookup_VectorStore",
            FileIds = { uploadedFileId }
        });
        string vectorStoreId = vectorStoreCreate.Value.Id;

        // Wait for vector store indexing to complete before using it
        await WaitForVectorStoreReadyAsync(vectorStoreClient, vectorStoreId);

        var fileSearchTool = new HostedFileSearchTool() { Inputs = [new HostedVectorStoreContent(vectorStoreId)] };

        var agent = createMechanism switch
        {
            "CreateWithChatClientAgentOptionsAsync" => await this._assistantClient.CreateAIAgentAsync(
                model: TestConfiguration.GetRequiredValue(TestSettings.OpenAIChatModelName),
                options: new ChatClientAgentOptions()
                {
                    ChatOptions = new()
                    {
                        Instructions = Instructions,
                        Tools = [fileSearchTool]
                    }
                }),
            "CreateWithChatClientAgentOptionsSync" => await this._assistantClient.CreateAIAgentAsync(
                model: TestConfiguration.GetRequiredValue(TestSettings.OpenAIChatModelName),
                options: new ChatClientAgentOptions()
                {
                    ChatOptions = new()
                    {
                        Instructions = Instructions,
                        Tools = [fileSearchTool]
                    }
                }),
            "CreateWithParamsAsync" => await this._assistantClient.CreateAIAgentAsync(
                model: TestConfiguration.GetRequiredValue(TestSettings.OpenAIChatModelName),
                instructions: Instructions,
                tools: [fileSearchTool]),
            _ => throw new InvalidOperationException($"Unknown create mechanism: {createMechanism}")
        };

        try
        {
            // Act - ask about banana code which must be retrieved via file search.
            var response = await agent.RunAsync("Can you give me the documented code for 'banana'?");
            var text = response.ToString();
            Assert.Contains("673457", text);
        }
        finally
        {
            await this._assistantClient.DeleteAssistantAsync(agent.Id);
            await vectorStoreClient.DeleteVectorStoreAsync(vectorStoreId);
            await this._fileClient.DeleteFileAsync(uploadedFileId);
            File.Delete(searchFilePath);
        }
    }

    /// <summary>
    /// Waits for a vector store to complete indexing by polling its status.
    /// </summary>
    /// <param name="client">The vector store client.</param>
    /// <param name="vectorStoreId">The ID of the vector store.</param>
    /// <param name="maxWaitSeconds">Maximum time to wait in seconds (default: 30).</param>
    /// <returns>A task that completes when the vector store is ready or throws on timeout/failure.</returns>
    private static async Task WaitForVectorStoreReadyAsync(
        VectorStoreClient client,
        string vectorStoreId,
        int maxWaitSeconds = 30)
    {
        Stopwatch sw = Stopwatch.StartNew();
        while (sw.Elapsed.TotalSeconds < maxWaitSeconds)
        {
            VectorStore vectorStore = await client.GetVectorStoreAsync(vectorStoreId);
            VectorStoreStatus status = vectorStore.Status;

            if (status == VectorStoreStatus.Completed)
            {
                if (vectorStore.FileCounts.Failed > 0)
                {
                    throw new InvalidOperationException("Vector store indexing failed for some files");
                }

                return;
            }

            if (status == VectorStoreStatus.Expired)
            {
                throw new InvalidOperationException("Vector store has expired");
            }

            await Task.Delay(1000);
        }

        throw new TimeoutException($"Vector store did not complete indexing within {maxWaitSeconds}s");
    }
}