Skip to content

Commit

Permalink
test: added tests for tools
Browse files Browse the repository at this point in the history 
Signed-off-by: rawnly <[email protected]>
  • Loading branch information
@rawnly
rawnly committed Nov 19, 2023
1 parent 723a688 commit 5e90958
Show file tree
Hide file tree
Showing 3 changed files with 170 additions and 19 deletions.
68 changes: 56 additions & 12 deletions Tests/OpenAITests/OpenAITests.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,25 +104,69 @@ class OpenAITests: XCTestCase {

func testChats() async throws {
let query = ChatQuery(model: .gpt4, messages: [
.init(role: .system, content: "You are Librarian-GPT. You know everything about the books."),
.init(role: .user, content: "Who wrote Harry Potter?")
.init(role: .system, content: "You are Librarian-GPT. You know everything about the books.", toolCalls: nil),
.init(role: .user, content: "Who wrote Harry Potter?", toolCalls: nil),
])
let chatResult = ChatResult(id: "id-12312", object: "foo", created: 100, model: .gpt3_5Turbo, choices: [
.init(index: 0, message: .init(role: .system, content: "bar"), finishReason: "baz"),
.init(index: 0, message: .init(role: .user, content: "bar1"), finishReason: "baz1"),
.init(index: 0, message: .init(role: .assistant, content: "bar2"), finishReason: "baz2")
.init(index: 0, message: .init(role: .system, content: "bar", toolCalls: nil), finishReason: "baz"),
.init(index: 0, message: .init(role: .user, content: "bar1", toolCalls: nil), finishReason: "baz1"),
.init(index: 0, message: .init(role: .assistant, content: "bar2", toolCalls: nil), finishReason: "baz2")
], usage: .init(promptTokens: 100, completionTokens: 200, totalTokens: 300))
try self.stub(result: chatResult)

let result = try await openAI.chats(query: query)

XCTAssertEqual(result, chatResult)
}

func testChatsTools() async throws {
let tools: [ChatTool] = [
ChatTool(type: .function, value: .function(
.init(
name: "get_weather",
description: "Get the current weather in the given location",
parameters: .init(
type: .object,
properties: [
"location": .string(description: "The city and state, e.g San Francisco, CA"),
"unit": .string(enumValues: ["celsius", "fahrenheit"])
],
required: ["location"]
)
)
))
]

let messages: [Message] = [
.init(role: .system, content: "You are Weather-GPT. You know everything about the weather.", toolCalls: nil),
.init(role: .user, content: "What's the weather like in Boston?", toolCalls: nil),
]

let query = ChatQuery(model: .gpt3_5Turbo_1106, messages: messages, tools: tools)

let chatResult = ChatResult(
id: "id-12312",
object: "foo",
created: 100,
model: .gpt3_5Turbo,
choices: [
.init(index: 0, message: .init(role: .system, content: "bar", toolCalls: nil), finishReason: "baz"),
.init(index: 0, message: .init(role: .user, content: "bar1", toolCalls: nil), finishReason: "baz1"),
.init(index: 0, message: .init(role: .assistant, content: "bar2", toolCalls: nil), finishReason: "baz2")
],
usage: .init(promptTokens: 100, completionTokens: 200, totalTokens: 300)
)

try self.stub(result: chatResult)

let result = try await openAI.chats(query: query)
XCTAssertEqual(result, chatResult)
}

func testChatsFunction() async throws {
let query = ChatQuery(model: .gpt3_5Turbo_1106, messages: [
.init(role: .system, content: "You are Weather-GPT. You know everything about the weather."),
.init(role: .user, content: "What's the weather like in Boston?"),
.init(role: .system, content: "You are Weather-GPT. You know everything about the weather.", toolCalls: nil),
.init(role: .user, content: "What's the weather like in Boston?", toolCalls: nil),
], functions: [
.init(name: "get_current_weather", description: "Get the current weather in a given location", parameters: .init(type: .object, properties: [
"location": .init(type: .string, description: "The city and state, e.g. San Francisco, CA"),
Expand All @@ -131,9 +175,9 @@ class OpenAITests: XCTestCase {
], functionCall: .auto)

let chatResult = ChatResult(id: "id-12312", object: "foo", created: 100, model: .gpt3_5Turbo, choices: [
.init(index: 0, message: .init(role: .system, content: "bar"), finishReason: "baz"),
.init(index: 0, message: .init(role: .user, content: "bar1"), finishReason: "baz1"),
.init(index: 0, message: .init(role: .assistant, content: "bar2"), finishReason: "baz2")
.init(index: 0, message: .init(role: .system, content: "bar", toolCalls: nil), finishReason: "baz"),
.init(index: 0, message: .init(role: .user, content: "bar1", toolCalls: nil), finishReason: "baz1"),
.init(index: 0, message: .init(role: .assistant, content: "bar2", toolCalls: nil), finishReason: "baz2")
], usage: .init(promptTokens: 100, completionTokens: 200, totalTokens: 300))
try self.stub(result: chatResult)

Expand All @@ -143,8 +187,8 @@ class OpenAITests: XCTestCase {

func testChatsError() async throws {
let query = ChatQuery(model: .gpt4, messages: [
.init(role: .system, content: "You are Librarian-GPT. You know everything about the books."),
.init(role: .user, content: "Who wrote Harry Potter?")
.init(role: .system, content: "You are Librarian-GPT. You know everything about the books.", toolCalls: nil),
.init(role: .user, content: "Who wrote Harry Potter?", toolCalls: nil)
])
let inError = APIError(message: "foo", type: "bar", param: "baz", code: "100")
self.stub(error: inError)
Expand Down
10 changes: 5 additions & 5 deletions Tests/OpenAITests/OpenAITestsCombine.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,13 +39,13 @@ final class OpenAITestsCombine: XCTestCase {

func testChats() throws {
let query = ChatQuery(model: .gpt4, messages: [
.init(role: .system, content: "You are Librarian-GPT. You know everything about the books."),
.init(role: .user, content: "Who wrote Harry Potter?")
.init(role: .system, content: "You are Librarian-GPT. You know everything about the books.", toolCalls: nil),
.init(role: .user, content: "Who wrote Harry Potter?", toolCalls: nil)
])
let chatResult = ChatResult(id: "id-12312", object: "foo", created: 100, model: .gpt3_5Turbo, choices: [
.init(index: 0, message: .init(role: .system, content: "bar"), finishReason: "baz"),
.init(index: 0, message: .init(role: .user, content: "bar1"), finishReason: "baz1"),
.init(index: 0, message: .init(role: .assistant, content: "bar2"), finishReason: "baz2")
.init(index: 0, message: .init(role: .system, content: "bar", toolCalls: nil), finishReason: "baz"),
.init(index: 0, message: .init(role: .user, content: "bar1", toolCalls: nil), finishReason: "baz1"),
.init(index: 0, message: .init(role: .assistant, content: "bar2", toolCalls: nil), finishReason: "baz2")
], usage: .init(promptTokens: 100, completionTokens: 200, totalTokens: 300))
try self.stub(result: chatResult)
let result = try awaitPublisher(openAI.chats(query: query))
Expand Down
111 changes: 109 additions & 2 deletions Tests/OpenAITests/OpenAITestsDecoder.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ class OpenAITestsDecoder: XCTestCase {
"""

let expectedValue = ChatResult(id: "chatcmpl-123", object: "chat.completion", created: 1677652288, model: .gpt4, choices: [
.init(index: 0, message: Message(role: .assistant, content: "Hello, world!"), finishReason: "stop")
.init(index: 0, message: Message(role: .assistant, content: "Hello, world!", toolCalls: nil), finishReason: "stop")
], usage: .init(promptTokens: 9, completionTokens: 12, totalTokens: 21))
try decode(data, expectedValue)
}
Expand Down Expand Up @@ -140,12 +140,57 @@ class OpenAITestsDecoder: XCTestCase {

XCTAssertEqual(imageQueryAsDict, expectedValueAsDict)
}

func testToolEncode() async throws {
let value: ChatTool = ChatTool(type: .function, value: .function(
.init(
name: "test",
description: "test",
parameters: .init(
type: .object,
properties: [
"location": .string(description: "location.desc"),
"unit": .string(enumValues: ["unit.A", "unit.B"])
],
required: ["location"]
)
)
))


let expectedValue = """
{
"type": "function",
"function": {
"name": "test",
"description": "test",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "location.desc"
},
"unit": { "type": "string", "enum": ["unit.A", "unit.B"] }
},
"required": ["location"]
}
}
}
"""

// To compare serialized JSONs we first convert them both into NSDictionary which are comparable (unline native swift dictionaries)
let result = try jsonDataAsNSDictionary(JSONEncoder().encode(value))
let expectedResult = try jsonDataAsNSDictionary(expectedValue.data(using: .utf8)!)

XCTAssertEqual(result, expectedResult)
}

func testChatQueryWithFunctionCall() async throws {
let chatQuery = ChatQuery(
model: .gpt3_5Turbo,
messages: [
Message(role: .user, content: "What's the weather like in Boston?")
Message(role: .user, content: "What's the weather like in Boston?", toolCalls: nil)
],
responseFormat: .init(type: .jsonObject),
functions: [
Expand Down Expand Up @@ -236,6 +281,68 @@ class OpenAITestsDecoder: XCTestCase {

XCTAssertEqual(resultDict, expectedValueAsDict)
}

func testChatResultWithToolCall() async throws {
let data = """
{
"id": "chatcmpl-1234",
"object": "chat.completion",
"created": 1677652288,
"model": "gpt-3.5-turbo",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": null,
"tool_calls": [
{
"index": 0,
"id": "id",
"type": "function",
"function": {
"name": "get_current_weather"
}
}
]
},
"finish_reason": "stop"
}
],
"usage": {
"prompt_tokens": 82,
"completion_tokens": 18,
"total_tokens": 100
}
}
"""

let expectedValue = ChatResult(
id: "chatcmpl-1234",
object: "chat.completion",
created: 1677652288,
model: .gpt3_5Turbo,
choices: [
.init(
index: 0,
message: Message(role: .assistant, toolCalls: [
.init(
index: 0,
id: "id",
type: .function,
value: .function(
.withName("get_current_weather")
)
)
]),
finishReason: "stop"
)
],
usage: .init(promptTokens: 82, completionTokens: 18, totalTokens: 100)
)

try decode(data, expectedValue)
}

func testChatCompletionWithFunctionCall() async throws {
let data = """
Expand Down

0 comments on commit 5e90958

Please sign in to comment.