go-openai/vector_store_test.go

package openai_test

import (
	"context"

	openai "git.vaala.cloud/VaalaCat/go-openai"
	"git.vaala.cloud/VaalaCat/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")
	})
}