llama-stack 0.4.2__py3-none-any.whl → 0.4.4__py3-none-any.whl

llama_stack/providers/inline/ios/inference/LocalInferenceImpl/LocalInference.swift

@@ -0,0 +1,189 @@
+import Foundation
+
+import LLaMARunner
+import LlamaStackClient
+
+class RunnerHolder: ObservableObject {
+  var runner: Runner?
+}
+
+public class LocalInference: Inference {
+  private var runnerHolder = RunnerHolder()
+  private let runnerQueue: DispatchQueue
+
+  public init (queue: DispatchQueue) {
+    runnerQueue = queue
+  }
+
+  public func loadModel(modelPath: String, tokenizerPath: String, completion: @escaping (Result<Void, Error>) -> Void) {
+    runnerHolder.runner = runnerHolder.runner ?? Runner(
+      modelPath: modelPath,
+      tokenizerPath: tokenizerPath
+    )
+
+
+    runnerQueue.async {
+      let runner = self.runnerHolder.runner
+      do {
+        try runner!.load()
+        completion(.success(()))
+      } catch let loadError {
+        print("error: " + loadError.localizedDescription)
+        completion(.failure(loadError))
+      }
+    }
+  }
+
+  public func stop() {
+    runnerHolder.runner?.stop()
+  }
+
+  public func chatCompletion(request: Components.Schemas.ChatCompletionRequest) -> AsyncStream<Components.Schemas.ChatCompletionResponseStreamChunk> {
+    return AsyncStream { continuation in
+      let workItem = DispatchWorkItem {
+        do {
+          var tokens: [String] = []
+
+          let prompt = try encodeDialogPrompt(messages: prepareMessages(request: request))
+          var stopReason: Components.Schemas.CompletionMessage.stop_reasonPayload? = nil
+          var buffer = ""
+          var ipython = false
+          var echoDropped = false
+
+          try self.runnerHolder.runner?.generate(prompt, sequenceLength: 4096) { token in
+            buffer += token
+
+            // HACK: Workaround until LlamaRunner exposes echo param
+            if (!echoDropped) {
+              if (buffer.hasPrefix(prompt)) {
+                buffer = String(buffer.dropFirst(prompt.count))
+                echoDropped = true
+              }
+              return
+            }
+
+            tokens.append(token)
+
+            if !ipython && (buffer.starts(with: "<|python_tag|>") || buffer.starts(with: "[") ) {
+              ipython = true
+              continuation.yield(
+                Components.Schemas.ChatCompletionResponseStreamChunk(
+                  event: Components.Schemas.ChatCompletionResponseEvent(
+                    event_type: .progress,
+                    delta: .tool_call(Components.Schemas.ToolCallDelta(
+                      _type: Components.Schemas.ToolCallDelta._typePayload.tool_call,
+                      tool_call: .case1(""),
+                      parse_status: Components.Schemas.ToolCallDelta.parse_statusPayload.started
+                      )
+                    )
+                  )
+                )
+              )
+
+              if (buffer.starts(with: "<|python_tag|>")) {
+                buffer = String(buffer.dropFirst("<|python_tag|>".count))
+              }
+            }
+
+            // TODO: Non-streaming lobprobs
+
+            var text = ""
+            if token == "<|eot_id|>" {
+              stopReason = Components.Schemas.CompletionMessage.stop_reasonPayload.end_of_turn
+            } else if token == "<|eom_id|>" {
+              stopReason = Components.Schemas.CompletionMessage.stop_reasonPayload.end_of_message
+            } else {
+              text = token
+            }
+
+            var delta: Components.Schemas.ContentDelta
+            if ipython {
+              delta = .tool_call(Components.Schemas.ToolCallDelta(
+                _type: .tool_call,
+                tool_call: .case1(text),
+                parse_status: .in_progress
+              ))
+            } else {
+              delta = .text(Components.Schemas.TextDelta(
+                _type: Components.Schemas.TextDelta._typePayload.text,
+                text: text
+                )
+              )
+            }
+
+            if stopReason == nil {
+              continuation.yield(
+                Components.Schemas.ChatCompletionResponseStreamChunk(
+                  event: Components.Schemas.ChatCompletionResponseEvent(
+                    event_type: .progress,
+                    delta: delta
+                  )
+                )
+              )
+            }
+          }
+
+          if stopReason == nil {
+            stopReason = Components.Schemas.CompletionMessage.stop_reasonPayload.out_of_tokens
+          }
+
+          let message = decodeAssistantMessage(tokens: tokens.joined(), stopReason: stopReason!)
+          // TODO: non-streaming support
+
+          let didParseToolCalls = message.tool_calls?.count ?? 0 > 0
+          if ipython && !didParseToolCalls {
+            continuation.yield(
+              Components.Schemas.ChatCompletionResponseStreamChunk(
+                event: Components.Schemas.ChatCompletionResponseEvent(
+                  event_type: .progress,
+                  delta: .tool_call(Components.Schemas.ToolCallDelta(
+                    _type: Components.Schemas.ToolCallDelta._typePayload.tool_call,
+                    tool_call: .case1(""),
+                    parse_status: Components.Schemas.ToolCallDelta.parse_statusPayload.failed
+                    )
+                  )
+                )
+                // TODO: stopReason
+              )
+            )
+          }
+
+          for toolCall in message.tool_calls! {
+            continuation.yield(
+              Components.Schemas.ChatCompletionResponseStreamChunk(
+                event: Components.Schemas.ChatCompletionResponseEvent(
+                  event_type: .progress,
+                  delta: .tool_call(Components.Schemas.ToolCallDelta(
+                    _type: Components.Schemas.ToolCallDelta._typePayload.tool_call,
+                    tool_call: Components.Schemas.ToolCallDelta.tool_callPayload.ToolCall(toolCall),
+                    parse_status: Components.Schemas.ToolCallDelta.parse_statusPayload.succeeded
+                    )
+                  )
+                )
+                // TODO: stopReason
+              )
+            )
+          }
+
+          continuation.yield(
+            Components.Schemas.ChatCompletionResponseStreamChunk(
+              event: Components.Schemas.ChatCompletionResponseEvent(
+                event_type: .complete,
+                delta: .text(Components.Schemas.TextDelta(
+                  _type: Components.Schemas.TextDelta._typePayload.text,
+                  text: ""
+                  )
+                )
+              )
+              // TODO: stopReason
+            )
+          )
+        }
+        catch (let error) {
+          print("Inference error: " + error.localizedDescription)
+        }
+      }
+      runnerQueue.async(execute: workItem)
+    }
+  }
+}

