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

llama_stack/providers/inline/agents/meta_reference/responses/openai_responses.py

@@ -324,6 +324,125 @@ class OpenAIResponsesImpl:
             messages=messages,
         )
 
+    def _prepare_input_items_for_storage(
+        self,
+        input: str | list[OpenAIResponseInput],
+    ) -> list[OpenAIResponseInput]:
+        """Prepare input items for storage, adding IDs where needed.
+
+        This method is called once at the start of streaming to prepare input items
+        that will be reused across multiple persistence calls during streaming.
+        """
+        new_input_id = f"msg_{uuid.uuid4()}"
+        input_items_data: list[OpenAIResponseInput] = []
+
+        if isinstance(input, str):
+            input_content = OpenAIResponseInputMessageContentText(text=input)
+            input_content_item = OpenAIResponseMessage(
+                role="user",
+                content=[input_content],
+                id=new_input_id,
+            )
+            input_items_data = [input_content_item]
+        else:
+            for input_item in input:
+                if isinstance(input_item, OpenAIResponseMessage):
+                    input_item_dict = input_item.model_dump()
+                    if "id" not in input_item_dict:
+                        input_item_dict["id"] = new_input_id
+                    input_items_data.append(OpenAIResponseMessage(**input_item_dict))
+                else:
+                    input_items_data.append(input_item)
+
+        return input_items_data
+
+    async def _persist_streaming_state(
+        self,
+        stream_chunk: OpenAIResponseObjectStream,
+        orchestrator,
+        input_items: list[OpenAIResponseInput],
+        output_items: list,
+    ) -> None:
+        """Persist response state at significant streaming events.
+
+        This enables clients to poll GET /v1/responses/{response_id} during streaming
+        to see in-progress turn state instead of empty results.
+
+        Persistence occurs at:
+        - response.in_progress: Initial INSERT with empty output
+        - response.output_item.done: UPDATE with accumulated output items
+        - response.completed/response.incomplete: Final UPDATE with complete state
+        - response.failed: UPDATE with error state
+
+        :param stream_chunk: The current streaming event.
+        :param orchestrator: The streaming orchestrator (for snapshotting response).
+        :param input_items: Pre-prepared input items for storage.
+        :param output_items: Accumulated output items so far.
+        """
+        try:
+            match stream_chunk.type:
+                case "response.in_progress":
+                    # Initial persistence when response starts
+                    in_progress_response = stream_chunk.response
+                    await self.responses_store.upsert_response_object(
+                        response_object=in_progress_response,
+                        input=input_items,
+                        messages=[],
+                    )
+
+                case "response.output_item.done":
+                    # Incremental update when an output item completes (tool call, message)
+                    current_snapshot = orchestrator._snapshot_response(
+                        status="in_progress",
+                        outputs=output_items,
+                    )
+                    # Get current messages (filter out system messages)
+                    messages_to_store = list(
+                        filter(
+                            lambda x: not isinstance(x, OpenAISystemMessageParam),
+                            orchestrator.final_messages or orchestrator.ctx.messages,
+                        )
+                    )
+                    await self.responses_store.upsert_response_object(
+                        response_object=current_snapshot,
+                        input=input_items,
+                        messages=messages_to_store,
+                    )
+
+                case "response.completed" | "response.incomplete":
+                    # Final persistence when response finishes
+                    final_response = stream_chunk.response
+                    messages_to_store = list(
+                        filter(
+                            lambda x: not isinstance(x, OpenAISystemMessageParam),
+                            orchestrator.final_messages,
+                        )
+                    )
+                    await self.responses_store.upsert_response_object(
+                        response_object=final_response,
+                        input=input_items,
+                        messages=messages_to_store,
+                    )
+
+                case "response.failed":
+                    # Persist failed state so GET shows error
+                    failed_response = stream_chunk.response
+                    # Preserve any accumulated non-system messages for failed responses
+                    messages_to_store = list(
+                        filter(
+                            lambda x: not isinstance(x, OpenAISystemMessageParam),
+                            orchestrator.final_messages or orchestrator.ctx.messages,
+                        )
+                    )
+                    await self.responses_store.upsert_response_object(
+                        response_object=failed_response,
+                        input=input_items,
+                        messages=messages_to_store,
+                    )
+        except Exception as e:
+            # Best-effort persistence: log error but don't fail the stream
+            logger.warning(f"Failed to persist streaming state for {stream_chunk.type}: {e}")
+
     async def create_openai_response(
         self,
         input: str | list[OpenAIResponseInput],
@@ -528,6 +647,10 @@ class OpenAIResponsesImpl:
 
         # Type as ConversationItem to avoid list invariance issues
         output_items: list[ConversationItem] = []
+
+        # Prepare input items for storage once (used by all persistence calls)
+        input_items_for_storage = self._prepare_input_items_for_storage(all_input)
+
         try:
             async for stream_chunk in orchestrator.create_response():
                 match stream_chunk.type:
@@ -541,6 +664,16 @@ class OpenAIResponsesImpl:
                     case _:
                         pass  # Other event types
 
+                # Incremental persistence: persist on significant state changes
+                # This enables clients to poll GET /v1/responses/{response_id} during streaming
+                if store:
+                    await self._persist_streaming_state(
+                        stream_chunk=stream_chunk,
+                        orchestrator=orchestrator,
+                        input_items=input_items_for_storage,
+                        output_items=output_items,
+                    )
+
                 # Store and sync before yielding terminal events
                 # This ensures the storage/syncing happens even if the consumer breaks after receiving the event
                 if (
@@ -548,18 +681,10 @@ class OpenAIResponsesImpl:
                     and final_response
                     and failed_response is None
                 ):
-                    messages_to_store = list(
-                        filter(lambda x: not isinstance(x, OpenAISystemMessageParam), orchestrator.final_messages)
-                    )
-                    if store:
-                        # TODO: we really should work off of output_items instead of "final_messages"
-                        await self._store_response(
-                            response=final_response,
-                            input=all_input,
-                            messages=messages_to_store,
-                        )
-
                     if conversation:
+                        messages_to_store = list(
+                            filter(lambda x: not isinstance(x, OpenAISystemMessageParam), orchestrator.final_messages)
+                        )
                         await self._sync_response_to_conversation(conversation, input, output_items)
                         await self.responses_store.store_conversation_messages(conversation, messages_to_store)
 

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

@@ -0,0 +1,9 @@
+#import <Foundation/Foundation.h>
+
+//! Project version number for LocalInference.
+FOUNDATION_EXPORT double LocalInferenceVersionNumber;
+
+//! Project version string for LocalInference.
+FOUNDATION_EXPORT const unsigned char LocalInferenceVersionString[];
+
+// In this header, you should import all the public headers of your framework using statements like #import <LocalInference/PublicHeader.h>

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)
+  }
+}