Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion typesense/api/client_gen.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

8 changes: 8 additions & 0 deletions typesense/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,14 @@ func (c *Client) Preset(presetName string) PresetInterface {
return &preset{apiClient: c.apiClient, presetName: presetName}
}

func (c *Client) NLSearchModels() NLSearchModelsInterface {
return &nlSearchModels{apiClient: c.apiClient}
}

func (c *Client) NLSearchModel(modelID string) NLSearchModelInterface {
return &nlSearchModel{apiClient: c.apiClient, modelID: modelID}
}

func (c *Client) Stopwords() StopwordsInterface {
return &stopwords{apiClient: c.apiClient}
}
Expand Down
51 changes: 51 additions & 0 deletions typesense/nl_search_model.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
package typesense

import (
"context"

"github.com/typesense/typesense-go/v3/typesense/api"
)

type NLSearchModelInterface interface {
Retrieve(ctx context.Context) (*api.NLSearchModelSchema, error)
Update(ctx context.Context, model *api.NLSearchModelUpdateSchema) (*api.NLSearchModelSchema, error)
Delete(ctx context.Context) (*api.NLSearchModelDeleteSchema, error)
}

type nlSearchModel struct {
apiClient APIClientInterface
modelID string
}

func (n *nlSearchModel) Retrieve(ctx context.Context) (*api.NLSearchModelSchema, error) {
response, err := n.apiClient.RetrieveNLSearchModelWithResponse(ctx, n.modelID)
if err != nil {
return nil, err
}
if response.JSON200 == nil {
return nil, &HTTPError{Status: response.StatusCode(), Body: response.Body}
}
return response.JSON200, nil
}

func (n *nlSearchModel) Update(ctx context.Context, model *api.NLSearchModelUpdateSchema) (*api.NLSearchModelSchema, error) {
response, err := n.apiClient.UpdateNLSearchModelWithResponse(ctx, n.modelID, *model)
if err != nil {
return nil, err
}
if response.JSON200 == nil {
return nil, &HTTPError{Status: response.StatusCode(), Body: response.Body}
}
return response.JSON200, nil
}

func (n *nlSearchModel) Delete(ctx context.Context) (*api.NLSearchModelDeleteSchema, error) {
response, err := n.apiClient.DeleteNLSearchModelWithResponse(ctx, n.modelID)
if err != nil {
return nil, err
}
if response.JSON200 == nil {
return nil, &HTTPError{Status: response.StatusCode(), Body: response.Body}
}
return response.JSON200, nil
}
45 changes: 45 additions & 0 deletions typesense/nl_search_models.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
package typesense

import (
"context"

"github.com/typesense/typesense-go/v3/typesense/api"
)

type NLSearchModelsInterface interface {
Retrieve(ctx context.Context) ([]*api.NLSearchModelSchema, error)
Create(ctx context.Context, model *api.NLSearchModelCreateSchema) (*api.NLSearchModelSchema, error)
}

type nlSearchModels struct {
apiClient APIClientInterface
}

func (n *nlSearchModels) Retrieve(ctx context.Context) ([]*api.NLSearchModelSchema, error) {
response, err := n.apiClient.RetrieveAllNLSearchModelsWithResponse(ctx)
if err != nil {
return nil, err
}
if response.JSON200 == nil {
return nil, &HTTPError{Status: response.StatusCode(), Body: response.Body}
}

// Convert []NLSearchModelSchema to []*NLSearchModelSchema
result := make([]*api.NLSearchModelSchema, len(*response.JSON200))
for i, model := range *response.JSON200 {
modelCopy := model // Create a copy to get address
result[i] = &modelCopy
}
return result, nil
}

func (n *nlSearchModels) Create(ctx context.Context, model *api.NLSearchModelCreateSchema) (*api.NLSearchModelSchema, error) {
response, err := n.apiClient.CreateNLSearchModelWithResponse(ctx, *model)
if err != nil {
return nil, err
}
if response.JSON201 == nil {
return nil, &HTTPError{Status: response.StatusCode(), Body: response.Body}
}
return response.JSON201, nil
}
68 changes: 68 additions & 0 deletions typesense/test/dbhelpers_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ package test
import (
"context"
"fmt"
"os"
"testing"
"time"

Expand Down Expand Up @@ -441,3 +442,70 @@ func retrieveDocuments(t *testing.T, collectionName string, docIDs ...string) []
}
return results
}

func newNLSearchModelCreateSchema() *api.NLSearchModelCreateSchema {
apiKey := os.Getenv("NL_SEARCH_MODEL_API_KEY")

return &api.NLSearchModelCreateSchema{
ModelName: pointer.String("openai/gpt-3.5-turbo"),
ApiKey: pointer.String(apiKey),
MaxBytes: pointer.Int(1000),
Temperature: pointer.Float32(0.7),
SystemPrompt: pointer.String("You are a helpful assistant."),
TopP: pointer.Float32(0.9),
TopK: pointer.Int(40),
StopSequences: &[]string{"END", "STOP"},
ApiVersion: pointer.String("v1"),
}
}