llama_stack/providers/inline/ios/inference/LocalInferenceImpl/Parsing.swift

@@ -0,0 +1,238 @@
+import Foundation
+
+import LlamaStackClient
+
+func encodeHeader(role: String) -> String {
+  return "<|start_header_id|>\(role)<|end_header_id|>\n\n"
+}
+
+func encodeDialogPrompt(messages: [Components.Schemas.Message]) -> String {
+  var prompt = ""
+
+  prompt.append("<|begin_of_text|>")
+  for message in messages {
+    let msg = encodeMessage(message: message)
+    prompt += msg
+  }
+
+  prompt.append(encodeHeader(role: "assistant"))
+
+  return prompt
+}
+
+func getRole(message: Components.Schemas.Message) -> String {
+  switch (message) {
+  case .user(let m):
+    return m.role.rawValue
+  case .system(let m):
+    return m.role.rawValue
+  case .tool(let m):
+    return m.role.rawValue
+  case .assistant(let m):
+    return m.role.rawValue
+  }
+}
+
+func encodeMessage(message: Components.Schemas.Message) -> String {
+  var prompt = encodeHeader(role: getRole(message: message))
+
+  switch (message) {
+  case .assistant(let m):
+    if (m.tool_calls?.count ?? 0 > 0) {
+      prompt += "<|python_tag|>"
+    }
+  default:0
+    break
+  }
+
+  func _processContent(_ content: Any) -> String {
+    func _process(_ c: Any) {
+      if let str = c as? String {
+        prompt += str
+      }
+    }
+
+    if let str = content as? String {
+      _process(str)
+    } else if let list = content as? [Any] {
+      for c in list {
+        _process(c)
+      }
+    }
+
+    return ""
+  }
+
+  switch (message) {
+  case .user(let m):
+    prompt += _processContent(m.content)
+  case .system(let m):
+    prompt += _processContent(m.content)
+  case .tool(let m):
+    prompt += _processContent(m.content)
+  case .assistant(let m):
+    prompt += _processContent(m.content)
+  }
+
+  var eom = false
+
+  switch (message) {
+  case .user(let m):
+    switch (m.content) {
+    case .case1(let c):
+      prompt += _processContent(c)
+    case .InterleavedContentItem(let c):
+      prompt += _processContent(c)
+    case .case3(let c):
+      prompt += _processContent(c)
+    }
+  case .assistant(let m):
+    // TODO: Support encoding past tool call history
+    // for t in m.tool_calls {
+    //  _processContent(t.)
+    //}
+    eom = m.stop_reason == Components.Schemas.CompletionMessage.stop_reasonPayload.end_of_message
+  case .system(_):
+    break
+  case .tool(_):
+    break
+  }
+
+  if (eom) {
+    prompt += "<|eom_id|>"
+  } else {
+    prompt += "<|eot_id|>"
+  }
+
+  return prompt
+}
+
+func prepareMessages(request: Components.Schemas.ChatCompletionRequest) throws -> [Components.Schemas.Message] {
+  var existingMessages = request.messages
+  var existingSystemMessage: Components.Schemas.Message?
+  // TODO: Existing system message
+
+  var messages: [Components.Schemas.Message] = []
+
+  let defaultGen = SystemDefaultGenerator()
+  let defaultTemplate = defaultGen.gen()
+
+  var sysContent = ""
+
+  // TODO: Built-in tools
+
+  sysContent += try defaultTemplate.render()
+
+  messages.append(.system(Components.Schemas.SystemMessage(
+    role: .system,
+    content: .case1(sysContent)
+    ))
+  )
+
+  if request.tools?.isEmpty == false {
+    // TODO: Separate built-ins and custom tools (right now everything treated as custom)
+    let toolGen = FunctionTagCustomToolGenerator()
+    let toolTemplate = try toolGen.gen(customTools: request.tools!)
+    let tools = try toolTemplate.render()
+    messages.append(.user(Components.Schemas.UserMessage(
+      role: .user,
+      content: .case1(tools))
+    ))
+  }
+
+  messages.append(contentsOf: existingMessages)
+
+  return messages
+}
+
+struct FunctionCall {
+    let name: String
+    let params: [String: Any]
+}
+
+public func maybeExtractCustomToolCalls(input: String) -> [Components.Schemas.ToolCall] {
+  guard input.hasPrefix("[") && input.hasSuffix("]") else {
+    return []
+  }
+
+  do {
+    let trimmed = input.trimmingCharacters(in: CharacterSet(charactersIn: "[]"))
+    let calls = trimmed.components(separatedBy: "),").map { $0.hasSuffix(")") ? $0 : $0 + ")" }
+
+    var result: [Components.Schemas.ToolCall] = []
+
+    for call in calls {
+      guard let nameEndIndex = call.firstIndex(of: "("),
+            let paramsStartIndex = call.firstIndex(of: "{"),
+            let paramsEndIndex = call.lastIndex(of: "}") else {
+        return []
+      }
+
+      let name = String(call[..<nameEndIndex]).trimmingCharacters(in: .whitespacesAndNewlines)
+      let paramsString = String(call[paramsStartIndex...paramsEndIndex])
+
+      guard let data = paramsString.data(using: .utf8),
+            let params = try? JSONSerialization.jsonObject(with: data, options: []) as? [String: Any] else {
+        return []
+      }
+
+      var props: [String : Components.Schemas.ToolCall.argumentsPayload.additionalPropertiesPayload] = [:]
+      for (param_name, param) in params {
+        switch (param) {
+        case let value as String:
+          props[param_name] = .case1(value)
+        case let value as Int:
+          props[param_name] = .case2(value)
+        case let value as Double:
+          props[param_name] = .case3(value)
+        case let value as Bool:
+          props[param_name] = .case4(value)
+        default:
+          return []
+        }
+      }
+
+      result.append(
+        Components.Schemas.ToolCall(
+          call_id: UUID().uuidString,
+          tool_name: .case2(name), // custom_tool
+          arguments: .init(additionalProperties: props)
+        )
+      )
+    }
+
+    return result.isEmpty ? [] : result
+  } catch {
+    return []
+  }
+}
+
+func decodeAssistantMessage(tokens: String, stopReason: Components.Schemas.CompletionMessage.stop_reasonPayload) -> Components.Schemas.CompletionMessage {
+  var content = tokens
+
+  let roles = ["user", "system", "assistant"]
+  for role in roles {
+    let headerStr = encodeHeader(role: role)
+    if content.hasPrefix(headerStr) {
+      content = String(content.dropFirst(encodeHeader(role: role).count))
+    }
+  }
+
+  if content.hasPrefix("<|python_tag|>") {
+    content = String(content.dropFirst("<|python_tag|>".count))
+  }
+
+
+  if content.hasSuffix("<|eot_id|>") {
+    content = String(content.dropLast("<|eot_id|>".count))
+  } else {
+    content = String(content.dropLast("<|eom_id|>".count))
+  }
+
+  return Components.Schemas.CompletionMessage(
+    role: .assistant,
+    content: .case1(content),
+    stop_reason: stopReason,
+    tool_calls: maybeExtractCustomToolCalls(input: content)
+  )
+}

