[azopenaiassistants] Adding in streaming support for Assistants V2. (Azure#23093)

Adding in streaming support for all the Assistants endpoints.

Fixes Azure#22293

Author: richardpark-msft

sdk/ai/azopenaiassistants/autorest.md

@@ -212,6 +212,7 @@ directive:
       // POSTs
       $[threadsBase + "/messages/{messageId}"].post.parameters[2].schema["x-ms-client-name"] = "UpdateMessageBody";
       $[threadsBase + "/runs/{runId}"].post.parameters[2].schema["x-ms-client-name"] = "UpdateRunBody";
+      $[threadsBase + "/runs"].post.parameters[1]["x-ms-client-name"] = "CreateRunBody";
       $[threadsBase + "/runs/{runId}/submit_tool_outputs"].post.parameters[2].schema["x-ms-client-name"] = "SubmitToolOutputsToRunBody";
 
       // 

sdk/ai/azopenaiassistants/client.go

@@ -67,7 +67,7 @@ type GetFileContentOptions struct {
 //
 //   - fileID - The ID of the file to retrieve.
 //   - options - GetFileContentOptions contains the optional parameters for the Client.GetFileContent method.
-func (client *Client) GetFileContent(ctx context.Context, fileID string, options *GetFileContentOptions) (GetFileContentResponse, error) {
+func (client *Client) GetFileContent(ctx context.Context, fileID string, _ *GetFileContentOptions) (GetFileContentResponse, error) {
 	var err error
 
 	req, err := func() (*policy.Request, error) {

sdk/ai/azopenaiassistants/client_custom_files.go

@@ -14,7 +14,7 @@ import (
 
 // NewListMessagesPager returns a pager for messages associated with a thread.
 func (c *Client) NewListMessagesPager(threadID string, options *ListMessagesOptions) *runtime.Pager[ListMessagesResponse] {
-	nextPageFn := func(client *Client, ctx context.Context, opts *ListMessagesOptions) (ListMessagesResponse, error) {
+	nextPageFn := func(ctx context.Context, opts *ListMessagesOptions) (ListMessagesResponse, error) {
 		return c.internalListMessages(ctx, threadID, opts)
 	}
 	return newOpenAIPager(c, nextPageFn, options)
@@ -26,7 +26,7 @@ func (r ListMessagesResponse) hasMore() bool             { return *r.HasMore }
 
 // NewListAssistantsPager returns a pager for assistants.
 func (c *Client) NewListAssistantsPager(options *ListAssistantsOptions) *runtime.Pager[ListAssistantsResponse] {
-	nextPageFn := func(client *Client, ctx context.Context, opts *ListAssistantsOptions) (ListAssistantsResponse, error) {
+	nextPageFn := func(ctx context.Context, opts *ListAssistantsOptions) (ListAssistantsResponse, error) {
 		return c.internalListAssistants(ctx, opts)
 	}
 	return newOpenAIPager(c, nextPageFn, options)
@@ -38,7 +38,7 @@ func (r ListAssistantsResponse) hasMore() bool             { return *r.HasMore }
 
 // NewListRunStepsPager returns a pager for a Run's steps.
 func (c *Client) NewListRunStepsPager(threadID string, runID string, options *ListRunStepsOptions) *runtime.Pager[ListRunStepsResponse] {
-	nextPageFn := func(client *Client, ctx context.Context, opts *ListRunStepsOptions) (ListRunStepsResponse, error) {
+	nextPageFn := func(ctx context.Context, opts *ListRunStepsOptions) (ListRunStepsResponse, error) {
 		return c.internalListRunSteps(ctx, threadID, runID, opts)
 	}
 
@@ -51,7 +51,7 @@ func (r ListRunStepsResponse) hasMore() bool             { return *r.HasMore }
 
 // NewListRunsPager returns a pager for a Thread's runs.
 func (c *Client) NewListRunsPager(threadID string, options *ListRunsOptions) *runtime.Pager[ListRunsResponse] {
-	nextPageFn := func(client *Client, ctx context.Context, opts *ListRunsOptions) (ListRunsResponse, error) {
+	nextPageFn := func(ctx context.Context, opts *ListRunsOptions) (ListRunsResponse, error) {
 		return c.internalListRuns(ctx, threadID, opts)
 	}
 	return newOpenAIPager(c, nextPageFn, options)
@@ -63,7 +63,7 @@ func (r ListRunsResponse) hasMore() bool             { return *r.HasMore }
 
 // NewListVectorStoresPager returns a pager for a VectorStores.
 func (c *Client) NewListVectorStoresPager(options *ListVectorStoresOptions) *runtime.Pager[ListVectorStoresResponse] {
-	nextPageFn := func(client *Client, ctx context.Context, opts *ListVectorStoresOptions) (ListVectorStoresResponse, error) {
+	nextPageFn := func(ctx context.Context, opts *ListVectorStoresOptions) (ListVectorStoresResponse, error) {
 		return c.internalListVectorStores(ctx, opts)
 	}
 	return newOpenAIPager(c, nextPageFn, options)
@@ -75,7 +75,7 @@ func (r ListVectorStoresResponse) hasMore() bool             { return *r.HasMore
 
 // NewListVectorStoreFilesPager returns a pager for a vector store files.
 func (c *Client) NewListVectorStoreFilesPager(vectorStoreID string, options *ListVectorStoreFilesOptions) *runtime.Pager[ListVectorStoreFilesResponse] {
-	nextPageFn := func(client *Client, ctx context.Context, opts *ListVectorStoreFilesOptions) (ListVectorStoreFilesResponse, error) {
+	nextPageFn := func(ctx context.Context, opts *ListVectorStoreFilesOptions) (ListVectorStoreFilesResponse, error) {
 		return c.internalListVectorStoreFiles(ctx, vectorStoreID, opts)
 	}
 	return newOpenAIPager(c, nextPageFn, options)
@@ -87,7 +87,7 @@ func (r ListVectorStoreFilesResponse) hasMore() bool             { return *r.Has
 
 // NewListVectorStoreFileBatchFilesPager returns a pager for vector store files in a batch.
 func (c *Client) NewListVectorStoreFileBatchFilesPager(vectorStoreID string, batchID string, options *ListVectorStoreFileBatchFilesOptions) *runtime.Pager[ListVectorStoreFileBatchFilesResponse] {
-	nextPageFn := func(client *Client, ctx context.Context, opts *ListVectorStoreFileBatchFilesOptions) (ListVectorStoreFileBatchFilesResponse, error) {
+	nextPageFn := func(ctx context.Context, opts *ListVectorStoreFileBatchFilesOptions) (ListVectorStoreFileBatchFilesResponse, error) {
 		return c.internalListVectorStoreFileBatchFiles(ctx, vectorStoreID, batchID, opts)
 	}
 	return newOpenAIPager(c, nextPageFn, options)
@@ -112,7 +112,7 @@ func newOpenAIPager[ResponseT respType, OptionsT any, POptionsT interface {
 	updateAfter(after *string)
 }](
 	client *Client,
-	nextPageFn func(client *Client, ctx context.Context, opts POptionsT) (ResponseT, error),
+	nextPageFn func(ctx context.Context, opts POptionsT) (ResponseT, error),
 	options POptionsT) *runtime.Pager[ResponseT] {
 	var lastID *string
 
@@ -137,7 +137,7 @@ func newOpenAIPager[ResponseT respType, OptionsT any, POptionsT interface {
 
 			first = false
 
-			resp, err := nextPageFn(client, ctx, newOptions)
+			resp, err := nextPageFn(ctx, newOptions)
 
 			if err != nil {
 				var zero ResponseT

sdk/ai/azopenaiassistants/client_custom_pagers.go

@@ -0,0 +1,136 @@
+//go:build go1.18
+// +build go1.18
+
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+package azopenaiassistants
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
+)
+
+// API Reference: https://platform.openai.com/docs/api-reference/assistants-streaming
+
+// CreateThreadAndRunStreamResponse contains the response from [CreateThreadAndRunStream].
+type CreateThreadAndRunStreamResponse struct {
+	// Stream can be used to stream response events.
+	Stream *EventReader[StreamEvent]
+}
+
+// CreateThreadAndRunStreamOptions contains the optional parameters for [CreateThreadAndRunStream].
+type CreateThreadAndRunStreamOptions struct {
+	// for future expansion
+}
+
+// CreateThreadAndRunStream is the equivalent of [CreateThreadAndRun], but it returns a stream of responses instead of a
+// single response.
+func (client *Client) CreateThreadAndRunStream(ctx context.Context, body CreateAndRunThreadBody, _ *CreateThreadAndRunStreamOptions) (CreateThreadAndRunStreamResponse, error) {
+	// enable streaming.
+	// https: //platform.openai.com/docs/api-reference/runs/createThreadAndRun#runs-createthreadandrun-stream
+	body.stream = to.Ptr(true)
+
+	var err error
+
+	req, err := client.createThreadAndRunCreateRequest(ctx, body, nil)
+	if err != nil {
+		return CreateThreadAndRunStreamResponse{}, err
+	}
+
+	runtime.SkipBodyDownload(req) // we'll handle this.
+
+	httpResp, err := client.internal.Pipeline().Do(req)
+	if err != nil {
+		return CreateThreadAndRunStreamResponse{}, err
+	}
+	if !runtime.HasStatusCode(httpResp, http.StatusOK) {
+		err = runtime.NewResponseError(httpResp)
+		return CreateThreadAndRunStreamResponse{}, err
+	}
+
+	return CreateThreadAndRunStreamResponse{
+		Stream: newEventReader(httpResp.Body, unmarshalStreamEvent),
+	}, err
+}
+
+// API Reference: https://platform.openai.com/docs/api-reference/runs/createRun
+
+// CreateRunStreamResponse contains the response from [CreateRunStream].
+type CreateRunStreamResponse struct {
+	// Stream can be used to stream response events.
+	Stream *EventReader[StreamEvent]
+}
+
+// CreateRunStreamOptions contains the optional parameters for [CreateRunStream].
+type CreateRunStreamOptions struct {
+	// for future expansion
+}
+
+// CreateRunStream is the equivalent of [CreateRun], but it returns a stream of responses instead of a
+// single response.
+func (client *Client) CreateRunStream(ctx context.Context, threadID string, body CreateRunBody, _ *CreateRunStreamOptions) (CreateRunStreamResponse, error) {
+	var err error
+	body.stream = to.Ptr(true)
+
+	req, err := client.createRunCreateRequest(ctx, threadID, body, nil)
+	if err != nil {
+		return CreateRunStreamResponse{}, err
+	}
+
+	runtime.SkipBodyDownload(req) // we'll handle this.
+
+	httpResp, err := client.internal.Pipeline().Do(req)
+	if err != nil {
+		return CreateRunStreamResponse{}, err
+	}
+	if !runtime.HasStatusCode(httpResp, http.StatusOK) {
+		err = runtime.NewResponseError(httpResp)
+		return CreateRunStreamResponse{}, err
+	}
+
+	return CreateRunStreamResponse{
+		Stream: newEventReader[StreamEvent](httpResp.Body, unmarshalStreamEvent),
+	}, err
+}
+
+// SubmitToolOutputsToRunStreamResponse contains the response from [SubmitToolOutputsToRunStream].
+type SubmitToolOutputsToRunStreamResponse struct {
+	// Stream can be used to stream response events.
+	Stream *EventReader[StreamEvent]
+}
+
+// SubmitToolOutputsToRunStreamOptions contains the optional parameters for [SubmitToolOutputsToRunStream].
+type SubmitToolOutputsToRunStreamOptions struct {
+}
+
+// SubmitToolOutputsToRunStream is the equivalent of [SubmitToolOutputsToRun], but it returns a stream of responses instead of a
+// single response.
+func (client *Client) SubmitToolOutputsToRunStream(ctx context.Context, threadID string, runID string, body SubmitToolOutputsToRunBody, options *SubmitToolOutputsToRunOptions) (SubmitToolOutputsToRunStreamResponse, error) {
+	var err error
+
+	body.stream = to.Ptr(true)
+
+	req, err := client.submitToolOutputsToRunCreateRequest(ctx, threadID, runID, body, options)
+	if err != nil {
+		return SubmitToolOutputsToRunStreamResponse{}, err
+	}
+
+	runtime.SkipBodyDownload(req) // we'll handle this.
+
+	httpResp, err := client.internal.Pipeline().Do(req)
+	if err != nil {
+		return SubmitToolOutputsToRunStreamResponse{}, err
+	}
+	if !runtime.HasStatusCode(httpResp, http.StatusOK) {
+		err = runtime.NewResponseError(httpResp)
+		return SubmitToolOutputsToRunStreamResponse{}, err
+	}
+
+	return SubmitToolOutputsToRunStreamResponse{
+		Stream: newEventReader[StreamEvent](httpResp.Body, unmarshalStreamEvent),
+	}, nil
+}