Skip to content

Commit

Permalink
chore: use model type for listing models
Browse files Browse the repository at this point in the history
Signed-off-by: Donnie Adams <[email protected]>
  • Loading branch information
thedadams committed Jan 14, 2025
1 parent 4c3ef0b commit 07f7c7c
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 6 deletions.
35 changes: 33 additions & 2 deletions gptscript.go
Original file line number Diff line number Diff line change
Expand Up @@ -293,8 +293,34 @@ type ListModelsOptions struct {
CredentialOverrides []string
}

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"`
}

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"`
}

// 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...)
Expand All @@ -314,7 +340,12 @@ 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, fmt.Errorf("failed to parse models: %w", err)
}

return models, nil
}

func (g *GPTScript) Confirm(ctx context.Context, resp AuthResponse) error {
Expand Down
8 changes: 4 additions & 4 deletions gptscript_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
}
Expand Down Expand Up @@ -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)
}
}
}
Expand Down

0 comments on commit 07f7c7c

Please sign in to comment.