llama_stack/providers/inline/ios/inference/LocalInferenceImpl/PromptTemplate.swift

@@ -0,0 +1,12 @@
+import Foundation
+import Stencil
+
+public struct PromptTemplate {
+    let template: String
+    let data: [String: Any]
+
+  public func render() throws -> String {
+    let template = Template(templateString: self.template)
+    return try template.render(self.data)
+  }
+}

llama_stack/providers/inline/ios/inference/LocalInferenceImpl/SystemPrompts.swift

@@ -0,0 +1,89 @@
+import Foundation
+
+import LlamaStackClient
+
+func convertToNativeSwiftType(_ value: Any) -> Any {
+    switch value {
+    case let number as NSNumber:
+        if CFGetTypeID(number) == CFBooleanGetTypeID() {
+            return number.boolValue
+        }
+        if floor(number.doubleValue) == number.doubleValue {
+            return number.intValue
+        }
+        return number.doubleValue
+    case let string as String:
+        return string
+    case let array as [Any]:
+        return array.map(convertToNativeSwiftType)
+    case let dict as [String: Any]:
+        return dict.mapValues(convertToNativeSwiftType)
+    case is NSNull:
+        return NSNull()
+    default:
+        return value
+    }
+}
+
+public class SystemDefaultGenerator {
+  public init() {}
+
+  public func gen() -> PromptTemplate {
+    let templateStr = """
+            Cutting Knowledge Date: December 2023
+            Today Date: {{ today }}
+            """
+
+    let dateFormatter = DateFormatter()
+    dateFormatter.dateFormat = "dd MMMM yyyy"
+
+    return PromptTemplate(
+      template: templateStr,
+      data: ["today": dateFormatter.string(from: Date())]
+    )
+  }
+}
+
+
+public class FunctionTagCustomToolGenerator {
+  public init() {}
+
+  public func gen(customTools: [Components.Schemas.ToolDefinition]) throws -> PromptTemplate {
+    // TODO: required params
+    // TODO: {{#unless @last}},{{/unless}}
+
+    let templateStr = """
+            You are an expert in composing functions. You are given a question and a set of possible functions.
+            Based on the question, you will need to make one or more function/tool calls to achieve the purpose.
+            If none of the function can be used, point it out. If the given question lacks the parameters required by the function,
+            also point it out. You should only return the function call in tools call sections.
+
+            If you decide to invoke any of the function(s), you MUST put it in the format of [func_name1(params_name1=params_value1, params_name2=params_value2...), func_name2(params)]
+            You SHOULD NOT include any other text in the response.
+
+            Here is a list of functions in JSON format that you can invoke.
+
+            [
+            {% for t in custom_tools %}
+            {
+                "name": "{{t.tool_name}}",
+                "description": "{{t.description}}",
+                "input_schema": { {{t.input_schema}} }
+            }
+
+            {{/let}}
+            {% endfor -%}
+            ]
+            """
+
+    let encoder = JSONEncoder()
+    return PromptTemplate(
+      template: templateStr,
+      data: ["custom_tools": try customTools.map {
+        let data = try encoder.encode($0)
+        let obj = try JSONSerialization.jsonObject(with: data)
+        return convertToNativeSwiftType(obj)
+      }]
+    )
+  }
+}