Skip to content

Commit b49015d

Browse files
1amageek1amageek
committed
Refactors chat API to use new OKTool structure
1 parent 183be24 commit b49015d

File tree

3 files changed

+110
-143
lines changed

3 files changed

+110
-143
lines changed

‎Sources/OllamaKit/OllamaKit+Chat.swift‎

Lines changed: 61 additions & 140 deletions
Original file line numberDiff line numberDiff line change
@@ -9,95 +9,52 @@ import Combine
99
import Foundation
1010

1111
extension OllamaKit {
12-
/// Establishes an asynchronous stream for chat responses from the Ollama API, based on the provided data.
12+
/// Starts a stream for chat responses from the Ollama API.
1313
///
14-
/// This method sets up a streaming connection using Swift's concurrency features, allowing for real-time data handling as chat responses are generated by the Ollama API.
15-
///
16-
/// Example usage
17-
///
18-
/// ```swift
19-
/// let ollamaKit = OllamaKit()
20-
/// let chatData = OKChatRequestData(/* parameters */)
21-
///
22-
/// Task {
23-
/// do {
24-
/// for try await response in ollamaKit.chat(data: chatData) {
25-
/// // Handle each chat response
26-
/// }
27-
/// } catch {
28-
/// // Handle error
29-
/// }
30-
/// }
31-
/// ```
32-
///
33-
/// Example usage with tools
14+
/// This method allows real-time handling of chat responses using Swift's concurrency.
3415
///
16+
/// Example usage:
3517
/// ```swift
3618
/// let ollamaKit = OllamaKit()
3719
/// let chatData = OKChatRequestData(
38-
/// /* parameters */,
20+
/// model: "example-model",
21+
/// messages: [
22+
/// .user("What's the weather like in Tokyo?")
23+
/// ],
3924
/// tools: [
40-
/// .object([
41-
/// "type": .string("function"),
42-
/// "function": .object([
43-
/// "name": .string("get_current_weather"),
44-
/// "description": .string("Get the current weather for a location"),
45-
/// "parameters": .object([
46-
/// "type": .string("object"),
47-
/// "properties": .object([
48-
/// "location": .object([
49-
/// "type": .string("string"),
50-
/// "description": .string("The location to get the weather for, e.g. San Francisco, CA")
51-
/// ]),
52-
/// "format": .object([
53-
/// "type": .string("string"),
54-
/// "description": .string("The format to return the weather in, e.g. 'celsius' or 'fahrenheit'"),
55-
/// "enum": .array([.string("celsius"), .string("fahrenheit")])
56-
/// ])
25+
/// .function(
26+
/// OKFunction(
27+
/// name: "get_current_weather",
28+
/// description: "Fetch current weather information.",
29+
/// parameters: .object([
30+
/// "location": .object([
31+
/// "type": .string("string"),
32+
/// "description": .string("The location to get the weather for, e.g., Tokyo")
5733
/// ]),
58-
/// "required": .array([.string("location"), .string("format")])
59-
/// ])
60-
/// ])
61-
/// ])
34+
/// "format": .object([
35+
/// "type": .string("string"),
36+
/// "description": .string("The format for the weather, e.g., 'celsius'."),
37+
/// "enum": .array([.string("celsius"), .string("fahrenheit")])
38+
/// ])
39+
/// ], required: ["location", "format"])
40+
/// )
41+
/// )
6242
/// ]
6343
/// )
6444
///
6545
/// Task {
6646
/// do {
6747
/// for try await response in ollamaKit.chat(data: chatData) {
68-
/// if let toolCalls = response.message?.toolCalls {
69-
/// for toolCall in toolCalls {
70-
/// if let function = toolCall.function {
71-
/// print("Tool called: \(function.name ?? "")")
72-
///
73-
/// if let arguments = function.arguments {
74-
/// switch arguments {
75-
/// case .object(let argDict):
76-
/// if let location = argDict["location"], case .string(let locationValue) = location {
77-
/// print("Location: \(locationValue)")
78-
/// }
79-
///
80-
/// if let format = argDict["format"], case .string(let formatValue) = format {
81-
/// print("Format: \(formatValue)")
82-
/// }
83-
/// default:
84-
/// print("Unexpected arguments format")
85-
/// }
86-
/// } else {
87-
/// print("No arguments provided")
88-
/// }
89-
/// }
90-
/// }
91-
/// }
48+
/// // Handle each response here
49+
/// print(response)
9250
/// }
9351
/// } catch {
94-
/// // Handle error
52+
/// print("Error: \(error)")
9553
/// }
9654
/// }
9755
/// ```
98-
///
99-
/// - Parameter data: The ``OKChatRequestData`` used to initiate the chat streaming from the Ollama API.
100-
/// - Returns: An `AsyncThrowingStream<OKChatResponse, Error>` emitting the live stream of chat responses from the Ollama API.
56+
/// - Parameter data: The ``OKChatRequestData`` containing chat request details.
57+
/// - Returns: An `AsyncThrowingStream<OKChatResponse, Error>` emitting chat responses from the Ollama API.
10158
public func chat(data: OKChatRequestData) -> AsyncThrowingStream<OKChatResponse, Error> {
10259
do {
10360
let request = try OKRouter.chat(data: data).asURLRequest(with: baseURL)
@@ -109,91 +66,55 @@ extension OllamaKit {
10966
}
11067
}
11168

112-
/// Establishes a Combine publisher for streaming chat responses from the Ollama API, based on the provided data.
113-
///
114-
/// This method sets up a streaming connection using the Combine framework, facilitating real-time data handling as chat responses are generated by the Ollama API.
69+
/// Publishes a stream of chat responses from the Ollama API using Combine.
11570
///
116-
/// Example usage
117-
///
118-
/// ```swift
119-
/// let ollamaKit = OllamaKit()
120-
/// let chatData = OKChatRequestData(/* parameters */)
121-
///
122-
/// ollamaKit.chat(data: chatData)
123-
/// .sink(receiveCompletion: { completion in
124-
/// // Handle completion or error
125-
/// }, receiveValue: { chatResponse in
126-
/// // Handle each chat response
127-
/// })
128-
/// .store(in: &cancellables)
129-
/// ```
130-
///
131-
/// Example usage with tools
71+
/// Enables real-time data handling through Combine's reactive streams.
13272
///
73+
/// Example usage:
13374
/// ```swift
13475
/// let ollamaKit = OllamaKit()
13576
/// let chatData = OKChatRequestData(
136-
/// /* parameters */,
77+
/// model: "example-model",
78+
/// messages: [
79+
/// .user("What's the weather like in Tokyo?")
80+
/// ],
13781
/// tools: [
138-
/// .object([
139-
/// "type": .string("function"),
140-
/// "function": .object([
141-
/// "name": .string("get_current_weather"),
142-
/// "description": .string("Get the current weather for a location"),
143-
/// "parameters": .object([
144-
/// "type": .string("object"),
145-
/// "properties": .object([
146-
/// "location": .object([
147-
/// "type": .string("string"),
148-
/// "description": .string("The location to get the weather for, e.g. San Francisco, CA")
149-
/// ]),
150-
/// "format": .object([
151-
/// "type": .string("string"),
152-
/// "description": .string("The format to return the weather in, e.g. 'celsius' or 'fahrenheit'"),
153-
/// "enum": .array([.string("celsius"), .string("fahrenheit")])
154-
/// ])
82+
/// .function(
83+
/// OKFunction(
84+
/// name: "get_current_weather",
85+
/// description: "Fetch current weather information.",
86+
/// parameters: .object([
87+
/// "location": .object([
88+
/// "type": .string("string"),
89+
/// "description": .string("The location to get the weather for, e.g., Tokyo")
15590
/// ]),
156-
/// "required": .array([.string("location"), .string("format")])
157-
/// ])
158-
/// ])
159-
/// ])
91+
/// "format": .object([
92+
/// "type": .string("string"),
93+
/// "description": .string("The format for the weather, e.g., 'celsius'."),
94+
/// "enum": .array([.string("celsius"), .string("fahrenheit")])
95+
/// ])
96+
/// ], required: ["location", "format"])
97+
/// )
98+
/// )
16099
/// ]
161100
/// )
162101
///
163102
/// ollamaKit.chat(data: chatData)
164103
/// .sink(receiveCompletion: { completion in
165-
/// // Handle completion or error
166-
/// }, receiveValue: { chatResponse in
167-
/// if let toolCalls = chatResponse.message?.toolCalls {
168-
/// for toolCall in toolCalls {
169-
/// if let function = toolCall.function {
170-
/// print("Tool called: \(function.name ?? "")")
171-
///
172-
/// if let arguments = function.arguments {
173-
/// switch arguments {
174-
/// case .object(let argDict):
175-
/// if let location = argDict["location"], case .string(let locationValue) = location {
176-
/// print("Location: \(locationValue)")
177-
/// }
178-
///
179-
/// if let format = argDict["format"], case .string(let formatValue) = format {
180-
/// print("Format: \(formatValue)")
181-
/// }
182-
/// default:
183-
/// print("Unexpected arguments format")
184-
/// }
185-
/// } else {
186-
/// print("No arguments provided")
187-
/// }
188-
/// }
189-
/// }
104+
/// switch completion {
105+
/// case .finished:
106+
/// print("Stream finished")
107+
/// case .failure(let error):
108+
/// print("Error: \(error)")
190109
/// }
110+
/// }, receiveValue: { response in
111+
/// // Handle each response here
112+
/// print(response)
191113
/// })
192114
/// .store(in: &cancellables)
193115
/// ```
194-
///
195-
/// - Parameter data: The ``OKChatRequestData`` used to initiate the chat streaming from the Ollama API.
196-
/// - Returns: An `AnyPublisher<OKChatResponse, Error>` emitting the live stream of chat responses from the Ollama API.
116+
/// - Parameter data: The ``OKChatRequestData`` containing chat request details.
117+
/// - Returns: An `AnyPublisher<OKChatResponse, Error>` emitting chat responses from the Ollama API.
197118
public func chat(data: OKChatRequestData) -> AnyPublisher<OKChatResponse, Error> {
198119
do {
199120
let request = try OKRouter.chat(data: data).asURLRequest(with: baseURL)

‎Sources/OllamaKit/RequestData/OKChatRequestData.swift‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ public struct OKChatRequestData: Sendable {
1818
public let messages: [Message]
1919

2020
/// An optional array of ``OKJSONValue`` representing the tools available for tool calling in the chat.
21-
public let tools: [OKJSONValue]?
21+
public let tools: [OKTool]?
2222

2323
/// Optional ``OKJSONValue`` representing the JSON schema for the response.
2424
/// Be sure to also include "return as JSON" in your prompt
@@ -31,7 +31,7 @@ public struct OKChatRequestData: Sendable {
3131
public init(
3232
model: String,
3333
messages: [Message],
34-
tools: [OKJSONValue]? = nil,
34+
tools: [OKTool]? = nil,
3535
format: OKJSONValue? = nil,
3636
options: OKCompletionOptions? = nil
3737
) {
@@ -46,7 +46,7 @@ public struct OKChatRequestData: Sendable {
4646
public init(
4747
model: String,
4848
messages: [Message],
49-
tools: [OKJSONValue]? = nil,
49+
tools: [OKTool]? = nil,
5050
format: OKJSONValue? = nil,
5151
with configureOptions: @Sendable (inout OKCompletionOptions) -> Void
5252
) {

‎Sources/OllamaKit/Utils/OKTool.swift‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
//
2+
// OKTool.swift
3+
// OllamaKit
4+
//
5+
// Created by Norikazu Muramoto on 2025/01/11.
6+
//
7+
8+
import Foundation
9+
10+
/// Represents a tool that can be used in the Ollama API chat.
11+
public struct OKTool: Encodable, Sendable {
12+
/// The type of the tool (e.g., "function").
13+
public let type: String
14+
15+
/// The function details associated with the tool.
16+
public let function: OKFunction
17+
18+
public init(type: String, function: OKFunction) {
19+
self.type = type
20+
self.function = function
21+
}
22+
23+
/// Convenience method for creating a tool with type "function".
24+
public static func function(_ function: OKFunction) -> OKTool {
25+
return OKTool(type: "function", function: function)
26+
}
27+
}
28+
29+
/// Represents a function used as a tool in the Ollama API chat.
30+
public struct OKFunction: Encodable, Sendable {
31+
/// The name of the function.
32+
public let name: String
33+
34+
/// A description of what the function does.
35+
public let description: String
36+
37+
/// Parameters required by the function, defined as a JSON schema.
38+
public let parameters: OKJSONValue
39+
40+
public init(name: String, description: String, parameters: OKJSONValue) {
41+
self.name = name
42+
self.description = description
43+
self.parameters = parameters
44+
}
45+
}
46+

0 commit comments

Comments
 (0)