Skip to content

Commit a5ccc1e

Browse files
matttmattt
committed
Refactor tests and fixtures
1 parent 04650e1 commit a5ccc1e

File tree

7 files changed

+299
-446
lines changed

7 files changed

+299
-446
lines changed

‎Sources/AnyLanguageModel/LanguageModelSession.swift‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,16 @@ public final class LanguageModelSession: @unchecked Sendable {
159159
public let content: Content
160160
public let rawContent: GeneratedContent
161161
public let transcriptEntries: ArraySlice<Transcript.Entry>
162+
163+
init(
164+
content: Content,
165+
rawContent: GeneratedContent,
166+
transcriptEntries: ArraySlice<Transcript.Entry>
167+
) {
168+
self.content = content
169+
self.rawContent = rawContent
170+
self.transcriptEntries = transcriptEntries
171+
}
162172
}
163173

164174
@discardableResult

‎Tests/AnyLanguageModelTests/GenerableMacroTests.swift‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ private struct PrimitiveContainer {
5555
var count: Int
5656
}
5757

58+
@Generable
5859
private enum TestEnum: Equatable {
5960
case optionA
6061
case optionB

‎Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -328,5 +328,78 @@ import Testing
328328
#expect(error == .insufficientMemory)
329329
}
330330
}
331+
332+
@Test func structuredGenerationBasicStruct() async throws {
333+
let session = LanguageModelSession(
334+
model: model,
335+
instructions: "You are a helpful assistant that generates structured data."
336+
)
337+
let response = try await session.respond(
338+
to: "Generate a person with name Alice, age 30, active status true, and score 95.5",
339+
generating: BasicStruct.self
340+
)
341+
#expect(!response.content.name.isEmpty)
342+
#expect(response.content.age >= 0)
343+
}
344+
345+
@Test func structuredGenerationNestedStruct() async throws {
346+
let session = LanguageModelSession(
347+
model: model,
348+
instructions: "You are a helpful assistant that generates structured data."
349+
)
350+
let response = try await session.respond(
351+
to: "Generate a person named John, age 25, living at 123 Main St, Springfield, 12345",
352+
generating: StructuredPerson.self
353+
)
354+
#expect(!response.content.name.isEmpty)
355+
#expect(response.content.age >= 0)
356+
#expect(!response.content.address.street.isEmpty)
357+
#expect(!response.content.address.city.isEmpty)
358+
}
359+
360+
@Test func structuredGenerationStructWithEnum() async throws {
361+
let session = LanguageModelSession(
362+
model: model,
363+
instructions: "You are a helpful assistant that generates structured data."
364+
)
365+
let response = try await session.respond(
366+
to: "Generate a task titled 'Complete project' with high priority, not completed",
367+
generating: TaskItem.self
368+
)
369+
#expect(!response.content.title.isEmpty)
370+
#expect([Priority.low, Priority.medium, Priority.high].contains(response.content.priority))
371+
}
372+
373+
@Test func structuredGenerationSimpleArray() async throws {
374+
let session = LanguageModelSession(
375+
model: model,
376+
instructions: "You are a helpful assistant that generates structured data."
377+
)
378+
let response = try await session.respond(
379+
to: "Generate a list of 3 color names: red, green, blue",
380+
generating: SimpleArray.self
381+
)
382+
#expect(!response.content.colors.isEmpty)
383+
}
384+
385+
@Test func structuredGenerationStructWithArray() async throws {
386+
let session = LanguageModelSession(
387+
model: model,
388+
instructions: "You are a helpful assistant that generates structured data."
389+
)
390+
let response = try await session.respond(
391+
to: """
392+
Generate a quiz question:
393+
- Question: What is the capital of France?
394+
- Choices: London, Paris, Berlin, Madrid
395+
- Answer: Paris
396+
- Explanation: Paris is the capital city of France
397+
""",
398+
generating: MultiChoiceQuestion.self
399+
)
400+
#expect(!response.content.text.isEmpty)
401+
#expect(response.content.choices.count == 4)
402+
#expect(!response.content.answer.isEmpty)
403+
}
331404
}
332405
#endif // Llama

