Skip to content

Commit

Permalink
Refactors chat API to use new OKTool structure
Browse files Browse the repository at this point in the history
  • Loading branch information
@1amageek
1amageek committed Jan 11, 2025
1 parent 183be24 commit b49015d
Show file tree
Hide file tree
Showing 3 changed files with 110 additions and 143 deletions.
201 changes: 61 additions & 140 deletions Sources/OllamaKit/OllamaKit+Chat.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,95 +9,52 @@ import Combine
import Foundation

extension OllamaKit {
/// Establishes an asynchronous stream for chat responses from the Ollama API, based on the provided data.
/// Starts a stream for chat responses from the Ollama API.
///
/// 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.
///
/// Example usage
///
/// ```swift
/// let ollamaKit = OllamaKit()
/// let chatData = OKChatRequestData(/* parameters */)
///
/// Task {
/// do {
/// for try await response in ollamaKit.chat(data: chatData) {
/// // Handle each chat response
/// }
/// } catch {
/// // Handle error
/// }
/// }
/// ```
///
/// Example usage with tools
/// This method allows real-time handling of chat responses using Swift's concurrency.
///
/// Example usage:
/// ```swift
/// let ollamaKit = OllamaKit()
/// let chatData = OKChatRequestData(
/// /* parameters */,
/// model: "example-model",
/// messages: [
/// .user("What's the weather like in Tokyo?")
/// ],
/// tools: [
/// .object([
/// "type": .string("function"),
/// "function": .object([
/// "name": .string("get_current_weather"),
/// "description": .string("Get the current weather for a location"),
/// "parameters": .object([
/// "type": .string("object"),
/// "properties": .object([
/// "location": .object([
/// "type": .string("string"),
/// "description": .string("The location to get the weather for, e.g. San Francisco, CA")
/// ]),
/// "format": .object([
/// "type": .string("string"),
/// "description": .string("The format to return the weather in, e.g. 'celsius' or 'fahrenheit'"),
/// "enum": .array([.string("celsius"), .string("fahrenheit")])
/// ])
/// .function(
/// OKFunction(
/// name: "get_current_weather",
/// description: "Fetch current weather information.",
/// parameters: .object([
/// "location": .object([
/// "type": .string("string"),
/// "description": .string("The location to get the weather for, e.g., Tokyo")
/// ]),
/// "required": .array([.string("location"), .string("format")])
/// ])
/// ])
/// ])
/// "format": .object([
/// "type": .string("string"),
/// "description": .string("The format for the weather, e.g., 'celsius'."),
/// "enum": .array([.string("celsius"), .string("fahrenheit")])
/// ])
/// ], required: ["location", "format"])
/// )
/// )
/// ]
/// )
///
/// Task {
/// do {
/// for try await response in ollamaKit.chat(data: chatData) {
/// if let toolCalls = response.message?.toolCalls {
/// for toolCall in toolCalls {
/// if let function = toolCall.function {
/// print("Tool called: \(function.name ?? "")")
///
/// if let arguments = function.arguments {
/// switch arguments {
/// case .object(let argDict):
/// if let location = argDict["location"], case .string(let locationValue) = location {
/// print("Location: \(locationValue)")
/// }
///
/// if let format = argDict["format"], case .string(let formatValue) = format {
/// print("Format: \(formatValue)")
/// }
/// default:
/// print("Unexpected arguments format")
/// }
/// } else {
/// print("No arguments provided")
/// }
/// }
/// }
/// }
/// // Handle each response here
/// print(response)
/// }
/// } catch {
/// // Handle error
/// print("Error: \(error)")
/// }
/// }
/// ```
///
/// - Parameter data: The ``OKChatRequestData`` used to initiate the chat streaming from the Ollama API.
/// - Returns: An `AsyncThrowingStream<OKChatResponse, Error>` emitting the live stream of chat responses from the Ollama API.
/// - Parameter data: The ``OKChatRequestData`` containing chat request details.
/// - Returns: An `AsyncThrowingStream<OKChatResponse, Error>` emitting chat responses from the Ollama API.
public func chat(data: OKChatRequestData) -> AsyncThrowingStream<OKChatResponse, Error> {
do {
let request = try OKRouter.chat(data: data).asURLRequest(with: baseURL)
Expand All @@ -109,91 +66,55 @@ extension OllamaKit {
}
}

/// Establishes a Combine publisher for streaming chat responses from the Ollama API, based on the provided data.
///
/// 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.
/// Publishes a stream of chat responses from the Ollama API using Combine.
///
/// Example usage
///
/// ```swift
/// let ollamaKit = OllamaKit()
/// let chatData = OKChatRequestData(/* parameters */)
///
/// ollamaKit.chat(data: chatData)
/// .sink(receiveCompletion: { completion in
/// // Handle completion or error
/// }, receiveValue: { chatResponse in
/// // Handle each chat response
/// })
/// .store(in: &cancellables)
/// ```
///
/// Example usage with tools
/// Enables real-time data handling through Combine's reactive streams.
///
/// Example usage:
/// ```swift
/// let ollamaKit = OllamaKit()
/// let chatData = OKChatRequestData(
/// /* parameters */,
/// model: "example-model",
/// messages: [
/// .user("What's the weather like in Tokyo?")
/// ],
/// tools: [
/// .object([
/// "type": .string("function"),
/// "function": .object([
/// "name": .string("get_current_weather"),
/// "description": .string("Get the current weather for a location"),
/// "parameters": .object([
/// "type": .string("object"),
/// "properties": .object([
/// "location": .object([
/// "type": .string("string"),
/// "description": .string("The location to get the weather for, e.g. San Francisco, CA")
/// ]),
/// "format": .object([
/// "type": .string("string"),
/// "description": .string("The format to return the weather in, e.g. 'celsius' or 'fahrenheit'"),
/// "enum": .array([.string("celsius"), .string("fahrenheit")])
/// ])
/// .function(
/// OKFunction(
/// name: "get_current_weather",
/// description: "Fetch current weather information.",
/// parameters: .object([
/// "location": .object([
/// "type": .string("string"),
/// "description": .string("The location to get the weather for, e.g., Tokyo")
/// ]),
/// "required": .array([.string("location"), .string("format")])
/// ])
/// ])
/// ])
/// "format": .object([
/// "type": .string("string"),
/// "description": .string("The format for the weather, e.g., 'celsius'."),
/// "enum": .array([.string("celsius"), .string("fahrenheit")])
/// ])
/// ], required: ["location", "format"])
/// )
/// )
/// ]
/// )
///
/// ollamaKit.chat(data: chatData)
/// .sink(receiveCompletion: { completion in
/// // Handle completion or error
/// }, receiveValue: { chatResponse in
/// if let toolCalls = chatResponse.message?.toolCalls {
/// for toolCall in toolCalls {
/// if let function = toolCall.function {
/// print("Tool called: \(function.name ?? "")")
///
/// if let arguments = function.arguments {
/// switch arguments {
/// case .object(let argDict):
/// if let location = argDict["location"], case .string(let locationValue) = location {
/// print("Location: \(locationValue)")
/// }
///
/// if let format = argDict["format"], case .string(let formatValue) = format {
/// print("Format: \(formatValue)")
/// }
/// default:
/// print("Unexpected arguments format")
/// }
/// } else {
/// print("No arguments provided")
/// }
/// }
/// }
/// switch completion {
/// case .finished:
/// print("Stream finished")
/// case .failure(let error):
/// print("Error: \(error)")
/// }
/// }, receiveValue: { response in
/// // Handle each response here
/// print(response)
/// })
/// .store(in: &cancellables)
/// ```
///
/// - Parameter data: The ``OKChatRequestData`` used to initiate the chat streaming from the Ollama API.
/// - Returns: An `AnyPublisher<OKChatResponse, Error>` emitting the live stream of chat responses from the Ollama API.
/// - Parameter data: The ``OKChatRequestData`` containing chat request details.
/// - Returns: An `AnyPublisher<OKChatResponse, Error>` emitting chat responses from the Ollama API.
public func chat(data: OKChatRequestData) -> AnyPublisher<OKChatResponse, Error> {
do {
let request = try OKRouter.chat(data: data).asURLRequest(with: baseURL)
Expand Down
6 changes: 3 additions & 3 deletions Sources/OllamaKit/RequestData/OKChatRequestData.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ public struct OKChatRequestData: Sendable {
public let messages: [Message]

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

/// Optional ``OKJSONValue`` representing the JSON schema for the response.
/// Be sure to also include "return as JSON" in your prompt
Expand All @@ -31,7 +31,7 @@ public struct OKChatRequestData: Sendable {
public init(
model: String,
messages: [Message],
tools: [OKJSONValue]? = nil,
tools: [OKTool]? = nil,
format: OKJSONValue? = nil,
options: OKCompletionOptions? = nil
) {
Expand All @@ -46,7 +46,7 @@ public struct OKChatRequestData: Sendable {
public init(
model: String,
messages: [Message],
tools: [OKJSONValue]? = nil,
tools: [OKTool]? = nil,
format: OKJSONValue? = nil,
with configureOptions: @Sendable (inout OKCompletionOptions) -> Void
) {
Expand Down
46 changes: 46 additions & 0 deletions Sources/OllamaKit/Utils/OKTool.swift
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
//
// OKTool.swift
// OllamaKit
//
// Created by Norikazu Muramoto on 2025/01/11.
//

import Foundation

/// Represents a tool that can be used in the Ollama API chat.
public struct OKTool: Encodable, Sendable {
/// The type of the tool (e.g., "function").
public let type: String

/// The function details associated with the tool.
public let function: OKFunction

public init(type: String, function: OKFunction) {
self.type = type
self.function = function
}

/// Convenience method for creating a tool with type "function".
public static func function(_ function: OKFunction) -> OKTool {
return OKTool(type: "function", function: function)
}
}

/// Represents a function used as a tool in the Ollama API chat.
public struct OKFunction: Encodable, Sendable {
/// The name of the function.
public let name: String

/// A description of what the function does.
public let description: String

/// Parameters required by the function, defined as a JSON schema.
public let parameters: OKJSONValue

public init(name: String, description: String, parameters: OKJSONValue) {
self.name = name
self.description = description
self.parameters = parameters
}
}

0 comments on commit b49015d

Please sign in to comment.