feat: completion async (#10)

Authored-by: Ann Vasileva <[email protected]>

auth.go

@@ -0,0 +1,31 @@
+package yandexgpt
+
+import (
+	"context"
+	"net/http"
+)
+
+const getIAMUrl = "https://iam.api.cloud.yandex.net/iam/v1/tokens"
+
+// Updates IAM token.
+//
+// Always call it before creating a request.
+//
+// If you will use it when API key is specified, method CreateRequest(...) will always use API key.
+func (c *YandexGPTClient) GetIAMToken(ctx context.Context) error {
+	iamRq := YandexIAMRequest{OAuthToken: c.config.OAuthToken}
+	req, err := c.newRequest(ctx, http.MethodPost, getIAMUrl, iamRq)
+	if err != nil {
+		return err
+	}
+
+	var resp YandexIAMResponse
+	err = c.sendRequest(req, &resp)
+	if err != nil {
+		return err
+	}
+
+	//set new IAMToken
+	c.config.SetIAMToken(resp.IAMToken)
+	return nil
+}

client.go

@@ -53,11 +53,25 @@ func NewYandexGPTClientWithAPIKey(
 	}
 }
 
+// Creates new YandexGPT Client.
+//
+// You will need to specify your own OAuth key.
+func NewYandexGPTClientWithOAuthToken(
+	oauthToken string,
+) *YandexGPTClient {
+	config := NewYandexGPTClientConfigWithOAuthToken(oauthToken)
+
+	return &YandexGPTClient{
+		config:         config,
+		requestBuilder: internal.NewRequestBuilder(),
+	}
+}
+
 func (c *YandexGPTClient) newRequest(
 	ctx context.Context,
 	method,
 	url string,
-	req YandexGPTRequest,
+	req any,
 ) (*http.Request, error) {
 	request, err := c.requestBuilder.Build(ctx, method, url, req)
 	if err != nil {
@@ -94,15 +108,16 @@ func (c *YandexGPTClient) sendRequest(
 func (c *YandexGPTClient) setHeaders(request *http.Request) {
 	request.Header.Set("Content-Type", "application/json")
 
-	if c.config.ApiKey != "" {
+	if c.config.IAMToken != "" {
 		request.Header.Set(
 			"Authorization",
-			fmt.Sprintf("Api-Key %s", c.config.ApiKey),
+			fmt.Sprintf("Bearer %s", c.config.IAMToken),
 		)
-	} else if c.config.IAMToken != "" {
+	}
+	if c.config.ApiKey != "" {
 		request.Header.Set(
 			"Authorization",
-			fmt.Sprintf("Bearer %s", c.config.IAMToken),
+			fmt.Sprintf("Bearer %s", c.config.ApiKey),
 		)
 	}
 }
@@ -120,35 +135,3 @@ func (c *YandexGPTClient) handleResponseError(response *http.Response) error {
 		errResponse.Error.Message,
 	)
 }
-
-// Updates IAM token.
-//
-// Always call it before creating a request.
-//
-// If you will use it when API key is specified, method CreateRequest(...) will always use API key.
-func (c *YandexGPTClient) UpdateIAMToken(iamToken string) {
-	c.config.updateIAMToken(iamToken)
-}
-
-// Creates request to YandexGPT.
-//
-// If you're using IAM token, make sure to update client's IAM token by calling
-// UpdateIAMToken(iamToken string) method first.
-//
-// Keep in mind that if for some strange reason you provided  API key and IAM token to the client,
-// this method will use API key
-func (c *YandexGPTClient) CreateRequest(
-	ctx context.Context,
-	request YandexGPTRequest,
-) (response YandexGPTResponse, err error) {
-	//TODO:
-	//1. Validate Request
-
-	req, err := c.newRequest(ctx, http.MethodPost, c.config.BaseURL, request)
-	if err != nil {
-		return
-	}
-	err = c.sendRequest(req, &response)
-
-	return
-}

completions.go

@@ -0,0 +1,54 @@
+package yandexgpt
+
+import (
+	"context"
+	"net/http"
+)
+
+const completionURL = "https://llm.api.cloud.yandex.net/foundationModels/v1"
+
+// Get completion from YandexGPT.
+//
+// If you're using IAM token, make sure to update client's IAM token by calling
+// GetIAMToken(iamToken string) method first.
+//
+// Keep in mind that if for some strange reason you provided  API key and IAM token to the client,
+// this method will use API key.
+func (c *YandexGPTClient) GetCompletion(
+	ctx context.Context,
+	request YandexGPTRequest,
+) (response YandexGPTResponse, err error) {
+
+	endpoint := completionURL + "/completion"
+
+	req, err := c.newRequest(ctx, http.MethodPost, endpoint, request)
+	if err != nil {
+		return
+	}
+	err = c.sendRequest(req, &response)
+
+	return
+}
+
+// Get completion from YandexGPT with async method.
+//
+// If you're using IAM token, make sure to update client's IAM token by calling
+// GetIAMToken(iamToken string) method first.
+//
+// Keep in mind that if for some strange reason you provided  API key and IAM token to the client,
+// this method will use API key.
+func (c *YandexGPTClient) RunCompletionAsync(
+	ctx context.Context,
+	request YandexGPTRequest,
+) (response YandexCompletionResponse, err error) {
+	endpoint := completionURL + "/completionAsync"
+
+	req, err := c.newRequest(ctx, http.MethodPost, endpoint, request)
+	if err != nil {
+		return
+	}
+
+	err = c.sendRequest(req, &response)
+
+	return
+}

config.go

@@ -2,20 +2,15 @@ package yandexgpt
 
 import "net/http"
 
-const (
-	completionURL = "https://llm.api.cloud.yandex.net/foundationModels/v1/completion"
-)
-
 type YandexGPTClientConfig struct {
+	OAuthToken string
 	ApiKey     string
 	IAMToken   string
-	BaseURL    string
 	HTTPClient *http.Client
 }
 
 func NewYandexGPTClientConfig() *YandexGPTClientConfig {
 	return &YandexGPTClientConfig{
-		BaseURL:    completionURL,
 		HTTPClient: &http.Client{},
 	}
 }
@@ -25,7 +20,6 @@ func NewYandexGPTClientConfigWithIAMToken(
 ) *YandexGPTClientConfig {
 	return &YandexGPTClientConfig{
 		IAMToken:   iamToken,
-		BaseURL:    completionURL,
 		HTTPClient: &http.Client{},
 	}
 }
@@ -35,11 +29,22 @@ func NewYandexGPTClientConfigWithAPIKey(
 ) *YandexGPTClientConfig {
 	return &YandexGPTClientConfig{
 		ApiKey:     apiKey,
-		BaseURL:    completionURL,
 		HTTPClient: &http.Client{},
 	}
 }
 
-func (c *YandexGPTClientConfig) updateIAMToken(iamToken string) {
+func NewYandexGPTClientConfigWithOAuthToken(
+	oauthToken string,
+) *YandexGPTClientConfig {
+	return &YandexGPTClientConfig{
+		OAuthToken: oauthToken,
+		HTTPClient: &http.Client{},
+	}
+}
+
+// Setter for IAM token in config.
+//
+// Use it for manually updating token in config.
+func (c *YandexGPTClientConfig) SetIAMToken(iamToken string) {
 	c.IAMToken = iamToken
 }

error.go

@@ -1,7 +1,8 @@
 package yandexgpt
 
 type YandexGPTError struct {
-	HTTPCode   int    `json:"httpCode"`
-	Message    string `json:"message"`
-	HTTPStatus string `json:"httpStatus"`
+	HTTPCode   int             `json:"httpCode"`
+	Message    string          `json:"message"`
+	HTTPStatus string          `json:"httpStatus"`
+	Details    DetailsResponse `json:"error.details"`
 }

examples/completion/main.go

@@ -0,0 +1,45 @@
+package main
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/sheeiavellie/go-yandexgpt"
+)
+
+func main() {
+	client := yandexgpt.NewYandexGPTClientWithOAuthToken("OAUTH_TOKEN")
+
+	// get, update and set iam token
+	ctx := context.Background()
+	err := client.GetIAMToken(ctx)
+	if err != nil {
+		panic(err)
+	}
+
+	request := yandexgpt.YandexGPTRequest{
+		ModelURI: yandexgpt.MakeModelURI("CATALOG_ID", yandexgpt.YandexGPT4ModelLite),
+		CompletionOptions: yandexgpt.YandexGPTCompletionOptions{
+			Stream:      false,
+			Temperature: 0.7,
+			MaxTokens:   100,
+		},
+		Messages: []yandexgpt.YandexGPTMessage{
+			{
+				Role: yandexgpt.YandexGPTMessageRoleSystem,
+				Text: "Every time you get ONE you answer just TWO",
+			},
+			{
+				Role: yandexgpt.YandexGPTMessageRoleUser,
+				Text: "ONE",
+			},
+		},
+	}
+
+	response, err := client.GetCompletion(ctx, request)
+	if err != nil {
+		panic(err)
+	}
+
+	fmt.Println(response.Result.Alternatives[0].Message.Text)
+}

examples/completionAsync/main.go

@@ -0,0 +1,61 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"time"
+
+	"github.com/sheeiavellie/go-yandexgpt"
+)
+
+func main() {
+	client := yandexgpt.NewYandexGPTClientWithOAuthToken("OAUTH_TOKEN")
+
+	// get, update and set iam token
+	ctx := context.Background()
+	err := client.GetIAMToken(ctx)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	request := yandexgpt.YandexGPTRequest{
+		ModelURI: yandexgpt.MakeModelURI("CATALOG_ID", yandexgpt.YandexGPT4ModelLite),
+		CompletionOptions: yandexgpt.YandexGPTCompletionOptions{
+			Stream:      false,
+			Temperature: 0.7,
+			MaxTokens:   100,
+		},
+		Messages: []yandexgpt.YandexGPTMessage{
+			{
+				Role: yandexgpt.YandexGPTMessageRoleSystem,
+				Text: "Every time you get ONE you answer just TWO",
+			},
+			{
+				Role: yandexgpt.YandexGPTMessageRoleUser,
+				Text: "ONE",
+			},
+		},
+	}
+
+	response, err := client.RunCompletionAsync(ctx, request)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	isCompleted := false
+	for !isCompleted {
+		status, err := client.GetOperationStatus(ctx, response.ID)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		if status.Done {
+			isCompleted = true
+			fmt.Println("\n Chat answer: \n")
+			fmt.Println(status.Response.Alternatives[0].Message.Text)
+		} else {
+			time.Sleep(5 * time.Second)
+		}
+	}
+}

message.go

@@ -1,6 +1,24 @@
 package yandexgpt
 
 type YandexGPTMessage struct {
-	Role yandexGPTRole `json:"role"`
-	Text string        `json:"text"`
+	Role         yandexGPTRole `json:"role"`
+	Text         string        `json:"text"`
+	ToolCallList *ToolCallList `json:"toolCallList"`
+}
+
+type YandexToolResultList struct {
+	ToolResults []YandexFunctionResult `json:"toolResults"`
+}
+
+type YandexFunctionResult struct {
+	Name    string `json:"name"`
+	Content string `json:"content"`
+}
+
+type ToolCallList struct {
+	ToolCalls []FunctionCall `json:"toolCalls"`
+}
+type FunctionCall struct {
+	Name      string `json:"name"`
+	Arguments any    `json:"arguments"`
 }

operations.go

@@ -0,0 +1,24 @@
+package yandexgpt
+
+import (
+	"context"
+	"net/http"
+)
+
+const operationURL = "https://operation.api.cloud.yandex.net/operations"
+
+func (c *YandexGPTClient) GetOperationStatus(
+	ctx context.Context,
+	operationID string,
+) (response OperationResponse, err error) {
+	endpoint := operationURL + "/" + operationID
+
+	req, err := c.newRequest(ctx, http.MethodGet, endpoint, nil)
+	if err != nil {
+		return
+	}
+
+	err = c.sendRequest(req, &response)
+
+	return
+}