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
16
assistant.go
16
assistant.go
@@ -24,6 +24,7 @@ type Assistant struct {
|
||||
Tools []AssistantTool `json:"tools"`
|
||||
FileIDs []string `json:"file_ids,omitempty"`
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
ToolResources *AssistantToolResource `json:"tool_resources,omitempty"`
|
||||
|
||||
httpHeader
|
||||
}
|
||||
@@ -34,6 +35,7 @@ const (
|
||||
AssistantToolTypeCodeInterpreter AssistantToolType = "code_interpreter"
|
||||
AssistantToolTypeRetrieval AssistantToolType = "retrieval"
|
||||
AssistantToolTypeFunction AssistantToolType = "function"
|
||||
AssistantToolTypeFileSearch AssistantToolType = "file_search"
|
||||
)
|
||||
|
||||
type AssistantTool struct {
|
||||
@@ -41,6 +43,19 @@ type AssistantTool struct {
|
||||
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.
|
||||
// When modifying the tools the API functions as the following:
|
||||
// If Tools is undefined, no changes are made to the Assistant's tools.
|
||||
@@ -54,6 +69,7 @@ type AssistantRequest struct {
|
||||
Tools []AssistantTool `json:"-"`
|
||||
FileIDs []string `json:"file_ids,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
|
||||
|
||||
@@ -24,7 +24,7 @@ const (
|
||||
|
||||
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.
|
||||
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