Skip to content

Commit 07f7c7c

Browse files
committed
chore: use model type for listing models
Signed-off-by: Donnie Adams <[email protected]>
1 parent 4c3ef0b commit 07f7c7c

File tree

2 files changed

+37
-6
lines changed

2 files changed

+37
-6
lines changed

gptscript.go

+33-2
Original file line numberDiff line numberDiff line change
@@ -293,8 +293,34 @@ type ListModelsOptions struct {
293293
CredentialOverrides []string
294294
}
295295

296+
type Model struct {
297+
CreatedAt int64 `json:"created"`
298+
ID string `json:"id"`
299+
Object string `json:"object"`
300+
OwnedBy string `json:"owned_by"`
301+
Permission []Permission `json:"permission"`
302+
Root string `json:"root"`
303+
Parent string `json:"parent"`
304+
Metadata map[string]string `json:"metadata"`
305+
}
306+
307+
type Permission struct {
308+
CreatedAt int64 `json:"created"`
309+
ID string `json:"id"`
310+
Object string `json:"object"`
311+
AllowCreateEngine bool `json:"allow_create_engine"`
312+
AllowSampling bool `json:"allow_sampling"`
313+
AllowLogprobs bool `json:"allow_logprobs"`
314+
AllowSearchIndices bool `json:"allow_search_indices"`
315+
AllowView bool `json:"allow_view"`
316+
AllowFineTuning bool `json:"allow_fine_tuning"`
317+
Organization string `json:"organization"`
318+
Group interface{} `json:"group"`
319+
IsBlocking bool `json:"is_blocking"`
320+
}
321+
296322
// ListModels will list all the available models.
297-
func (g *GPTScript) ListModels(ctx context.Context, opts ...ListModelsOptions) ([]string, error) {
323+
func (g *GPTScript) ListModels(ctx context.Context, opts ...ListModelsOptions) ([]Model, error) {
298324
var o ListModelsOptions
299325
for _, opt := range opts {
300326
o.Providers = append(o.Providers, opt.Providers...)
@@ -314,7 +340,12 @@ func (g *GPTScript) ListModels(ctx context.Context, opts ...ListModelsOptions) (
314340
return nil, err
315341
}
316342

317-
return strings.Split(strings.TrimSpace(out), "\n"), nil
343+
var models []Model
344+
if err = json.Unmarshal([]byte(out), &models); err != nil {
345+
return nil, fmt.Errorf("failed to parse models: %w", err)
346+
}
347+
348+
return models, nil
318349
}
319350

320351
func (g *GPTScript) Confirm(ctx context.Context, resp AuthResponse) error {

gptscript_test.go

+4-4
Original file line numberDiff line numberDiff line change
@@ -98,8 +98,8 @@ func TestListModelsWithProvider(t *testing.T) {
9898
}
9999

100100
for _, model := range models {
101-
if !strings.HasPrefix(model, "claude-3-") || !strings.HasSuffix(model, "from github.com/gptscript-ai/claude3-anthropic-provider") {
102-
t.Errorf("Unexpected model name: %s", model)
101+
if !strings.HasPrefix(model.ID, "claude-3-") || !strings.HasSuffix(model.ID, "from github.com/gptscript-ai/claude3-anthropic-provider") {
102+
t.Errorf("Unexpected model name: %s", model.ID)
103103
}
104104
}
105105
}
@@ -128,8 +128,8 @@ func TestListModelsWithDefaultProvider(t *testing.T) {
128128
}
129129

130130
for _, model := range models {
131-
if !strings.HasPrefix(model, "claude-3-") || !strings.HasSuffix(model, "from github.com/gptscript-ai/claude3-anthropic-provider") {
132-
t.Errorf("Unexpected model name: %s", model)
131+
if !strings.HasPrefix(model.ID, "claude-3-") || !strings.HasSuffix(model.ID, "from github.com/gptscript-ai/claude3-anthropic-provider") {
132+
t.Errorf("Unexpected model name: %s", model.ID)
133133
}
134134
}
135135
}

0 commit comments

Comments
 (0)