From 948404a0412e662e786df332791231838bf80617 Mon Sep 17 00:00:00 2001 From: Grant Linville Date: Tue, 17 Dec 2024 15:15:25 -0500 Subject: [PATCH] enhance: get more information about models Signed-off-by: Grant Linville --- gptscript.go | 8 ++++++-- gptscript_test.go | 8 ++++---- openai.go | 27 +++++++++++++++++++++++++++ 3 files changed, 37 insertions(+), 6 deletions(-) create mode 100644 openai.go diff --git a/gptscript.go b/gptscript.go index 1e30d95..0fa173d 100644 --- a/gptscript.go +++ b/gptscript.go @@ -294,7 +294,7 @@ type ListModelsOptions struct { } // ListModels will list all the available models. -func (g *GPTScript) ListModels(ctx context.Context, opts ...ListModelsOptions) ([]string, error) { +func (g *GPTScript) ListModels(ctx context.Context, opts ...ListModelsOptions) ([]Model, error) { var o ListModelsOptions for _, opt := range opts { o.Providers = append(o.Providers, opt.Providers...) @@ -314,7 +314,11 @@ func (g *GPTScript) ListModels(ctx context.Context, opts ...ListModelsOptions) ( return nil, err } - return strings.Split(strings.TrimSpace(out), "\n"), nil + var models []Model + if err = json.Unmarshal([]byte(out), &models); err != nil { + return nil, err + } + return models, nil } func (g *GPTScript) Confirm(ctx context.Context, resp AuthResponse) error { diff --git a/gptscript_test.go b/gptscript_test.go index 32adff9..834181e 100644 --- a/gptscript_test.go +++ b/gptscript_test.go @@ -98,8 +98,8 @@ func TestListModelsWithProvider(t *testing.T) { } for _, model := range models { - if !strings.HasPrefix(model, "claude-3-") || !strings.HasSuffix(model, "from github.com/gptscript-ai/claude3-anthropic-provider") { - t.Errorf("Unexpected model name: %s", model) + if !strings.HasPrefix(model.ID, "claude-3-") || !strings.HasSuffix(model.ID, "from github.com/gptscript-ai/claude3-anthropic-provider") { + t.Errorf("Unexpected model name: %s", model.ID) } } } @@ -128,8 +128,8 @@ func TestListModelsWithDefaultProvider(t *testing.T) { } for _, model := range models { - if !strings.HasPrefix(model, "claude-3-") || !strings.HasSuffix(model, "from github.com/gptscript-ai/claude3-anthropic-provider") { - t.Errorf("Unexpected model name: %s", model) + if !strings.HasPrefix(model.ID, "claude-3-") || !strings.HasSuffix(model.ID, "from github.com/gptscript-ai/claude3-anthropic-provider") { + t.Errorf("Unexpected model name: %s", model.ID) } } } diff --git a/openai.go b/openai.go new file mode 100644 index 0000000..eb970a6 --- /dev/null +++ b/openai.go @@ -0,0 +1,27 @@ +package gptscript + +type Permission struct { + CreatedAt int64 `json:"created"` + ID string `json:"id"` + Object string `json:"object"` + AllowCreateEngine bool `json:"allow_create_engine"` + AllowSampling bool `json:"allow_sampling"` + AllowLogprobs bool `json:"allow_logprobs"` + AllowSearchIndices bool `json:"allow_search_indices"` + AllowView bool `json:"allow_view"` + AllowFineTuning bool `json:"allow_fine_tuning"` + Organization string `json:"organization"` + Group interface{} `json:"group"` + IsBlocking bool `json:"is_blocking"` +} + +type Model struct { + CreatedAt int64 `json:"created"` + ID string `json:"id"` + Object string `json:"object"` + OwnedBy string `json:"owned_by"` + Permission []Permission `json:"permission"` + Root string `json:"root"` + Parent string `json:"parent"` + Metadata map[string]string `json:"metadata"` +}