‎Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,5 +121,87 @@ import Testing
121121
let response = try await session.respond(to: "")
122122
#expect(!response.content.isEmpty)
123123
}
124+
125+
@Test func structuredGenerationSimpleString() async throws {
126+
let session = LanguageModelSession(
127+
model: model,
128+
instructions: "You are a helpful assistant that generates structured data."
129+
)
130+
let response = try await session.respond(
131+
to: "Generate a greeting message that says hello",
132+
generating: SimpleString.self
133+
)
134+
#expect(!response.content.message.isEmpty)
135+
}
136+
137+
@Test func structuredGenerationSimpleInt() async throws {
138+
let session = LanguageModelSession(
139+
model: model,
140+
instructions: "You are a helpful assistant that generates structured data."
141+
)
142+
let response = try await session.respond(
143+
to: "Generate a count value of 42",
144+
generating: SimpleInt.self
145+
)
146+
#expect(response.content.count >= 0)
147+
}
148+
149+
@Test func structuredGenerationSimpleDouble() async throws {
150+
let session = LanguageModelSession(
151+
model: model,
152+
instructions: "You are a helpful assistant that generates structured data."
153+
)
154+
let response = try await session.respond(
155+
to: "Generate a temperature value of 72.5 degrees",
156+
generating: SimpleDouble.self
157+
)
158+
#expect(!response.content.temperature.isNaN)
159+
}
160+
161+
@Test func structuredGenerationSimpleBool() async throws {
162+
let session = LanguageModelSession(
163+
model: model,
164+
instructions: "You are a helpful assistant that generates structured data."
165+
)
166+
let response = try await session.respond(
167+
to: "Generate a boolean value: true",
168+
generating: SimpleBool.self
169+
)
170+
let jsonData = response.rawContent.jsonString.data(using: .utf8)
171+
#expect(jsonData != nil)
172+
if let jsonData {
173+
let json = try JSONSerialization.jsonObject(with: jsonData)
174+
let dictionary = json as? [String: Any]
175+
let boolValue = dictionary?["value"] as? Bool
176+
#expect(boolValue != nil)
177+
}
178+
}
179+
180+
@Test func structuredGenerationOptionalFields() async throws {
181+
let session = LanguageModelSession(
182+
model: model,
183+
instructions: "You are a helpful assistant that generates structured data."
184+
)
185+
let response = try await session.respond(
186+
to: "Generate a person named Alex with nickname 'Lex'. Nickname may be omitted if unsure.",
187+
generating: OptionalFields.self
188+
)
189+
#expect(!response.content.name.isEmpty)
190+
if let nickname = response.content.nickname {
191+
#expect(!nickname.isEmpty)
192+
}
193+
}
194+
195+
@Test func structuredGenerationEnum() async throws {
196+
let session = LanguageModelSession(
197+
model: model,
198+
instructions: "You are a helpful assistant that generates structured data."
199+
)
200+
let response = try await session.respond(
201+
to: "Generate a high priority value",
202+
generating: Priority.self
203+
)
204+
#expect([Priority.low, Priority.medium, Priority.high].contains(response.content))
205+
}
124206
}
125207
#endif // MLX

‎Tests/AnyLanguageModelTests/Shared/StructuredGenerationFixtures.swift‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
import Foundation
2+
3+
@testable import AnyLanguageModel
4+
5+
@Generable
6+
enum Priority: Equatable {
7+
case low
8+
case medium
9+
case high
10+
}
11+
12+
@Generable
13+
struct SimpleString: Equatable {
14+
@Guide(description: "A greeting message")
15+
var message: String
16+
}
17+
18+
@Generable
19+
struct SimpleInt: Equatable {
20+
@Guide(description: "A count value", .minimum(0))
21+
var count: Int
22+
}
23+
24+
@Generable
25+
struct SimpleBool: Equatable {
26+
@Guide(description: "A boolean flag")
27+
var value: Bool
28+
}
29+
30+
@Generable
31+
struct SimpleDouble: Equatable {
32+
@Guide(description: "A temperature value")
33+
var temperature: Double
34+
}
35+
36+
@Generable
37+
struct OptionalFields: Equatable {
38+
@Guide(description: "A required name")
39+
var name: String
40+
41+
@Guide(description: "An optional nickname")
42+
var nickname: String?
43+
}
44+
45+
@Generable
46+
struct BasicStruct: Equatable {
47+
@Guide(description: "Person's name")
48+
var name: String
49+
50+
@Guide(description: "Person's age", .minimum(0))
51+
var age: Int
52+
53+
@Guide(description: "Is the person active")
54+
var isActive: Bool
55+
56+
@Guide(description: "Score value")
57+
var score: Double
58+
}
59+
60+
@Generable
61+
struct Address: Equatable {
62+
@Guide(description: "Street name")
63+
var street: String
64+
65+
@Guide(description: "City name")
66+
var city: String
67+
68+
@Guide(description: "Postal code")
69+
var postalCode: String
70+
}
71+
72+
@Generable
73+
struct ReusedNestedStruct: Equatable {
74+
@Guide(description: "Some text")
75+
var text: String
76+
}
77+
78+
@Generable
79+
struct ContainerWithDuplicateNestedType: Equatable {
80+
var first: ReusedNestedStruct
81+
var second: ReusedNestedStruct
82+
}
83+
84+
@Generable
85+
struct StructuredPerson: Equatable {
86+
@Guide(description: "Person's name")
87+
var name: String
88+
89+
@Guide(description: "Person's age")
90+
var age: Int
91+
92+
var address: Address
93+
}
94+
95+
@Generable
96+
struct TaskItem: Equatable {
97+
@Guide(description: "Task title")
98+
var title: String
99+
100+
var priority: Priority
101+
102+
@Guide(description: "Is completed")
103+
var isCompleted: Bool
104+
}
105+
106+
@Generable
107+
struct SimpleArray: Equatable {
108+
@Guide(description: "A list of color names")
109+
var colors: [String]
110+
}
111+
112+
@Generable
113+
struct MultiChoiceQuestion: Equatable {
114+
@Guide(description: "The quiz question")
115+
var text: String
116+
117+
@Guide(.count(4))
118+
var choices: [String]
119+
120+
var answer: String
121+
122+
@Guide(description: "A brief explanation of why the answer is correct")
123+
var explanation: String
124+
}

0 commit comments

Comments
 (0)