Adding upload file and poll operation with sample and test (#31928)

Co-authored-by: Zachary King <[email protected]>

sdk/ai/ai-projects/review/ai-projects.api.md

@@ -144,6 +144,7 @@ export interface AgentsOperations {
     updateRun: (threadId: string, runId: string, options?: UpdateRunOptions, requestParams?: OptionalRequestParameters) => Promise<ThreadRunOutput>;
     updateThread: (threadId: string, options?: UpdateAgentThreadOptions, requestParams?: OptionalRequestParameters) => Promise<AgentThreadOutput>;
     uploadFile: (data: ReadableStream | NodeJS.ReadableStream, purpose: FilePurpose, fileName?: string, requestParams?: OptionalRequestParameters) => Promise<OpenAIFileOutput>;
+    uploadFileAndPoll: (data: ReadableStream | NodeJS.ReadableStream, purpose: FilePurpose, fileName?: string, pollingOptions?: PollingOptions, requestParams?: OptionalRequestParameters) => Promise<OpenAIFileOutput>;
 }
 
 // @public

sdk/ai/ai-projects/samples-dev/agents/files_polling.ts

@@ -0,0 +1,37 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import {AIProjectsClient} from "@azure/ai-projects"
+import { DefaultAzureCredential } from "@azure/identity";
+import * as dotenv from "dotenv";
+import { Readable } from "stream";
+dotenv.config();
+
+const connectionString = process.env["AZURE_AI_PROJECTS_CONNECTION_STRING"] || "<endpoint>>;<subscription>;<resource group>;<project>";
+
+export async function main(): Promise<void> {
+  const client = AIProjectsClient.fromConnectionString(connectionString || "", new DefaultAzureCredential());
+
+  // Set up abort controller (optional)
+  // Polling can be stopped by calling abortController.abort()
+  const abortController = new AbortController();
+
+  // Create file content
+  const fileContent = "Hello, World!";
+  const readable = new Readable();
+  readable.push(fileContent);
+  readable.push(null); // end the stream
+
+  // Upload file and poll
+  const pollingOptions = { sleepIntervalInMs: 1000, abortSignal: abortController.signal };
+  const file = await client.agents.uploadFileAndPoll(readable, "assistants", "my-polling-file", pollingOptions);
+  console.log(`Uploaded file with status ${file.status}, file ID : ${file.id}`);
+
+  // Delete file
+  await client.agents.deleteFile(file.id);
+  console.log(`Deleted file, file ID ${file.id}`);
+}
+
+main().catch((err) => {
+  console.error("The sample encountered an error:", err);
+});

sdk/ai/ai-projects/src/agents/files.ts

@@ -4,6 +4,8 @@
 import { Client, createRestError, StreamableMethod } from "@azure-rest/core-client";
 import { FileDeletionStatusOutput, FileListResponseOutput, OpenAIFileOutput } from "../generated/src/outputModels.js";
 import { DeleteFileParameters, GetFileContentParameters, GetFileParameters, ListFilesParameters, UploadFileParameters } from "../generated/src/parameters.js";
+import { OptionalRequestParameters, PollingOptions } from "./customModels.js";
+import { AgentsPoller } from "./poller.js";
 
 const expectedStatuses = ["200"];
 
@@ -42,6 +44,28 @@ export async function uploadFile(
   return result.body; 
 }
 
+/** Uploads a file for use by other operations. */
+export async function uploadFileAndPoll(
+  context: Client, 
+  options: UploadFileParameters, 
+  pollingOptions?: PollingOptions,
+  requestParams?: OptionalRequestParameters,
+): Promise<OpenAIFileOutput> {
+  async function updateUploadFileAndPoll(currentResult?: OpenAIFileOutput): Promise<{result: OpenAIFileOutput; completed: boolean}> {
+    let file: OpenAIFileOutput;
+    if (!currentResult) {
+      file = await uploadFile(context, {...options, ...requestParams});
+    } else {
+      file = await getFile(context, currentResult.id, options);
+    }
+    return { result: file, completed: file.status === "uploaded" || file.status === "processed" || file.status === "deleted" };
+  }
+
+  const poller = new AgentsPoller<OpenAIFileOutput>({ update: updateUploadFileAndPoll, pollingOptions: pollingOptions,});
+
+  return poller.pollUntilDone(); 
+}
+
 /** Delete a previously uploaded file. */
 export async function deleteFile(
   context: Client,

sdk/ai/ai-projects/src/agents/index.ts

@@ -4,7 +4,7 @@
 import { Client, StreamableMethod } from "@azure-rest/core-client";
 import { AgentDeletionStatusOutput, AgentOutput, AgentThreadOutput, FileDeletionStatusOutput, FileListResponseOutput, OpenAIFileOutput, OpenAIPageableListOfAgentOutput, OpenAIPageableListOfRunStepOutput, OpenAIPageableListOfThreadMessageOutput, OpenAIPageableListOfThreadRunOutput, OpenAIPageableListOfVectorStoreFileOutput, OpenAIPageableListOfVectorStoreOutput, RunStepOutput, ThreadDeletionStatusOutput, ThreadMessageOutput, ThreadRunOutput, VectorStoreDeletionStatusOutput, VectorStoreFileBatchOutput, VectorStoreFileDeletionStatusOutput, VectorStoreFileOutput, VectorStoreOutput } from "../generated/src/outputModels.js";
 import { createAgent, deleteAgent, getAgent, listAgents, updateAgent } from "./assistants.js";
-import { deleteFile, getFile, getFileContent, listFiles, uploadFile } from "./files.js";
+import { deleteFile, getFile, getFileContent, listFiles, uploadFile, uploadFileAndPoll } from "./files.js";
 import { createThread, deleteThread, getThread, updateThread } from "./threads.js";
 import { cancelRun, createRun, createThreadAndRun, getRun, listRuns, submitToolOutputsToRun, updateRun } from "./runs.js";
 import { createMessage, listMessages, updateMessage } from "./messages.js";
@@ -158,6 +158,9 @@ export interface AgentsOperations {
   ) => Promise<FileListResponseOutput>;
   /** Uploads a file for use by other operations. */
   uploadFile: (data: ReadableStream | NodeJS.ReadableStream, purpose: FilePurpose, fileName?: string, requestParams?: OptionalRequestParameters
+  ) => Promise<OpenAIFileOutput>
+    /** Uploads a file for use by other operations. */
+  uploadFileAndPoll: (data: ReadableStream | NodeJS.ReadableStream, purpose: FilePurpose, fileName?: string, pollingOptions?: PollingOptions , requestParams?: OptionalRequestParameters
   ) => Promise<OpenAIFileOutput>
   /** Delete a previously uploaded file. */
   deleteFile: (
@@ -358,6 +361,12 @@ function getAgents(context: Client): AgentsOperations {
         ...(requestParams as { [key: string]: any; }),
         contentType: "multipart/form-data"
       }),
+    uploadFileAndPoll: ( content: ReadableStream | NodeJS.ReadableStream, purpose: FilePurpose, fileName?: string, pollingOptions?: PollingOptions, requestParams?: OptionalRequestParameters) =>
+      uploadFileAndPoll(context, {
+        body: [{ name: "file" as const, body: content, filename: fileName }, { name: "purpose" as const, body: purpose }],
+        ...(requestParams as { [key: string]: any; }),
+        contentType: "multipart/form-data"
+      }, pollingOptions, requestParams),
     deleteFile: (fileId: string, requestParams?: OptionalRequestParameters) =>
       deleteFile(context, fileId, requestParams),
     getFile: (fileId: string, requestParams?: OptionalRequestParameters) =>

sdk/ai/ai-projects/src/agents/poller.ts

@@ -7,7 +7,7 @@ import { AbortSignalLike } from "@azure/abort-controller";
 import { PollingOptions } from "./customModels.js";
 
 interface PollResult {
-  status: string;
+  status?: string;
 }
 
 interface AgentsPollOperationState<T extends PollResult> extends PollOperationState<T> {

sdk/ai/ai-projects/test/public/agents/files.spec.ts

@@ -42,6 +42,18 @@ describe("Agents - files", () => {
     assert.isNotEmpty(file);
   });
 
+  it("should upload file and poll", async function () {
+    const fileContent = new ReadableStream({
+      start(controller) {
+        controller.enqueue(new TextEncoder().encode("fileContent"));
+        controller.close();
+      }
+    });
+    const file = await agents.uploadFileAndPoll(fileContent, "assistants", 1000, "fileName");
+    assert.notInclude(["uploaded","pending","running"], file.status);
+    assert.isNotEmpty(file);
+  });
+
   it("should delete file", async function () {
     const fileContent = new ReadableStream({
       start(controller) {