From 64a2d4eb4a6a54e046f9d80902fe211a757ae7b1 Mon Sep 17 00:00:00 2001
From: Richard Weerasinghe <richard@skyscrapr.io>
Date: Tue, 7 May 2024 09:50:55 +1200
Subject: [PATCH] update support for assistants with tool_resources, top_p and
 temperature

---
 openai/assistant.go      | 40 +++++++++++++++++++++++++++++-----------
 openai/assistant_test.go | 15 ++++++++++++---
 2 files changed, 41 insertions(+), 14 deletions(-)

diff --git a/openai/assistant.go b/openai/assistant.go
index 8a857bb..d664a1a 100644
--- a/openai/assistant.go
+++ b/openai/assistant.go
@@ -23,16 +23,18 @@ func (c *Client) Assistants() *AssistantsEndpoint {
 }
 
 type Assistant struct {
-	Id           string            `json:"id"`
-	Object       string            `json:"object"` // The object type, which is always assistant.
-	CreatedAt    int64             `json:"created_at"`
-	Name         *string           `json:"name"`
-	Description  *string           `json:"description"`
-	Model        string            `json:"model"`
-	Instructions *string           `json:"instructions"`
-	Tools        []AssistantTool   `json:"tools,omitempty"`
-	FileIds      []string          `json:"file_ids"`
-	MetaData     map[string]string `json:"metadata,omitempty"`
+	Id            string                  `json:"id"`
+	Object        string                  `json:"object"` // The object type, which is always assistant.
+	CreatedAt     int64                   `json:"created_at"`
+	Name          *string                 `json:"name"`
+	Description   *string                 `json:"description"`
+	Model         string                  `json:"model"`
+	Instructions  *string                 `json:"instructions"`
+	Tools         []AssistantTool         `json:"tools,omitempty"`
+	ToolResources *AssistantToolResources `json:"tool_resources,omitempty"`
+	MetaData      map[string]string       `json:"metadata,omitempty"`
+	Temperature   float64                 `json:"temperature,omitempty"`
+	TopP          float64                 `json:"top_p,omitempty"`
 }
 
 type Assistants struct {
@@ -59,10 +61,13 @@ type AssistantRequest struct {
 	Tools []AssistantTool `json:"tools,omitempty"`
 
 	// A list of file IDs attached to this assistant. There can be a maximum of 20 files attached to the assistant. Files are ordered by their creation date in ascending order.
-	FileIds []string `json:"file_ids,omitempty"`
+	ToolResources *AssistantToolResources `json:"tool_resources,omitempty"`
 
 	// Set of 16 key-value pairs that can be attached to an object. This can be useful for storing additional information about the object in a structured format. Keys can be a maximum of 64 characters long and va
 	MetaData map[string]string `json:"metadata,omitempty"`
+
+	Temperature float64 `json:"temperature,omitempty"`
+	TopP        float64 `json:"top_p,omitempty"`
 }
 
 type AssistantFile struct {
@@ -92,6 +97,19 @@ type AssistantTool struct {
 	} `json:"function,omitempty"`
 }
 
+type AssistantToolResources struct {
+	CodeInterpreter *struct {
+		FileIDs []string `json:"file_ids"`
+	} `json:"code_interpreter,omitempty"`
+	FileSearch *struct {
+		VectorStoreIDs []string `json:"vector_store_ids"`
+		VectorStore    *struct {
+			FileIDs  []string          `json:"file_ids"`
+			MetaData map[string]string `json:"metadata,omitempty"`
+		} `json:"vector_store,omitempty"`
+	} `json:"file_search,omitempty"`
+}
+
 // Create an assistant with a model and instructions.
 // [OpenAI Documentation]: https://platform.openai.com/docs/api-reference/assistants/createAssistant
 func (e *AssistantsEndpoint) CreateAssistant(req *AssistantRequest) (*Assistant, error) {
diff --git a/openai/assistant_test.go b/openai/assistant_test.go
index 8ef60fc..b1f06db 100644
--- a/openai/assistant_test.go
+++ b/openai/assistant_test.go
@@ -7,7 +7,7 @@ import (
 	"testing"
 
 	"github.com/skyscrapr/openai-sdk-go/openai"
-	"github.com/skyscrapr/openai-sdk-go/openai/test"
+	openai_test "github.com/skyscrapr/openai-sdk-go/openai/test"
 )
 
 func TestCreateAssistant(t *testing.T) {
@@ -32,8 +32,17 @@ func TestCreateAssistant(t *testing.T) {
 		Name:         &name,
 		Description:  &description,
 		Instructions: &instructions,
-		FileIds:      []string{"test_file_id_1", "test_file_id_2"},
-		MetaData:     map[string]string{"test_key_1": "test_value_1"},
+		Tools: []openai.AssistantTool{{
+			Type: "code_intepreter",
+		}},
+		ToolResources: &openai.AssistantToolResources{
+			CodeInterpreter: &struct {
+				FileIDs []string "json:\"file_ids\""
+			}{
+				FileIDs: []string{"file_1", "file_2"},
+			},
+		},
+		MetaData: map[string]string{"test_key_1": "test_value_1"},
 	}
 	_, err := client.Assistants().CreateAssistant(&req)
 	t.Helper()