Skip to content

Commit 8155caa

Browse files
ericcurtinericcurtin
committed
Refactor list and run commands to support OpenAI
Change --openai flag from boolean to URL string in both commands. Implement OpenAI endpoint connection for model listing and chat. Update internal functions to accept interface parameters instead of concrete types for better abstraction. Add validation to prevent incompatible flag combinations. Move OpenAI-specific logic to separate functions for clarity Signed-off-by: Eric Curtin <[email protected]>
1 parent 50c9b8a commit 8155caa

File tree

8 files changed

+448
-98
lines changed

8 files changed

+448
-98
lines changed

‎cmd/cli/commands/df.go‎

Lines changed: 22 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"github.com/docker/model-runner/cmd/cli/commands/completion"
88
"github.com/docker/model-runner/cmd/cli/desktop"
99
"github.com/olekukonko/tablewriter"
10+
"github.com/olekukonko/tablewriter/tw"
1011
"github.com/spf13/cobra"
1112
)
1213

@@ -29,21 +30,27 @@ func newDFCmd() *cobra.Command {
2930

3031
func diskUsageTable(df desktop.DiskUsage) string {
3132
var buf bytes.Buffer
32-
table := tablewriter.NewWriter(&buf)
33-
34-
table.SetHeader([]string{"TYPE", "SIZE"})
35-
36-
table.SetBorder(false)
37-
table.SetColumnSeparator("")
38-
table.SetHeaderLine(false)
39-
table.SetTablePadding(" ")
40-
table.SetNoWhiteSpace(true)
41-
42-
table.SetColumnAlignment([]int{
43-
tablewriter.ALIGN_LEFT, // TYPE
44-
tablewriter.ALIGN_LEFT, // SIZE
45-
})
46-
table.SetHeaderAlignment(tablewriter.ALIGN_LEFT)
33+
table := tablewriter.NewTable(&buf,
34+
tablewriter.WithHeader([]string{"TYPE", "SIZE"}),
35+
tablewriter.WithHeaderAlignment(tw.AlignLeft),
36+
tablewriter.WithRendition(tw.Rendition{
37+
Borders: tw.Border{Left: tw.Off, Right: tw.Off, Top: tw.Off, Bottom: tw.Off},
38+
Settings: tw.Settings{
39+
Separators: tw.Separators{
40+
ShowHeader: tw.Off,
41+
BetweenRows: tw.Off,
42+
BetweenColumns: tw.Off,
43+
},
44+
Lines: tw.Lines{
45+
ShowTop: tw.Off,
46+
ShowBottom: tw.Off,
47+
ShowHeaderLine: tw.Off,
48+
ShowFooterLine: tw.Off,
49+
},
50+
},
51+
}),
52+
tablewriter.WithPadding(tw.Padding{Left: "", Right: " "}),
53+
)
4754

4855
table.Append([]string{"Models", units.CustomSize("%.2f%s", float64(df.ModelsDiskUsage), 1000.0, []string{"B", "kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"})})
4956
if df.DefaultBackendDiskUsage != 0 {

‎cmd/cli/commands/list.go‎

Lines changed: 101 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -16,35 +16,50 @@ import (
1616
"github.com/docker/model-runner/cmd/cli/pkg/standalone"
1717
dmrm "github.com/docker/model-runner/pkg/inference/models"
1818
"github.com/olekukonko/tablewriter"
19+
"github.com/olekukonko/tablewriter/tw"
1920
"github.com/spf13/cobra"
2021
)
2122

2223
func newListCmd() *cobra.Command {
23-
var jsonFormat, openai, quiet bool
24+
var jsonFormat, quiet bool
25+
var openaiURL string
2426
c := &cobra.Command{
2527
Use: "list [OPTIONS] [MODEL]",
2628
Aliases: []string{"ls"},
2729
Short: "List the models pulled to your local environment",
2830
Args: cobra.MaximumNArgs(1),
2931
RunE: func(cmd *cobra.Command, args []string) error {
30-
if openai && quiet {
31-
return fmt.Errorf("--quiet flag cannot be used with --openai flag or OpenAI backend")
32+
useOpenAI := openaiURL != ""
33+
if useOpenAI && quiet {
34+
return fmt.Errorf("--quiet flag cannot be used with --openai flag")
35+
}
36+
37+
var modelFilter string
38+
if len(args) > 0 {
39+
modelFilter = args[0]
40+
}
41+
42+
// If --openai URL is provided, connect to external OpenAI endpoint
43+
if useOpenAI {
44+
openaiClient := desktop.NewOpenAIClient(openaiURL)
45+
models, err := listModelsFromOpenAI(openaiClient, modelFilter)
46+
if err != nil {
47+
return err
48+
}
49+
fmt.Fprint(cmd.OutOrStdout(), models)
50+
return nil
3251
}
3352

3453
// If we're doing an automatic install, only show the installation
3554
// status if it won't corrupt machine-readable output.
3655
var standaloneInstallPrinter standalone.StatusPrinter
37-
if !jsonFormat && !openai && !quiet {
56+
if !jsonFormat && !quiet {
3857
standaloneInstallPrinter = asPrinter(cmd)
3958
}
4059
if _, err := ensureStandaloneRunnerAvailable(cmd.Context(), standaloneInstallPrinter, false); err != nil {
4160
return fmt.Errorf("unable to initialize standalone model runner: %w", err)
4261
}
43-
var modelFilter string
44-
if len(args) > 0 {
45-
modelFilter = args[0]
46-
}
47-
models, err := listModels(openai, desktopClient, quiet, jsonFormat, modelFilter)
62+
models, err := listModels(desktopClient, quiet, jsonFormat, modelFilter)
4863
if err != nil {
4964
return err
5065
}
@@ -54,7 +69,7 @@ func newListCmd() *cobra.Command {
5469
ValidArgsFunction: completion.ModelNamesAndTags(getDesktopClient, 1),
5570
}
5671
c.Flags().BoolVar(&jsonFormat, "json", false, "List models in a JSON format")
57-
c.Flags().BoolVar(&openai, "openai", false, "List models in an OpenAI format")
72+
c.Flags().StringVar(&openaiURL, "openai", "", "List models from an OpenAI-compatible endpoint URL")
5873
c.Flags().BoolVarP(&quiet, "quiet", "q", false, "Only show model IDs")
5974
return c
6075
}
@@ -74,25 +89,64 @@ func matchesModelFilter(tag, filter string) bool {
7489
return repository == filter
7590
}
7691

77-
func listModels(openai bool, desktopClient *desktop.Client, quiet bool, jsonFormat bool, modelFilter string) (string, error) {
78-
if openai {
79-
models, err := desktopClient.ListOpenAI()
80-
if err != nil {
81-
return "", handleClientError(err, "Failed to list models")
82-
}
83-
if modelFilter != "" {
84-
filter := normalizeModelFilter(modelFilter)
85-
filtered := models.Data[:0]
86-
for _, m := range models.Data {
87-
if matchesModelFilter(m.ID, filter) {
88-
filtered = append(filtered, m)
89-
}
92+
func listModelsFromOpenAI(openaiClient *desktop.OpenAIClient, modelFilter string) (string, error) {
93+
models, err := openaiClient.ListModels()
94+
if err != nil {
95+
return "", handleClientError(err, "Failed to list models")
96+
}
97+
if modelFilter != "" {
98+
var filteredData []*dmrm.OpenAIModel
99+
for _, m := range models.Data {
100+
if matchesModelFilter(m.ID, modelFilter) {
101+
filteredData = append(filteredData, m)
90102
}
91-
models.Data = filtered
92103
}
93-
return formatter.ToStandardJSON(models)
104+
models.Data = filteredData
105+
}
106+
return prettyPrintOpenAIModels(models), nil
107+
}
108+
109+
func prettyPrintOpenAIModels(models dmrm.OpenAIModelList) string {
110+
// Sort models by ID
111+
sort.Slice(models.Data, func(i, j int) bool {
112+
return strings.ToLower(models.Data[i].ID) < strings.ToLower(models.Data[j].ID)
113+
})
114+
115+
var buf bytes.Buffer
116+
table := tablewriter.NewTable(&buf,
117+
tablewriter.WithHeader([]string{"MODEL NAME", "OWNED BY"}),
118+
tablewriter.WithHeaderAlignment(tw.AlignLeft),
119+
tablewriter.WithRendition(tw.Rendition{
120+
Borders: tw.Border{Left: tw.Off, Right: tw.Off, Top: tw.Off, Bottom: tw.Off},
121+
Settings: tw.Settings{
122+
Separators: tw.Separators{
123+
ShowHeader: tw.Off,
124+
BetweenRows: tw.Off,
125+
BetweenColumns: tw.Off,
126+
},
127+
Lines: tw.Lines{
128+
ShowTop: tw.Off,
129+
ShowBottom: tw.Off,
130+
ShowHeaderLine: tw.Off,
131+
ShowFooterLine: tw.Off,
132+
},
133+
},
134+
}),
135+
tablewriter.WithPadding(tw.Padding{Left: "", Right: " "}),
136+
)
137+
138+
for _, m := range models.Data {
139+
table.Append([]string{
140+
m.ID,
141+
m.OwnedBy,
142+
})
94143
}
95144

145+
table.Render()
146+
return buf.String()
147+
}
148+
149+
func listModels(desktopClient *desktop.Client, quiet bool, jsonFormat bool, modelFilter string) (string, error) {
96150
models, err := desktopClient.List()
97151
if err != nil {
98152
return "", handleClientError(err, "Failed to list models")
@@ -197,27 +251,28 @@ func prettyPrintModels(models []dmrm.Model) string {
197251
})
198252

199253
var buf bytes.Buffer
200-
table := tablewriter.NewWriter(&buf)
201-
202-
table.SetHeader([]string{"MODEL NAME", "PARAMETERS", "QUANTIZATION", "ARCHITECTURE", "MODEL ID", "CREATED", "CONTEXT", "SIZE"})
203-
204-
table.SetBorder(false)
205-
table.SetColumnSeparator("")
206-
table.SetHeaderLine(false)
207-
table.SetTablePadding(" ")
208-
table.SetNoWhiteSpace(true)
209-
210-
table.SetColumnAlignment([]int{
211-
tablewriter.ALIGN_LEFT, // MODEL
212-
tablewriter.ALIGN_LEFT, // PARAMETERS
213-
tablewriter.ALIGN_LEFT, // QUANTIZATION
214-
tablewriter.ALIGN_LEFT, // ARCHITECTURE
215-
tablewriter.ALIGN_LEFT, // MODEL ID
216-
tablewriter.ALIGN_LEFT, // CREATED
217-
tablewriter.ALIGN_RIGHT, // CONTEXT
218-
tablewriter.ALIGN_LEFT, // SIZE
219-
})
220-
table.SetHeaderAlignment(tablewriter.ALIGN_LEFT)
254+
table := tablewriter.NewTable(&buf,
255+
tablewriter.WithHeader([]string{"MODEL NAME", "PARAMETERS", "QUANTIZATION", "ARCHITECTURE", "MODEL ID", "CREATED", "CONTEXT", "SIZE"}),
256+
tablewriter.WithHeaderAlignment(tw.AlignLeft),
257+
tablewriter.WithRendition(tw.Rendition{
258+
Borders: tw.Border{Left: tw.Off, Right: tw.Off, Top: tw.Off, Bottom: tw.Off},
259+
Settings: tw.Settings{
260+
Separators: tw.Separators{
261+
ShowHeader: tw.Off,
262+
BetweenRows: tw.Off,
263+
BetweenColumns: tw.Off,
264+
},
265+
Lines: tw.Lines{
266+
ShowTop: tw.Off,
267+
ShowBottom: tw.Off,
268+
ShowHeaderLine: tw.Off,
269+
ShowFooterLine: tw.Off,
270+
},
271+
},
272+
}),
273+
tablewriter.WithPadding(tw.Padding{Left: "", Right: " "}),
274+
tablewriter.WithAlignment(tw.MakeAlign(8, tw.AlignLeft).Set(6, tw.AlignRight)),
275+
)
221276

222277
for _, row := range rows {
223278
appendRow(table, row.tag, row.model)

‎cmd/cli/commands/ps.go‎

Lines changed: 22 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"github.com/docker/model-runner/cmd/cli/commands/completion"
1010
"github.com/docker/model-runner/cmd/cli/desktop"
1111
"github.com/olekukonko/tablewriter"
12+
"github.com/olekukonko/tablewriter/tw"
1213
"github.com/spf13/cobra"
1314
)
1415

@@ -31,23 +32,27 @@ func newPSCmd() *cobra.Command {
3132

3233
func psTable(ps []desktop.BackendStatus) string {
3334
var buf bytes.Buffer
34-
table := tablewriter.NewWriter(&buf)
35-
36-
table.SetHeader([]string{"MODEL NAME", "BACKEND", "MODE", "LAST USED"})
37-
38-
table.SetBorder(false)
39-
table.SetColumnSeparator("")
40-
table.SetHeaderLine(false)
41-
table.SetTablePadding(" ")
42-
table.SetNoWhiteSpace(true)
43-
44-
table.SetColumnAlignment([]int{
45-
tablewriter.ALIGN_LEFT, // MODEL
46-
tablewriter.ALIGN_LEFT, // BACKEND
47-
tablewriter.ALIGN_LEFT, // MODE
48-
tablewriter.ALIGN_LEFT, // LAST USED
49-
})
50-
table.SetHeaderAlignment(tablewriter.ALIGN_LEFT)
35+
table := tablewriter.NewTable(&buf,
36+
tablewriter.WithHeader([]string{"MODEL NAME", "BACKEND", "MODE", "LAST USED"}),
37+
tablewriter.WithHeaderAlignment(tw.AlignLeft),
38+
tablewriter.WithRendition(tw.Rendition{
39+
Borders: tw.Border{Left: tw.Off, Right: tw.Off, Top: tw.Off, Bottom: tw.Off},
40+
Settings: tw.Settings{
41+
Separators: tw.Separators{
42+
ShowHeader: tw.Off,
43+
BetweenRows: tw.Off,
44+
BetweenColumns: tw.Off,
45+
},
46+
Lines: tw.Lines{
47+
ShowTop: tw.Off,
48+
ShowBottom: tw.Off,
49+
ShowHeaderLine: tw.Off,
50+
ShowFooterLine: tw.Off,
51+
},
52+
},
53+
}),
54+
tablewriter.WithPadding(tw.Padding{Left: "", Right: " "}),
55+
)
5156

5257
for _, status := range ps {
5358
modelName := status.ModelName

0 commit comments

Comments
 (0)