Skip to content

Commit 4428c55

Browse files
committed
nicer model output
1 parent 3ed199a commit 4428c55

File tree

7 files changed

+361
-5
lines changed

7 files changed

+361
-5
lines changed

Package.resolved

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Sources/GenKit/Services/Mistral/MistralDecoders.swift

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,4 +72,141 @@ extension MistralService {
7272
return .cancelled
7373
}
7474
}
75+
76+
func decode(model: Mistral.ModelResponse) -> Model {
77+
var name = String?.none
78+
var family = String?.none
79+
var maxOutput = Int?.none
80+
var contextWindow = Int?.none
81+
82+
if model.id.hasPrefix("codestral") {
83+
family = "Codestral"
84+
switch model.id {
85+
case "codestral-latest":
86+
name = "Codestral (latest)"
87+
case "codestral-2405":
88+
name = "Codestral (2405)"
89+
case "codestral-mamba-latest":
90+
name = "Codestral Mamba (latest)"
91+
case "codestral-mamba-2407":
92+
name = "Codestral Mamba (2407)"
93+
default:
94+
name = model.id
95+
}
96+
}
97+
98+
if model.id.hasPrefix("mistral-large") {
99+
family = "Mistral Large"
100+
switch model.id {
101+
case "mistral-large-latest":
102+
name = "Mistral Large (latest)"
103+
case "mistral-large-2402":
104+
name = "Mistral Large (2402)"
105+
case "mistral-large-2407":
106+
name = "Mistral Large (2407)"
107+
default:
108+
name = model.id
109+
}
110+
}
111+
112+
if model.id.hasPrefix("mistral-medium") {
113+
family = "Mistral Medium"
114+
switch model.id {
115+
case "mistral-medium":
116+
name = "Mistral Medium"
117+
case "mistral-medium-latest":
118+
name = "Mistral Medium (latest)"
119+
case "mistral-medium-2312":
120+
name = "Mistral Medium (2312)"
121+
default:
122+
name = model.id
123+
}
124+
}
125+
126+
if model.id.hasPrefix("mistral-small") {
127+
family = "Mistral Small"
128+
switch model.id {
129+
case "mistral-small":
130+
name = "Mistral Small"
131+
case "mistral-small-latest":
132+
name = "Mistral Small (latest)"
133+
case "mistral-small-2312":
134+
name = "Mistral Small (2312)"
135+
case "mistral-small-2402":
136+
name = "Mistral Small (2402)"
137+
case "mistral-small-2409":
138+
name = "Mistral Small (2409)"
139+
default:
140+
name = model.id
141+
}
142+
}
143+
144+
if model.id.hasPrefix("mistral-tiny") {
145+
family = "Mistral Tiny"
146+
switch model.id {
147+
case "mistral-tiny":
148+
name = "Mistral Tiny"
149+
case "mistral-tiny-latest":
150+
name = "Mistral Tiny (latest)"
151+
case "mistral-tiny-2312":
152+
name = "Mistral Tiny (2312)"
153+
case "mistral-tiny-2407":
154+
name = "Mistral Tiny (2407)"
155+
default:
156+
name = model.id
157+
}
158+
}
159+
160+
if model.id.hasPrefix("pixtral") {
161+
family = "Pixtral"
162+
switch model.id {
163+
case "pixtral-12b":
164+
name = "Pixtral 12b"
165+
case "pixtral-12b-latest":
166+
name = "Pixtral 12b (latest)"
167+
case "pixtral-12b-2409":
168+
name = "Pixtral 12b (2409)"
169+
default:
170+
name = model.id
171+
}
172+
}
173+
174+
if model.id.hasPrefix("open-mistral") {
175+
family = "Open Mistral"
176+
switch model.id {
177+
case "open-mistral-7b":
178+
name = "Open Mistral 7b"
179+
case "open-mistral-nemo":
180+
name = "Open Mistral Nemo"
181+
case "open-mistral-nemo-2407":
182+
name = "Open Mistral Nemo (2407)"
183+
default:
184+
name = model.id
185+
}
186+
}
187+
188+
if model.id.hasPrefix("open-mixtral") {
189+
family = "Open Mixtral"
190+
switch model.id {
191+
case "open-mixtral-8x22b":
192+
name = "Open Mixtral 8x22b"
193+
case "open-mixtral-8x22b-2404":
194+
name = "Open Mixtral 8x22b (2404)"
195+
case "open-mixtral-8x7b":
196+
name = "Open Mixtral 8x7b"
197+
default:
198+
name = model.id
199+
}
200+
}
201+
202+
return Model(
203+
id: model.id,
204+
family: family,
205+
name: name,
206+
owner: model.ownedBy,
207+
contextWindow: contextWindow,
208+
maxOutput: maxOutput,
209+
trainingCutoff: nil
210+
)
211+
}
75212
}

Sources/GenKit/Services/Ollama/Ollama.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ extension OllamaService: ModelService {
7878

7979
public func models() async throws -> [Model] {
8080
let result = try await client.models()
81-
return result.models.map { Model(id: $0.name, owner: "ollama") }
81+
return result.models.map { decode(model: $0) }
8282
}
8383
}
8484

Sources/GenKit/Services/Ollama/OllamaDecoders.swift

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,4 +54,16 @@ extension OllamaService {
5454
guard let done else { return nil }
5555
return done ? .stop : nil
5656
}
57+
58+
func decode(model: Ollama.ModelResponse) -> Model {
59+
.init(
60+
id: model.model,
61+
family: model.details?.family,
62+
name: model.name,
63+
owner: "ollama",
64+
contextWindow: nil,
65+
maxOutput: nil,
66+
trainingCutoff: nil
67+
)
68+
}
5769
}

Sources/GenKit/Services/OpenAI/OpenAI.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ extension OpenAIService: ModelService {
5656

5757
public func models() async throws -> [Model] {
5858
let result = try await client.models()
59-
return result.data.map { Model(id: $0.id, owner: $0.ownedBy) }
59+
return result.data.map { decode(model: $0) }
6060
}
6161
}
6262

0 commit comments

Comments
 (0)