feat: provide vector store (#772)
* implement vectore store feature * fix after integration testing * fix golint error * improve test to increare code coverage * fix golint anc code coverage problem * add tool_resource in assistant response * chore: code style * feat: use pagination param * feat: use pagination param * test: use pagination param * test: rm unused code --------- Co-authored-by: Denny Depok <61371551+kodernubie@users.noreply.github.com> Co-authored-by: eric.p <eric.p>
This commit is contained in:
Peng Guan-Cheng
GitHub
50
assistant.go
50
assistant.go
@@ -14,16 +14,17 @@ const (
|
|||||||
)
|
)
|
||||||
|
|
||||||
type Assistant struct {
|
type Assistant struct {
|
||||||
ID string `json:"id"`
|
ID string `json:"id"`
|
||||||
Object string `json:"object"`
|
Object string `json:"object"`
|
||||||
CreatedAt int64 `json:"created_at"`
|
CreatedAt int64 `json:"created_at"`
|
||||||
Name *string `json:"name,omitempty"`
|
Name *string `json:"name,omitempty"`
|
||||||
Description *string `json:"description,omitempty"`
|
Description *string `json:"description,omitempty"`
|
||||||
Model string `json:"model"`
|
Model string `json:"model"`
|
||||||
Instructions *string `json:"instructions,omitempty"`
|
Instructions *string `json:"instructions,omitempty"`
|
||||||
Tools []AssistantTool `json:"tools"`
|
Tools []AssistantTool `json:"tools"`
|
||||||
FileIDs []string `json:"file_ids,omitempty"`
|
FileIDs []string `json:"file_ids,omitempty"`
|
||||||
Metadata map[string]any `json:"metadata,omitempty"`
|
Metadata map[string]any `json:"metadata,omitempty"`
|
||||||
|
ToolResources *AssistantToolResource `json:"tool_resources,omitempty"`
|
||||||
|
|
||||||
httpHeader
|
httpHeader
|
||||||
}
|
}
|
||||||
@@ -34,6 +35,7 @@ const (
|
|||||||
AssistantToolTypeCodeInterpreter AssistantToolType = "code_interpreter"
|
AssistantToolTypeCodeInterpreter AssistantToolType = "code_interpreter"
|
||||||
AssistantToolTypeRetrieval AssistantToolType = "retrieval"
|
AssistantToolTypeRetrieval AssistantToolType = "retrieval"
|
||||||
AssistantToolTypeFunction AssistantToolType = "function"
|
AssistantToolTypeFunction AssistantToolType = "function"
|
||||||
|
AssistantToolTypeFileSearch AssistantToolType = "file_search"
|
||||||
)
|
)
|
||||||
|
|
||||||
type AssistantTool struct {
|
type AssistantTool struct {
|
||||||
@@ -41,19 +43,33 @@ type AssistantTool struct {
|
|||||||
Function *FunctionDefinition `json:"function,omitempty"`
|
Function *FunctionDefinition `json:"function,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type AssistantToolFileSearch struct {
|
||||||
|
VectorStoreIDs []string `json:"vector_store_ids"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type AssistantToolCodeInterpreter struct {
|
||||||
|
FileIDs []string `json:"file_ids"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type AssistantToolResource struct {
|
||||||
|
FileSearch *AssistantToolFileSearch `json:"file_search,omitempty"`
|
||||||
|
CodeInterpreter *AssistantToolCodeInterpreter `json:"code_interpreter,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
// AssistantRequest provides the assistant request parameters.
|
// AssistantRequest provides the assistant request parameters.
|
||||||
// When modifying the tools the API functions as the following:
|
// When modifying the tools the API functions as the following:
|
||||||
// If Tools is undefined, no changes are made to the Assistant's tools.
|
// If Tools is undefined, no changes are made to the Assistant's tools.
|
||||||
// If Tools is empty slice it will effectively delete all of the Assistant's tools.
|
// If Tools is empty slice it will effectively delete all of the Assistant's tools.
|
||||||
// If Tools is populated, it will replace all of the existing Assistant's tools with the provided tools.
|
// If Tools is populated, it will replace all of the existing Assistant's tools with the provided tools.
|
||||||
type AssistantRequest struct {
|
type AssistantRequest struct {
|
||||||
Model string `json:"model"`
|
Model string `json:"model"`
|
||||||
Name *string `json:"name,omitempty"`
|
Name *string `json:"name,omitempty"`
|
||||||
Description *string `json:"description,omitempty"`
|
Description *string `json:"description,omitempty"`
|
||||||
Instructions *string `json:"instructions,omitempty"`
|
Instructions *string `json:"instructions,omitempty"`
|
||||||
Tools []AssistantTool `json:"-"`
|
Tools []AssistantTool `json:"-"`
|
||||||
FileIDs []string `json:"file_ids,omitempty"`
|
FileIDs []string `json:"file_ids,omitempty"`
|
||||||
Metadata map[string]any `json:"metadata,omitempty"`
|
Metadata map[string]any `json:"metadata,omitempty"`
|
||||||
|
ToolResources *AssistantToolResource `json:"tool_resources,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// MarshalJSON provides a custom marshaller for the assistant request to handle the API use cases
|
// MarshalJSON provides a custom marshaller for the assistant request to handle the API use cases
|
||||||
|
|||||||
@@ -24,7 +24,7 @@ const (
|
|||||||
|
|
||||||
const AzureAPIKeyHeader = "api-key"
|
const AzureAPIKeyHeader = "api-key"
|
||||||
|
|
||||||
const defaultAssistantVersion = "v1" // This will be deprecated by the end of 2024.
|
const defaultAssistantVersion = "v2" // upgrade to v2 to support vector store
|
||||||
|
|
||||||
// ClientConfig is a configuration of a client.
|
// ClientConfig is a configuration of a client.
|
||||||
type ClientConfig struct {
|
type ClientConfig struct {
|
||||||
|
|||||||
345
vector_store.go
Normal file
345
vector_store.go
Normal file
@@ -0,0 +1,345 @@
|
|||||||
|
package openai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"net/url"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
vectorStoresSuffix = "/vector_stores"
|
||||||
|
vectorStoresFilesSuffix = "/files"
|
||||||
|
vectorStoresFileBatchesSuffix = "/file_batches"
|
||||||
|
)
|
||||||
|
|
||||||
|
type VectorStoreFileCount struct {
|
||||||
|
InProgress int `json:"in_progress"`
|
||||||
|
Completed int `json:"completed"`
|
||||||
|
Failed int `json:"failed"`
|
||||||
|
Cancelled int `json:"cancelled"`
|
||||||
|
Total int `json:"total"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type VectorStore struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Object string `json:"object"`
|
||||||
|
CreatedAt int64 `json:"created_at"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
UsageBytes int `json:"usage_bytes"`
|
||||||
|
FileCounts VectorStoreFileCount `json:"file_counts"`
|
||||||
|
Status string `json:"status"`
|
||||||
|
ExpiresAfter *VectorStoreExpires `json:"expires_after"`
|
||||||
|
ExpiresAt *int `json:"expires_at"`
|
||||||
|
Metadata map[string]any `json:"metadata"`
|
||||||
|
|
||||||
|
httpHeader
|
||||||
|
}
|
||||||
|
|
||||||
|
type VectorStoreExpires struct {
|
||||||
|
Anchor string `json:"anchor"`
|
||||||
|
Days int `json:"days"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// VectorStoreRequest provides the vector store request parameters.
|
||||||
|
type VectorStoreRequest struct {
|
||||||
|
Name string `json:"name,omitempty"`
|
||||||
|
FileIDs []string `json:"file_ids,omitempty"`
|
||||||
|
ExpiresAfter *VectorStoreExpires `json:"expires_after,omitempty"`
|
||||||
|
Metadata map[string]any `json:"metadata,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// VectorStoresList is a list of vector store.
|
||||||
|
type VectorStoresList struct {
|
||||||
|
VectorStores []VectorStore `json:"data"`
|
||||||
|
LastID *string `json:"last_id"`
|
||||||
|
FirstID *string `json:"first_id"`
|
||||||
|
HasMore bool `json:"has_more"`
|
||||||
|
httpHeader
|
||||||
|
}
|
||||||
|
|
||||||
|
type VectorStoreDeleteResponse struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Object string `json:"object"`
|
||||||
|
Deleted bool `json:"deleted"`
|
||||||
|
|
||||||
|
httpHeader
|
||||||
|
}
|
||||||
|
|
||||||
|
type VectorStoreFile struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Object string `json:"object"`
|
||||||
|
CreatedAt int64 `json:"created_at"`
|
||||||
|
VectorStoreID string `json:"vector_store_id"`
|
||||||
|
UsageBytes int `json:"usage_bytes"`
|
||||||
|
Status string `json:"status"`
|
||||||
|
|
||||||
|
httpHeader
|
||||||
|
}
|
||||||
|
|
||||||
|
type VectorStoreFileRequest struct {
|
||||||
|
FileID string `json:"file_id"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type VectorStoreFilesList struct {
|
||||||
|
VectorStoreFiles []VectorStoreFile `json:"data"`
|
||||||
|
|
||||||
|
httpHeader
|
||||||
|
}
|
||||||
|
|
||||||
|
type VectorStoreFileBatch struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Object string `json:"object"`
|
||||||
|
CreatedAt int64 `json:"created_at"`
|
||||||
|
VectorStoreID string `json:"vector_store_id"`
|
||||||
|
Status string `json:"status"`
|
||||||
|
FileCounts VectorStoreFileCount `json:"file_counts"`
|
||||||
|
|
||||||
|
httpHeader
|
||||||
|
}
|
||||||
|
|
||||||
|
type VectorStoreFileBatchRequest struct {
|
||||||
|
FileIDs []string `json:"file_ids"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// CreateVectorStore creates a new vector store.
|
||||||
|
func (c *Client) CreateVectorStore(ctx context.Context, request VectorStoreRequest) (response VectorStore, err error) {
|
||||||
|
req, _ := c.newRequest(
|
||||||
|
ctx,
|
||||||
|
http.MethodPost,
|
||||||
|
c.fullURL(vectorStoresSuffix),
|
||||||
|
withBody(request),
|
||||||
|
withBetaAssistantVersion(c.config.AssistantVersion),
|
||||||
|
)
|
||||||
|
|
||||||
|
err = c.sendRequest(req, &response)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// RetrieveVectorStore retrieves an vector store.
|
||||||
|
func (c *Client) RetrieveVectorStore(
|
||||||
|
ctx context.Context,
|
||||||
|
vectorStoreID string,
|
||||||
|
) (response VectorStore, err error) {
|
||||||
|
urlSuffix := fmt.Sprintf("%s/%s", vectorStoresSuffix, vectorStoreID)
|
||||||
|
req, _ := c.newRequest(ctx, http.MethodGet, c.fullURL(urlSuffix),
|
||||||
|
withBetaAssistantVersion(c.config.AssistantVersion))
|
||||||
|
|
||||||
|
err = c.sendRequest(req, &response)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// ModifyVectorStore modifies a vector store.
|
||||||
|
func (c *Client) ModifyVectorStore(
|
||||||
|
ctx context.Context,
|
||||||
|
vectorStoreID string,
|
||||||
|
request VectorStoreRequest,
|
||||||
|
) (response VectorStore, err error) {
|
||||||
|
urlSuffix := fmt.Sprintf("%s/%s", vectorStoresSuffix, vectorStoreID)
|
||||||
|
req, _ := c.newRequest(ctx, http.MethodPost, c.fullURL(urlSuffix), withBody(request),
|
||||||
|
withBetaAssistantVersion(c.config.AssistantVersion))
|
||||||
|
|
||||||
|
err = c.sendRequest(req, &response)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteVectorStore deletes an vector store.
|
||||||
|
func (c *Client) DeleteVectorStore(
|
||||||
|
ctx context.Context,
|
||||||
|
vectorStoreID string,
|
||||||
|
) (response VectorStoreDeleteResponse, err error) {
|
||||||
|
urlSuffix := fmt.Sprintf("%s/%s", vectorStoresSuffix, vectorStoreID)
|
||||||
|
req, _ := c.newRequest(ctx, http.MethodDelete, c.fullURL(urlSuffix),
|
||||||
|
withBetaAssistantVersion(c.config.AssistantVersion))
|
||||||
|
|
||||||
|
err = c.sendRequest(req, &response)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListVectorStores Lists the currently available vector store.
|
||||||
|
func (c *Client) ListVectorStores(
|
||||||
|
ctx context.Context,
|
||||||
|
pagination Pagination,
|
||||||
|
) (response VectorStoresList, err error) {
|
||||||
|
urlValues := url.Values{}
|
||||||
|
|
||||||
|
if pagination.After != nil {
|
||||||
|
urlValues.Add("after", *pagination.After)
|
||||||
|
}
|
||||||
|
if pagination.Order != nil {
|
||||||
|
urlValues.Add("order", *pagination.Order)
|
||||||
|
}
|
||||||
|
if pagination.Limit != nil {
|
||||||
|
urlValues.Add("limit", fmt.Sprintf("%d", *pagination.Limit))
|
||||||
|
}
|
||||||
|
if pagination.Before != nil {
|
||||||
|
urlValues.Add("before", *pagination.Before)
|
||||||
|
}
|
||||||
|
|
||||||
|
encodedValues := ""
|
||||||
|
if len(urlValues) > 0 {
|
||||||
|
encodedValues = "?" + urlValues.Encode()
|
||||||
|
}
|
||||||
|
|
||||||
|
urlSuffix := fmt.Sprintf("%s%s", vectorStoresSuffix, encodedValues)
|
||||||
|
req, _ := c.newRequest(ctx, http.MethodGet, c.fullURL(urlSuffix),
|
||||||
|
withBetaAssistantVersion(c.config.AssistantVersion))
|
||||||
|
|
||||||
|
err = c.sendRequest(req, &response)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// CreateVectorStoreFile creates a new vector store file.
|
||||||
|
func (c *Client) CreateVectorStoreFile(
|
||||||
|
ctx context.Context,
|
||||||
|
vectorStoreID string,
|
||||||
|
request VectorStoreFileRequest,
|
||||||
|
) (response VectorStoreFile, err error) {
|
||||||
|
urlSuffix := fmt.Sprintf("%s/%s%s", vectorStoresSuffix, vectorStoreID, vectorStoresFilesSuffix)
|
||||||
|
req, _ := c.newRequest(ctx, http.MethodPost, c.fullURL(urlSuffix),
|
||||||
|
withBody(request),
|
||||||
|
withBetaAssistantVersion(c.config.AssistantVersion))
|
||||||
|
|
||||||
|
err = c.sendRequest(req, &response)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// RetrieveVectorStoreFile retrieves a vector store file.
|
||||||
|
func (c *Client) RetrieveVectorStoreFile(
|
||||||
|
ctx context.Context,
|
||||||
|
vectorStoreID string,
|
||||||
|
fileID string,
|
||||||
|
) (response VectorStoreFile, err error) {
|
||||||
|
urlSuffix := fmt.Sprintf("%s/%s%s/%s", vectorStoresSuffix, vectorStoreID, vectorStoresFilesSuffix, fileID)
|
||||||
|
req, _ := c.newRequest(ctx, http.MethodGet, c.fullURL(urlSuffix),
|
||||||
|
withBetaAssistantVersion(c.config.AssistantVersion))
|
||||||
|
|
||||||
|
err = c.sendRequest(req, &response)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// DeleteVectorStoreFile deletes an existing file.
|
||||||
|
func (c *Client) DeleteVectorStoreFile(
|
||||||
|
ctx context.Context,
|
||||||
|
vectorStoreID string,
|
||||||
|
fileID string,
|
||||||
|
) (err error) {
|
||||||
|
urlSuffix := fmt.Sprintf("%s/%s%s/%s", vectorStoresSuffix, vectorStoreID, vectorStoresFilesSuffix, fileID)
|
||||||
|
req, _ := c.newRequest(ctx, http.MethodDelete, c.fullURL(urlSuffix),
|
||||||
|
withBetaAssistantVersion(c.config.AssistantVersion))
|
||||||
|
|
||||||
|
err = c.sendRequest(req, nil)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListVectorStoreFiles Lists the currently available files for a vector store.
|
||||||
|
func (c *Client) ListVectorStoreFiles(
|
||||||
|
ctx context.Context,
|
||||||
|
vectorStoreID string,
|
||||||
|
pagination Pagination,
|
||||||
|
) (response VectorStoreFilesList, err error) {
|
||||||
|
urlValues := url.Values{}
|
||||||
|
if pagination.After != nil {
|
||||||
|
urlValues.Add("after", *pagination.After)
|
||||||
|
}
|
||||||
|
if pagination.Limit != nil {
|
||||||
|
urlValues.Add("limit", fmt.Sprintf("%d", *pagination.Limit))
|
||||||
|
}
|
||||||
|
if pagination.Before != nil {
|
||||||
|
urlValues.Add("before", *pagination.Before)
|
||||||
|
}
|
||||||
|
if pagination.Order != nil {
|
||||||
|
urlValues.Add("order", *pagination.Order)
|
||||||
|
}
|
||||||
|
|
||||||
|
encodedValues := ""
|
||||||
|
if len(urlValues) > 0 {
|
||||||
|
encodedValues = "?" + urlValues.Encode()
|
||||||
|
}
|
||||||
|
|
||||||
|
urlSuffix := fmt.Sprintf("%s/%s%s%s", vectorStoresSuffix, vectorStoreID, vectorStoresFilesSuffix, encodedValues)
|
||||||
|
req, _ := c.newRequest(ctx, http.MethodGet, c.fullURL(urlSuffix),
|
||||||
|
withBetaAssistantVersion(c.config.AssistantVersion))
|
||||||
|
|
||||||
|
err = c.sendRequest(req, &response)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// CreateVectorStoreFileBatch creates a new vector store file batch.
|
||||||
|
func (c *Client) CreateVectorStoreFileBatch(
|
||||||
|
ctx context.Context,
|
||||||
|
vectorStoreID string,
|
||||||
|
request VectorStoreFileBatchRequest,
|
||||||
|
) (response VectorStoreFileBatch, err error) {
|
||||||
|
urlSuffix := fmt.Sprintf("%s/%s%s", vectorStoresSuffix, vectorStoreID, vectorStoresFileBatchesSuffix)
|
||||||
|
req, _ := c.newRequest(ctx, http.MethodPost, c.fullURL(urlSuffix),
|
||||||
|
withBody(request),
|
||||||
|
withBetaAssistantVersion(c.config.AssistantVersion))
|
||||||
|
|
||||||
|
err = c.sendRequest(req, &response)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// RetrieveVectorStoreFileBatch retrieves a vector store file batch.
|
||||||
|
func (c *Client) RetrieveVectorStoreFileBatch(
|
||||||
|
ctx context.Context,
|
||||||
|
vectorStoreID string,
|
||||||
|
batchID string,
|
||||||
|
) (response VectorStoreFileBatch, err error) {
|
||||||
|
urlSuffix := fmt.Sprintf("%s/%s%s/%s", vectorStoresSuffix, vectorStoreID, vectorStoresFileBatchesSuffix, batchID)
|
||||||
|
req, _ := c.newRequest(ctx, http.MethodGet, c.fullURL(urlSuffix),
|
||||||
|
withBetaAssistantVersion(c.config.AssistantVersion))
|
||||||
|
|
||||||
|
err = c.sendRequest(req, &response)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// CancelVectorStoreFileBatch cancel a new vector store file batch.
|
||||||
|
func (c *Client) CancelVectorStoreFileBatch(
|
||||||
|
ctx context.Context,
|
||||||
|
vectorStoreID string,
|
||||||
|
batchID string,
|
||||||
|
) (response VectorStoreFileBatch, err error) {
|
||||||
|
urlSuffix := fmt.Sprintf("%s/%s%s/%s%s", vectorStoresSuffix,
|
||||||
|
vectorStoreID, vectorStoresFileBatchesSuffix, batchID, "/cancel")
|
||||||
|
req, _ := c.newRequest(ctx, http.MethodPost, c.fullURL(urlSuffix),
|
||||||
|
withBetaAssistantVersion(c.config.AssistantVersion))
|
||||||
|
|
||||||
|
err = c.sendRequest(req, &response)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListVectorStoreFiles Lists the currently available files for a vector store.
|
||||||
|
func (c *Client) ListVectorStoreFilesInBatch(
|
||||||
|
ctx context.Context,
|
||||||
|
vectorStoreID string,
|
||||||
|
batchID string,
|
||||||
|
pagination Pagination,
|
||||||
|
) (response VectorStoreFilesList, err error) {
|
||||||
|
urlValues := url.Values{}
|
||||||
|
if pagination.After != nil {
|
||||||
|
urlValues.Add("after", *pagination.After)
|
||||||
|
}
|
||||||
|
if pagination.Limit != nil {
|
||||||
|
urlValues.Add("limit", fmt.Sprintf("%d", *pagination.Limit))
|
||||||
|
}
|
||||||
|
if pagination.Before != nil {
|
||||||
|
urlValues.Add("before", *pagination.Before)
|
||||||
|
}
|
||||||
|
if pagination.Order != nil {
|
||||||
|
urlValues.Add("order", *pagination.Order)
|
||||||
|
}
|
||||||
|
|
||||||
|
encodedValues := ""
|
||||||
|
if len(urlValues) > 0 {
|
||||||
|
encodedValues = "?" + urlValues.Encode()
|
||||||
|
}
|
||||||
|
|
||||||
|
urlSuffix := fmt.Sprintf("%s/%s%s/%s%s%s", vectorStoresSuffix,
|
||||||
|
vectorStoreID, vectorStoresFileBatchesSuffix, batchID, "/files", encodedValues)
|
||||||
|
req, _ := c.newRequest(ctx, http.MethodGet, c.fullURL(urlSuffix),
|
||||||
|
withBetaAssistantVersion(c.config.AssistantVersion))
|
||||||
|
|
||||||
|
err = c.sendRequest(req, &response)
|
||||||
|
return
|
||||||
|
}
|
||||||
349
vector_store_test.go
Normal file
349
vector_store_test.go
Normal file
@@ -0,0 +1,349 @@
|
|||||||
|
package openai_test
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
|
||||||
|
openai "github.com/sashabaranov/go-openai"
|
||||||
|
"github.com/sashabaranov/go-openai/internal/test/checks"
|
||||||
|
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
// TestVectorStore Tests the vector store endpoint of the API using the mocked server.
|
||||||
|
func TestVectorStore(t *testing.T) {
|
||||||
|
vectorStoreID := "vs_abc123"
|
||||||
|
vectorStoreName := "TestStore"
|
||||||
|
vectorStoreFileID := "file-wB6RM6wHdA49HfS2DJ9fEyrH"
|
||||||
|
vectorStoreFileBatchID := "vsfb_abc123"
|
||||||
|
limit := 20
|
||||||
|
order := "desc"
|
||||||
|
after := "vs_abc122"
|
||||||
|
before := "vs_abc123"
|
||||||
|
|
||||||
|
client, server, teardown := setupOpenAITestServer()
|
||||||
|
defer teardown()
|
||||||
|
|
||||||
|
server.RegisterHandler(
|
||||||
|
"/v1/vector_stores/"+vectorStoreID+"/files/"+vectorStoreFileID,
|
||||||
|
func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method == http.MethodGet {
|
||||||
|
resBytes, _ := json.Marshal(openai.VectorStoreFile{
|
||||||
|
ID: vectorStoreFileID,
|
||||||
|
Object: "vector_store.file",
|
||||||
|
CreatedAt: 1234567890,
|
||||||
|
VectorStoreID: vectorStoreID,
|
||||||
|
Status: "completed",
|
||||||
|
})
|
||||||
|
fmt.Fprintln(w, string(resBytes))
|
||||||
|
} else if r.Method == http.MethodDelete {
|
||||||
|
fmt.Fprintln(w, `{
|
||||||
|
id: "file-wB6RM6wHdA49HfS2DJ9fEyrH",
|
||||||
|
object: "vector_store.file.deleted",
|
||||||
|
deleted: true
|
||||||
|
}`)
|
||||||
|
}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
server.RegisterHandler(
|
||||||
|
"/v1/vector_stores/"+vectorStoreID+"/files",
|
||||||
|
func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method == http.MethodGet {
|
||||||
|
resBytes, _ := json.Marshal(openai.VectorStoreFilesList{
|
||||||
|
VectorStoreFiles: []openai.VectorStoreFile{
|
||||||
|
{
|
||||||
|
ID: vectorStoreFileID,
|
||||||
|
Object: "vector_store.file",
|
||||||
|
CreatedAt: 1234567890,
|
||||||
|
VectorStoreID: vectorStoreID,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
fmt.Fprintln(w, string(resBytes))
|
||||||
|
} else if r.Method == http.MethodPost {
|
||||||
|
var request openai.VectorStoreFileRequest
|
||||||
|
err := json.NewDecoder(r.Body).Decode(&request)
|
||||||
|
checks.NoError(t, err, "Decode error")
|
||||||
|
|
||||||
|
resBytes, _ := json.Marshal(openai.VectorStoreFile{
|
||||||
|
ID: request.FileID,
|
||||||
|
Object: "vector_store.file",
|
||||||
|
CreatedAt: 1234567890,
|
||||||
|
VectorStoreID: vectorStoreID,
|
||||||
|
})
|
||||||
|
fmt.Fprintln(w, string(resBytes))
|
||||||
|
}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
server.RegisterHandler(
|
||||||
|
"/v1/vector_stores/"+vectorStoreID+"/file_batches/"+vectorStoreFileBatchID+"/files",
|
||||||
|
func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method == http.MethodGet {
|
||||||
|
resBytes, _ := json.Marshal(openai.VectorStoreFilesList{
|
||||||
|
VectorStoreFiles: []openai.VectorStoreFile{
|
||||||
|
{
|
||||||
|
ID: vectorStoreFileID,
|
||||||
|
Object: "vector_store.file",
|
||||||
|
CreatedAt: 1234567890,
|
||||||
|
VectorStoreID: vectorStoreID,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
fmt.Fprintln(w, string(resBytes))
|
||||||
|
}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
server.RegisterHandler(
|
||||||
|
"/v1/vector_stores/"+vectorStoreID+"/file_batches/"+vectorStoreFileBatchID+"/cancel",
|
||||||
|
func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method == http.MethodPost {
|
||||||
|
resBytes, _ := json.Marshal(openai.VectorStoreFileBatch{
|
||||||
|
ID: vectorStoreFileBatchID,
|
||||||
|
Object: "vector_store.file_batch",
|
||||||
|
CreatedAt: 1234567890,
|
||||||
|
VectorStoreID: vectorStoreID,
|
||||||
|
Status: "cancelling",
|
||||||
|
FileCounts: openai.VectorStoreFileCount{
|
||||||
|
InProgress: 0,
|
||||||
|
Completed: 1,
|
||||||
|
Failed: 0,
|
||||||
|
Cancelled: 0,
|
||||||
|
Total: 0,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
fmt.Fprintln(w, string(resBytes))
|
||||||
|
}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
server.RegisterHandler(
|
||||||
|
"/v1/vector_stores/"+vectorStoreID+"/file_batches/"+vectorStoreFileBatchID,
|
||||||
|
func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method == http.MethodGet {
|
||||||
|
resBytes, _ := json.Marshal(openai.VectorStoreFileBatch{
|
||||||
|
ID: vectorStoreFileBatchID,
|
||||||
|
Object: "vector_store.file_batch",
|
||||||
|
CreatedAt: 1234567890,
|
||||||
|
VectorStoreID: vectorStoreID,
|
||||||
|
Status: "completed",
|
||||||
|
FileCounts: openai.VectorStoreFileCount{
|
||||||
|
Completed: 1,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
fmt.Fprintln(w, string(resBytes))
|
||||||
|
} else if r.Method == http.MethodPost {
|
||||||
|
resBytes, _ := json.Marshal(openai.VectorStoreFileBatch{
|
||||||
|
ID: vectorStoreFileBatchID,
|
||||||
|
Object: "vector_store.file_batch",
|
||||||
|
CreatedAt: 1234567890,
|
||||||
|
VectorStoreID: vectorStoreID,
|
||||||
|
Status: "cancelling",
|
||||||
|
FileCounts: openai.VectorStoreFileCount{
|
||||||
|
Completed: 1,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
fmt.Fprintln(w, string(resBytes))
|
||||||
|
}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
server.RegisterHandler(
|
||||||
|
"/v1/vector_stores/"+vectorStoreID+"/file_batches",
|
||||||
|
func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method == http.MethodPost {
|
||||||
|
var request openai.VectorStoreFileBatchRequest
|
||||||
|
err := json.NewDecoder(r.Body).Decode(&request)
|
||||||
|
checks.NoError(t, err, "Decode error")
|
||||||
|
|
||||||
|
resBytes, _ := json.Marshal(openai.VectorStoreFileBatch{
|
||||||
|
ID: vectorStoreFileBatchID,
|
||||||
|
Object: "vector_store.file_batch",
|
||||||
|
CreatedAt: 1234567890,
|
||||||
|
VectorStoreID: vectorStoreID,
|
||||||
|
Status: "completed",
|
||||||
|
FileCounts: openai.VectorStoreFileCount{
|
||||||
|
InProgress: 0,
|
||||||
|
Completed: len(request.FileIDs),
|
||||||
|
Failed: 0,
|
||||||
|
Cancelled: 0,
|
||||||
|
Total: 0,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
fmt.Fprintln(w, string(resBytes))
|
||||||
|
}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
server.RegisterHandler(
|
||||||
|
"/v1/vector_stores/"+vectorStoreID,
|
||||||
|
func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
switch r.Method {
|
||||||
|
case http.MethodGet:
|
||||||
|
resBytes, _ := json.Marshal(openai.VectorStore{
|
||||||
|
ID: vectorStoreID,
|
||||||
|
Object: "vector_store",
|
||||||
|
CreatedAt: 1234567890,
|
||||||
|
Name: vectorStoreName,
|
||||||
|
})
|
||||||
|
fmt.Fprintln(w, string(resBytes))
|
||||||
|
case http.MethodPost:
|
||||||
|
var request openai.VectorStore
|
||||||
|
err := json.NewDecoder(r.Body).Decode(&request)
|
||||||
|
checks.NoError(t, err, "Decode error")
|
||||||
|
|
||||||
|
resBytes, _ := json.Marshal(openai.VectorStore{
|
||||||
|
ID: vectorStoreID,
|
||||||
|
Object: "vector_store",
|
||||||
|
CreatedAt: 1234567890,
|
||||||
|
Name: request.Name,
|
||||||
|
})
|
||||||
|
fmt.Fprintln(w, string(resBytes))
|
||||||
|
case http.MethodDelete:
|
||||||
|
fmt.Fprintln(w, `{
|
||||||
|
"id": "vectorstore_abc123",
|
||||||
|
"object": "vector_store.deleted",
|
||||||
|
"deleted": true
|
||||||
|
}`)
|
||||||
|
}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
server.RegisterHandler(
|
||||||
|
"/v1/vector_stores",
|
||||||
|
func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method == http.MethodPost {
|
||||||
|
var request openai.VectorStoreRequest
|
||||||
|
err := json.NewDecoder(r.Body).Decode(&request)
|
||||||
|
checks.NoError(t, err, "Decode error")
|
||||||
|
|
||||||
|
resBytes, _ := json.Marshal(openai.VectorStore{
|
||||||
|
ID: vectorStoreID,
|
||||||
|
Object: "vector_store",
|
||||||
|
CreatedAt: 1234567890,
|
||||||
|
Name: request.Name,
|
||||||
|
FileCounts: openai.VectorStoreFileCount{
|
||||||
|
InProgress: 0,
|
||||||
|
Completed: 0,
|
||||||
|
Failed: 0,
|
||||||
|
Cancelled: 0,
|
||||||
|
Total: 0,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
fmt.Fprintln(w, string(resBytes))
|
||||||
|
} else if r.Method == http.MethodGet {
|
||||||
|
resBytes, _ := json.Marshal(openai.VectorStoresList{
|
||||||
|
LastID: &vectorStoreID,
|
||||||
|
FirstID: &vectorStoreID,
|
||||||
|
VectorStores: []openai.VectorStore{
|
||||||
|
{
|
||||||
|
ID: vectorStoreID,
|
||||||
|
Object: "vector_store",
|
||||||
|
CreatedAt: 1234567890,
|
||||||
|
Name: vectorStoreName,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
fmt.Fprintln(w, string(resBytes))
|
||||||
|
}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
t.Run("create_vector_store", func(t *testing.T) {
|
||||||
|
_, err := client.CreateVectorStore(ctx, openai.VectorStoreRequest{
|
||||||
|
Name: vectorStoreName,
|
||||||
|
})
|
||||||
|
checks.NoError(t, err, "CreateVectorStore error")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("retrieve_vector_store", func(t *testing.T) {
|
||||||
|
_, err := client.RetrieveVectorStore(ctx, vectorStoreID)
|
||||||
|
checks.NoError(t, err, "RetrieveVectorStore error")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("delete_vector_store", func(t *testing.T) {
|
||||||
|
_, err := client.DeleteVectorStore(ctx, vectorStoreID)
|
||||||
|
checks.NoError(t, err, "DeleteVectorStore error")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("list_vector_store", func(t *testing.T) {
|
||||||
|
_, err := client.ListVectorStores(context.TODO(), openai.Pagination{
|
||||||
|
Limit: &limit,
|
||||||
|
Order: &order,
|
||||||
|
After: &after,
|
||||||
|
Before: &before,
|
||||||
|
})
|
||||||
|
checks.NoError(t, err, "ListVectorStores error")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("create_vector_store_file", func(t *testing.T) {
|
||||||
|
_, err := client.CreateVectorStoreFile(context.TODO(), vectorStoreID, openai.VectorStoreFileRequest{
|
||||||
|
FileID: vectorStoreFileID,
|
||||||
|
})
|
||||||
|
checks.NoError(t, err, "CreateVectorStoreFile error")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("list_vector_store_files", func(t *testing.T) {
|
||||||
|
_, err := client.ListVectorStoreFiles(ctx, vectorStoreID, openai.Pagination{
|
||||||
|
Limit: &limit,
|
||||||
|
Order: &order,
|
||||||
|
After: &after,
|
||||||
|
Before: &before,
|
||||||
|
})
|
||||||
|
checks.NoError(t, err, "ListVectorStoreFiles error")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("retrieve_vector_store_file", func(t *testing.T) {
|
||||||
|
_, err := client.RetrieveVectorStoreFile(ctx, vectorStoreID, vectorStoreFileID)
|
||||||
|
checks.NoError(t, err, "RetrieveVectorStoreFile error")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("delete_vector_store_file", func(t *testing.T) {
|
||||||
|
err := client.DeleteVectorStoreFile(ctx, vectorStoreID, vectorStoreFileID)
|
||||||
|
checks.NoError(t, err, "DeleteVectorStoreFile error")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("modify_vector_store", func(t *testing.T) {
|
||||||
|
_, err := client.ModifyVectorStore(ctx, vectorStoreID, openai.VectorStoreRequest{
|
||||||
|
Name: vectorStoreName,
|
||||||
|
})
|
||||||
|
checks.NoError(t, err, "ModifyVectorStore error")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("create_vector_store_file_batch", func(t *testing.T) {
|
||||||
|
_, err := client.CreateVectorStoreFileBatch(ctx, vectorStoreID, openai.VectorStoreFileBatchRequest{
|
||||||
|
FileIDs: []string{vectorStoreFileID},
|
||||||
|
})
|
||||||
|
checks.NoError(t, err, "CreateVectorStoreFileBatch error")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("retrieve_vector_store_file_batch", func(t *testing.T) {
|
||||||
|
_, err := client.RetrieveVectorStoreFileBatch(ctx, vectorStoreID, vectorStoreFileBatchID)
|
||||||
|
checks.NoError(t, err, "RetrieveVectorStoreFileBatch error")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("list_vector_store_files_in_batch", func(t *testing.T) {
|
||||||
|
_, err := client.ListVectorStoreFilesInBatch(
|
||||||
|
ctx,
|
||||||
|
vectorStoreID,
|
||||||
|
vectorStoreFileBatchID,
|
||||||
|
openai.Pagination{
|
||||||
|
Limit: &limit,
|
||||||
|
Order: &order,
|
||||||
|
After: &after,
|
||||||
|
Before: &before,
|
||||||
|
})
|
||||||
|
checks.NoError(t, err, "ListVectorStoreFilesInBatch error")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("cancel_vector_store_file_batch", func(t *testing.T) {
|
||||||
|
_, err := client.CancelVectorStoreFileBatch(ctx, vectorStoreID, vectorStoreFileBatchID)
|
||||||
|
checks.NoError(t, err, "CancelVectorStoreFileBatch error")
|
||||||
|
})
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user