func newNLSearchModelSchema(modelID string) *api.NLSearchModelSchema {
apiKey := os.Getenv("NL_SEARCH_MODEL_API_KEY")

return &api.NLSearchModelSchema{
Id: modelID,
ModelName: pointer.String("openai/gpt-3.5-turbo"),
ApiKey: pointer.String(apiKey),
MaxBytes: pointer.Int(1000),
Temperature: pointer.Float32(0.7),
SystemPrompt: pointer.String("You are a helpful assistant."),
TopP: pointer.Float32(0.9),
TopK: pointer.Int(40),
StopSequences: &[]string{"END", "STOP"},
ApiVersion: pointer.String("v1"),
}
}

func newNLSearchModelUpdateSchema() *api.NLSearchModelUpdateSchema {
apiKey := os.Getenv("NL_SEARCH_MODEL_API_KEY")

return &api.NLSearchModelUpdateSchema{
ModelName: pointer.String("openai/gpt-4"),
ApiKey: pointer.String(apiKey),
MaxBytes: pointer.Int(2000),
Temperature: pointer.Float32(0.5),
SystemPrompt: pointer.String("You are an expert assistant."),
TopP: pointer.Float32(0.8),
TopK: pointer.Int(50),
StopSequences: &[]string{"END", "STOP", "QUIT"},
ApiVersion: pointer.String("v1"),
}
}

func shouldSkipNLSearchModelTests(t *testing.T) {
if os.Getenv("NL_SEARCH_MODEL_API_KEY") == "" {
t.Skip("Skipping NL search model test: NL_SEARCH_MODEL_API_KEY not set")
}
}

func createNewNLSearchModel(t *testing.T) (string, *api.NLSearchModelSchema) {
t.Helper()
modelID := newUUIDName("nl-model-test")
modelSchema := newNLSearchModelCreateSchema()
modelSchema.Id = pointer.String(modelID)

result, err := typesenseClient.NLSearchModels().Create(context.Background(), modelSchema)

require.NoError(t, err)
return modelID, result
}
59 changes: 59 additions & 0 deletions typesense/test/nl_search_model_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
//go:build integration
// +build integration

package test

import (
"context"
"testing"

"github.com/stretchr/testify/require"
"github.com/typesense/typesense-go/v3/typesense/api/pointer"
)

func nlSearchModelsCleanUp() {
result, _ := typesenseClient.NLSearchModels().Retrieve(context.Background())
for _, model := range result {
typesenseClient.NLSearchModel(model.Id).Delete(context.Background())
}
}

func TestNLSearchModel(t *testing.T) {
shouldSkipNLSearchModelTests(t)
t.Cleanup(nlSearchModelsCleanUp)

t.Run("Retrieve", func(t *testing.T) {
modelID, expectedResult := createNewNLSearchModel(t)

result, err := typesenseClient.NLSearchModel(modelID).Retrieve(context.Background())

require.NoError(t, err)
require.Equal(t, expectedResult, result)
})

t.Run("Update", func(t *testing.T) {
modelID, originalModel := createNewNLSearchModel(t)

updateSchema := newNLSearchModelUpdateSchema()
updateSchema.Temperature = pointer.Float32(0.8)

result, err := typesenseClient.NLSearchModel(modelID).Update(context.Background(), updateSchema)

require.NoError(t, err)
require.Equal(t, "openai/gpt-4", *result.ModelName)
require.Equal(t, float32(0.8), *result.Temperature)
require.Equal(t, originalModel.Id, result.Id)
})

t.Run("Delete", func(t *testing.T) {
modelID, expectedResult := createNewNLSearchModel(t)

result, err := typesenseClient.NLSearchModel(modelID).Delete(context.Background())

require.NoError(t, err)
require.Equal(t, expectedResult.Id, result.Id)

_, err = typesenseClient.NLSearchModel(modelID).Retrieve(context.Background())
require.Error(t, err)
})
}
2 changes: 0 additions & 2 deletions typesense/test/search_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -156,8 +156,6 @@ func TestCollectionSearchWithPreset(t *testing.T) {
newDocument("123", withCompanyName("Company 1"), withNumEmployees(50)),
newDocument("125", withCompanyName("Company 2"), withNumEmployees(150)),
newDocument("127", withCompanyName("Company 3"), withNumEmployees(250)),
newDocument("129", withCompanyName("Stark Industries 4"), withNumEmployees(500)),
newDocument("131", withCompanyName("Stark Industries 5"), withNumEmployees(1000)),
}

params := &api.ImportDocumentsParams{Action: pointer.Any(api.Create)}
Expand Down
Loading