update

2026-08-13 18:14:46 +07:00
parent 510cf6126c
commit b8067c2d33
3 changed files with 582 additions and 174 deletions
@@ -392,30 +392,29 @@ function _processMessage(

    agentEventSink("_processMessage 11")
    # Extract tool calls from LLM response content blocks
-    hasToolCalls, toolCallList = _extractToolCalls(response)
+    hasToolCalls, toolCallList, assistant_msg = _extractToolCalls(response)
    agentEventSink("hasToolCalls: $hasToolCalls\ntoolCallList: $toolCallList")
-    error("debug marker")
+
    if hasToolCalls && length(toolCallList) > 0
-      # Build context and config for executeToolCalls
-      context = agentContext(
-        systemPrompt,
-        messages,
-        tools,
-      )
+      #WORKING Build context and config for executeToolCalls

      config = agentLoopConfig(
-        tools,
        beforeToolCall,
        afterToolCall,
        parallelToolExecute ? "parallel" : "sequential",
      )

-      signal = nothing
-      emit = agentEventSink
+      signal = abortSignal(false)
+      agentEventSink("_processMessage 12")

      # call executeToolCalls()
-      batch = executeToolCalls(context, response, toolCallList, config, signal, emit)

+        batch = executeToolCalls(preparedContext, assistant_msg, toolCallList, config, signal, 
+                               agentEventSink)
+
+      
+      agentEventSink("_processMessage 13")
+      error("debug marker")
      # save toolResults to messages
      for tool_result in batch.messages
        push!(messages, tool_result)
@@ -438,9 +437,9 @@ function _processMessage(
        final_response = assistantMessage(
          role="assistant",
          content=final_content,
-          api=response.api,
-          model=response.model,
-          usage=response.usage,
+          api=assistant_msg.api,
+          model=assistant_msg.model,
+          usage=assistant_msg.usage,
          stopReason="tool_use_terminated",
          errorMessage=if any(x -> x.isError, batch.messages)
            "One or more tool calls failed"
@@ -453,7 +452,7 @@ function _processMessage(
      end
    else
      # LLM did not use tool calls — this is the final response
-      final_response = response
+      final_response = assistant_msg
      break
    end
  end
@@ -530,9 +529,10 @@ function createToolResultMessage(f::finalizedOutcome)::toolResultMessage
 end

 """
-    _extractToolCalls(response) -> Tuple{Bool, Vector{agentToolCall}}
+    _extractToolCalls(response) -> Tuple{Bool, Vector{agentToolCall}, assistantMessage}

-Extracts tool calls from the LLM response. Supports two response formats:
+Extracts tool calls from the LLM response and constructs an `assistantMessage`.
+Supports two response formats:

 1. **Message format** (e.g. from LMStudio.jl / vLLM):
   `response["message"]["tool_calls"]` — array of tool call objects with
@@ -544,14 +544,19 @@ Extracts tool calls from the LLM response. Supports two response formats:
   contain a `"tool_calls"` array in format 1. Blocks with `"type" => "tool_call"`
   have `"name"`, `"arguments"`, `"id"` at the top level (already-parsed args dict).

-Returns `(hasToolCalls, toolCallList)` where `hasToolCalls` is `true` if any
-tool calls were found, and `toolCallList` is a vector of `agentToolCall` structs.
+The `assistantMessage` is constructed from:
+- `reasoning_content` (string) → added as a `textContent` block
+- `response.content` blocks (text/ reasoning) → added to content
+- Top-level `api`, `provider`, `model`, `usage` → copied to the message
+- `finish_reason` → used as `stopReason`

 # Arguments
 - `response`: LLM response object (Dict/JSON.Object or struct with `.content` field)

 # Returns
- `Tuple{Bool, Vector{agentToolCall}}`: `(hasToolCalls, toolCallList)`
+- `Tuple{Bool, Vector{agentToolCall}, assistantMessage}`: `(hasToolCalls, toolCallList, assistantMsg)`
+
+# Example response = JSON.Object{String, Any}("finish_reason" => "tool_calls", "index" => 0, "message" => JSON.Object{String, Any}("role" => "assistant", "content" => "", "reasoning_content" => "Here's a thinking process:\n\n1.  **Identify User Request**: The user is asking for the weather in Bangkok.\n2.  **Locate Relevant Tool**: I have a `getWeather` function available.\n3.  **Check Function Parameters**:\n   - `city` (required): City and country, e.g., 'San Francisco, CA' or 'Tokyo, Japan'\n   - `units` (optional, default \"celsius\"): Temperature scale (\"celsius\" or \"fahrenheit\")\n4.  **Prepare Parameters**:\n   - `city`: \"Bangkok, Thailand\" (adding country for clarity, though just \"Bangkok\" might work, following the example format is safer)\n   - `units`: Not specified, so I'll use the default (\"celsius\")\n5.  **Execute Tool Call**: Call `getWeather` with `city: \"Bangkok, Thailand\"`\n6.  **Anticipate Response**: The function will return current weather and forecast data for Bangkok. I'll then format it nicely for the user.\n   - *Self-Correction/Verification during thought*: The prompt says \"city: City and country, e.g., 'San Francisco, CA' or 'Tokyo, Japan'\". I'll use \"Bangkok, Thailand\". The `units` parameter is optional, so I'll omit it to use the default.\n   - Proceed with tool call.✅\n", "tool_calls" => Any[JSON.Object{String, Any}("type" => "function", "function" => JSON.Object{String, Any}("name" => "getWeather", "arguments" => "{\"city\":\"Bangkok, Thailand\"}"), "id" => "6fOilR5QPcdppbAAHluhkRyUDu3oWMAL")]))
 """
 function _extractToolCalls(response)
    hasToolCalls = false
@@ -600,9 +605,9 @@ function _extractToolCalls(response)
            for content_block in content
                if content_block isa AbstractDict
                    if get(content_block, "type", "") == "tool_calls"
-                        hasToolCalls = true
                        for tc_data in get(content_block, "tool_calls", [])
                            if tc_data isa AbstractDict
+                                hasToolCalls = true
                                push!(toolCallList, make_tc(tc_data))
                            end
                        end
@@ -621,7 +626,88 @@ function _extractToolCalls(response)
        end
    end

-    return hasToolCalls, toolCallList
+    # ── Construct assistantMessage from response ──────────────────────
+    # Check Format 1 nested message for reasoning_content, role, etc.
+    reasoning = get(response, "reasoning_content", nothing)
+    if reasoning === nothing && msg !== nothing && msg isa AbstractDict
+        reasoning = get(msg, "reasoning_content", nothing)
+    end
+    if reasoning isa String
+        reasoning_text = reasoning
+    elseif reasoning isa textContent
+        reasoning_text = reasoning.text
+    else
+        reasoning_text = ""
+    end
+    reasoning_block = !isempty(reasoning_text) ? [textContent(reasoning_text)] : textContent[]
+
+    finish_reason = get(response, "finish_reason", nothing)
+    stop_reason = finish_reason isa AbstractString ? String(finish_reason) : "end_turn"
+
+    api = get(response, "api", "")
+    provider = get(response, "provider", "")
+    model = get(response, "model", nothing)
+    usage = get(response, "usage", nothing)
+
+    content_from_response = get(response, "content", nothing)
+    content_blocks = Vector{messageContent}()
+    if content_from_response isa Vector
+        for block in content_from_response
+            if block isa AbstractDict
+                if get(block, "type", "") == "text"
+                    push!(content_blocks, textContent(get(block, "text", "")))
+                elseif get(block, "type", "") == "tool_call"
+                    # tool_call blocks — don't add text content for these
+                elseif get(block, "type", "") == "tool_calls"
+                    # tool_calls blocks — don't add text content for these
+                elseif get(block, "type", "") == "reasoning"
+                    push!(content_blocks, textContent(get(block, "text", "")))
+                else
+                    push!(content_blocks, textContent(get(block, "text", "")))
+                end
+            end
+        end
+    end
+
+    # Combine content blocks and reasoning
+    if !isempty(content_blocks) && !isempty(reasoning_block)
+        all_content = vcat(reasoning_block, content_blocks)
+    elseif !isempty(content_blocks)
+        all_content = content_blocks
+    elseif !isempty(reasoning_block)
+        all_content = reasoning_block
+    else
+        all_content = textContent[]
+    end
+
+    error_msg = get(response, "error_message", get(response, "errorMessage", nothing))
+
+    if usage === nothing || !(usage isa llmUsage)
+        usage = llmUsage(0, 0)
+    end
+
+    # Role: prefer Format 1 nested message, fallback to top-level, default "assistant"
+    role = get(response, "role", "assistant")
+    if role == "assistant" && msg !== nothing && msg isa AbstractDict
+        nested_role = get(msg, "role", nothing)
+        if nested_role isa AbstractString
+            role = nested_role
+        end
+    end
+
+    assistant_msg = assistantMessage(
+        role = role,
+        content = all_content,
+        api = api isa AbstractString ? String(api) : "",
+        provider = provider isa AbstractString ? String(provider) : "",
+        model = model,
+        usage = usage,
+        stopReason = stop_reason,
+        errorMessage = error_msg,
+        timestamp = now(),
+    )
+
+    return hasToolCalls, toolCallList, assistant_msg
 end

 """
@@ -759,37 +845,55 @@ function prepareToolCall(
  assistantMsg::assistantMessage,
  toolCall::agentToolCall,
  config::agentLoopConfig,
-  signal::Union{Nothing, abortSignal},
+  signal::abortSignal,
+  agentEventSink
 )::Union{preparedToolCall,immediateOutcome}
-
+  agentEventSink("prepareToolCall 1")
  tool = get(context.tools, toolCall.name, nothing)
  if tool === nothing
+    agentEventSink("prepareToolCall 2")
    return immediateOutcome(createErrorToolResult("Tool $toolCall.name not found"), true)
  end

  try
+    agentEventSink("prepareToolCall 3")
    # 1. prepare arguments (tool-specific transform)
    prepared = prepareToolCallArguments(tool, toolCall)
+    agentEventSink("prepareToolCall 4")
    validatedArgs = validateToolArguments(tool, prepared)
-
+    agentEventSink("prepareToolCall 5")
    # 2. beforeToolCall hook — can block
    if config.beforeToolCall !== nothing
+      agentEventSink("prepareToolCall 6")
+
      before = config.beforeToolCall(
        beforeToolCallContext(assistantMsg, toolCall, validatedArgs, context), 
        signal
        )
-      if signal !== nothing && signal.aborted
+      agentEventSink("prepareToolCall 7")
+      if signal.aborted
+        agentEventSink("prepareToolCall 8")
        return immediateOutcome(createErrorToolResult("Operation aborted"), true)
      end
+
      if before !== nothing && before.block
+        agentEventSink("prepareToolCall 9")
        return immediateOutcome(
          createErrorToolResult(get(before, :reason, "Tool execution was blocked")), true)
      end
    end
-
+    agentEventSink("prepareToolCall 10")
    return preparedToolCall(tool, toolCall, validatedArgs)
-  catch err
-    return immediateOutcome(createErrorToolResult(sprint(showerror, err)), true)
+  catch e
+    bt = catch_backtrace()
+    err_msg = sprint() do io
+        showerror(io, e, bt)
+        println(io)
+    end
+
+    agentEventSink(err_msg)
+
+    return immediateOutcome(createErrorToolResult(sprint(showerror, e)), true)
  end
 end

@@ -836,9 +940,9 @@ executePreparedToolCall(prep, nothing, emit)
 function executePreparedToolCall(
    prep::preparedToolCall,
    signal::Union{Nothing,abortSignal},
-    emit::Function,
-)::executedOutcome
-
+    agentEventSink,
+  )::executedOutcome
+  agentEventSink("executePreparedToolCall 1")
  updateEvents = promise[]
  accepting = true

@@ -848,7 +952,7 @@ function executePreparedToolCall(
      partialResult -> begin
        if accepting
          push!(updateEvents,
-                          emit(toolExecUpdateEvent(prep.toolCall.id, prep.toolCall.name,
+                agentEventSink(toolExecUpdateEvent(prep.toolCall.id, prep.toolCall.name,
                                          prep.toolCall.arguments, partialResult)))
        end
      end
@@ -1011,35 +1115,40 @@ function executeToolCallsSequential(
    assistantMsg::assistantMessage,
    toolCalls::Vector{agentToolCall},
    config::agentLoopConfig,
-    signal::Union{Nothing,abortSignal},
-    emit::Function,
+    signal::abortSignal,
+    agentEventSink,
 )::agentToolCallBatch
-
+  agentEventSink("executeToolCallsSequential 1")
  finalizedCalls = finalizedOutcome[]
  messages = toolResultMessage[]

  for tc in toolCalls
-        emit(toolExecStartEvent(tc.id, tc.name, tc.arguments))
-
-        prep = prepareToolCall(context, assistantMsg, tc, config, signal)
-
+      agentEventSink("start tool execute $(tc.id), $(tc.name), $(tc.arguments)")
+      prep = prepareToolCall(context, assistantMsg, tc, config, signal, agentEventSink)
+      agentEventSink("executeToolCallsSequential 2")
      if prep isa immediateOutcome
+        agentEventSink("executeToolCallsSequential 2-1")
        finalized = finalizedOutcome(tc, prep.result, prep.isError)
+        agentEventSink("executeToolCallsSequential 2-2")
      else
-            executed = executePreparedToolCall(prep, signal, emit)
-            finalized = finalizeExecutedToolCall(context, assistantMsg, prep, executed, config, signal)
+        agentEventSink("executeToolCallsSequential 3")
+        executed = executePreparedToolCall(prep, signal, agentEventSink)
+        agentEventSink("executeToolCallsSequential 3-1")
+        finalized = finalizeExecutedToolCall(context, assistantMsg, prep, executed, config, 
+                                            signal)
+        agentEventSink("executeToolCallsSequential 3-2")
      end
-
-        emit(toolExecEndEvent(finalized.toolCall.id, finalized.toolCall.name,
-                              finalized.result, finalized.isError))
+      agentEventSink("executeToolCallsSequential 4")
+      agentEventSink("$(finalized.toolCall.id), $(finalized.toolCall.name),
+                            $(finalized.result), $(finalized.isError)")
      push!(messages, createToolResultMessage(finalized))
      push!(finalizedCalls, finalized)
-
+      agentEventSink("executeToolCallsSequential 5")
      if signal !== nothing && signal.aborted
          break
      end
  end
-
+  agentEventSink("executeToolCallsSequential 6")
  return agentToolCallBatch(messages, shouldTerminate(finalizedCalls))
 end

@@ -1105,27 +1214,27 @@ function executeToolCallsParallel(
    assistantMsg::assistantMessage,
    toolCalls::Vector{agentToolCall},
    config::agentLoopConfig,
-    signal::Union{Nothing,abortSignal},
-    emit::Function,
+    signal::abortSignal,
+    agentEventSink,
 )::agentToolCallBatch

  entries = union{finalizedOutcome,task{finalizedOutcome}}[]

  for tc in toolCalls
-        emit(toolExecStartEvent(tc.id, tc.name, tc.arguments))
+    agentEventSink(toolExecStartEvent(tc.id, tc.name, tc.arguments))

    prep = prepareToolCall(context, assistantMsg, tc, config, signal)

    if prep isa immediateOutcome
      finalized = finalizedOutcome(tc, prep.result, prep.isError)
-            emit(toolExecEndEvent(finalized.toolCall.id, finalized.toolCall.name,
+      agentEventSink(toolExecEndEvent(finalized.toolCall.id, finalized.toolCall.name,
                            finalized.result, finalized.isError))
      push!(entries, finalized)
    else
      task = task() do
-                executed = executePreparedToolCall(prep, signal, emit)
+        executed = executePreparedToolCall(prep, signal, agentEventSink)
        finalized = finalizeExecutedToolCall(context, assistantMsg, prep, executed, config, signal)
-                emit(toolExecEndEvent(finalized.toolCall.id, finalized.toolCall.name,
+        agentEventSink(toolExecEndEvent(finalized.toolCall.id, finalized.toolCall.name,
                              finalized.result, finalized.isError))
        return finalized
      end
@@ -1207,10 +1316,11 @@ function executeToolCalls(
    assistantMsg::assistantMessage,
    toolCalls::Vector{agentToolCall},
    config::agentLoopConfig,
-    signal::Union{Nothing,abortSignal},
-    emit::Function,
-)::agentToolCallBatch
+    signal::abortSignal,
+    agentEventSink,
+  )::agentToolCallBatch

+  agentEventSink("_executeToolCalls 1")
  hasSequential = false
  for tc in toolCalls
      t = get(context.tools, tc.name, nothing)
@@ -1219,11 +1329,15 @@ function executeToolCalls(
          break
      end
  end
-
+  agentEventSink("_executeToolCalls 2")
  if config.toolExecution == "sequential" || hasSequential
-        return executeToolCallsSequential(context, assistantMsg, toolCalls, config, signal, emit)
+    agentEventSink("_executeToolCalls 3")
+    return executeToolCallsSequential(context, assistantMsg, toolCalls, config, signal, 
+                                      agentEventSink)
  else
-        return executeToolCallsParallel(context, assistantMsg, toolCalls, config, signal, emit)
+    agentEventSink("_executeToolCalls 4")
+    return executeToolCallsParallel(context, assistantMsg, toolCalls, config, signal, 
+                                    agentEventSink)
  end
 end

@@ -146,7 +146,8 @@ assistantMessage("assistant", [textContent("Hello!")], "", "", "gpt-4", ..., "en
 function assistantMessage(; role="assistant", content=Vector{messageContent}(),
    api="", provider="", model=nothing, usage=llmUsage(0, 0), stopReason="end_turn",
    errorMessage=nothing, timestamp=now())
-    return assistantMessage(role, content, api, provider, model, usage, stopReason, errorMessage, timestamp)
+    model_str = model isa AbstractString ? String(model) : ""
+    return assistantMessage(role, content, api, provider, model_str, usage, stopReason, errorMessage, timestamp)
 end

 struct toolResultMessage <: agentMessage    # Result returned from a tool execution
@@ -395,13 +396,11 @@ end
 Configuration for the agent tool execution loop.

 # Arguments
- `tools::OrderedDict{String, agentTool}`: Available tools keyed by name
 - `beforeToolCall::Union{Function, Nothing}`: Callback before tool execution
 - `afterToolCall::Union{Function, Nothing}`: Callback after tool execution
 - `toolExecution::String`: Execution mode — "sequential" or "parallel"
 """
 struct agentLoopConfig
-    tools::OrderedDict{String, agentTool}
    beforeToolCall::Union{Function, Nothing}
    afterToolCall::Union{Function, Nothing}
    toolExecution::String
@@ -9,9 +9,9 @@ import YiemAgent.agentCore: _extractToolCalls

@testset "_extractToolCalls" begin

-    # -------------------------------------------------------------- #
+    # ------------------------------------------------------------------ #
    # Format 1: response["message"]["tool_calls"] (LMStudio.jl style)     #
-    # -------------------------------------------------------------- #
+    # ------------------------------------------------------------------ #

    @testset "single tool call via message format" begin
        response = Dict{String,Any}(
@@ -33,12 +33,19 @@ import YiemAgent.agentCore: _extractToolCalls
                ],
            ),
        )
-        has_toolcalls, tc_list = _extractToolCalls(response)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getWeather"
        @test tc_list[1].id == "tc_001"
+        @test tc_list[1].type == "function"
        @test tc_list[1].arguments["city"] == "Bangkok, Thailand"
+        @test assistant_msg isa assistantMessage
+        @test assistant_msg.role == "assistant"
+        @test assistant_msg.stopReason == "tool_calls"
+        @test length(assistant_msg.content) == 1
+        @test assistant_msg.content[1] isa textContent
+        @test assistant_msg.content[1].text == "Let me check the weather."
    end

    @testset "multiple tool calls via message format" begin
@@ -66,13 +73,14 @@ import YiemAgent.agentCore: _extractToolCalls
                ],
            ),
        )
-        has_toolcalls, tc_list = _extractToolCalls(response)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 2
        @test tc_list[1].name == "getWeather"
        @test tc_list[1].arguments["city"] == "Tokyo, Japan"
        @test tc_list[2].name == "getTime"
        @test tc_list[2].arguments["timezone"] == "Asia/Tokyo"
+        @test assistant_msg.role == "assistant"
    end

    @testset "tool call with empty arguments string" begin
@@ -91,11 +99,12 @@ import YiemAgent.agentCore: _extractToolCalls
                ],
            ),
        )
-        has_toolcalls, tc_list = _extractToolCalls(response)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "listTools"
        @test tc_list[1].arguments == Dict{String,Any}()
+        @test assistant_msg.stopReason == "end_turn"
    end

    @testset "tool call with missing id falls back to uuid" begin
@@ -113,14 +122,14 @@ import YiemAgent.agentCore: _extractToolCalls
                ],
            ),
        )
-        has_toolcalls, tc_list = _extractToolCalls(response)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test !isempty(tc_list[1].id)
        @test tc_list[1].name == "getTime"
    end

-    @testset "tool call with non-string arguments (pre-parsed)" begin
+    @testset "tool call with non-string arguments (pre-parsed dict)" begin
        response = Dict{String,Any}(
            "message" => Dict{String,Any}(
                "role" => "assistant",
@@ -136,16 +145,42 @@ import YiemAgent.agentCore: _extractToolCalls
                ],
            ),
        )
-        has_toolcalls, tc_list = _extractToolCalls(response)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].arguments["city"] == "London"
        @test tc_list[1].arguments["units"] == "fahrenheit"
    end

-    # ----------------------------------------------------------- #
+    @testset "tool call with api/provider/model/usage metadata" begin
+        response = Dict{String,Any}(
+            "api" => "openai",
+            "provider" => "anthropic",
+            "model" => "claude-3-opus",
+            "message" => Dict{String,Any}(
+                "role" => "assistant",
+                "tool_calls" => Any[
+                    Dict{String,Any}(
+                        "type" => "function",
+                        "function" => Dict{String,Any}(
+                            "name" => "getTime",
+                            "arguments" => "{}",
+                        ),
+                        "id" => "tc_meta",
+                    ),
+                ],
+            ),
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test has_toolcalls == true
+        @test assistant_msg.api == "openai"
+        @test assistant_msg.provider == "anthropic"
+        @test assistant_msg.model == "claude-3-opus"
+    end
+
+    # --------------------------------------------------------------- #
    # Format 2: response.content blocks (OpenAI API style)            #
-    # ----------------------------------------------------------- #
+    # --------------------------------------------------------------- #

    @testset "content blocks with tool_calls" begin
        response = Dict{String,Any}(
@@ -166,11 +201,14 @@ import YiemAgent.agentCore: _extractToolCalls
                ),
            ],
        )
-        has_toolcalls, tc_list = _extractToolCalls(response)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getWeather"
        @test tc_list[1].arguments["city"] == "Paris"
+        # text block before tool_calls should be included in content
+        @test length(assistant_msg.content) == 1
+        @test assistant_msg.content[1].text == "Let me check."
    end

    @testset "content blocks with tool_call (single-call format)" begin
@@ -184,7 +222,7 @@ import YiemAgent.agentCore: _extractToolCalls
                ),
            ],
        )
-        has_toolcalls, tc_list = _extractToolCalls(response)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getTime"
@@ -192,9 +230,90 @@ import YiemAgent.agentCore: _extractToolCalls
        @test tc_list[1].arguments["timezone"] == "Europe/London"
    end

-    # ----------------------------------------------------------- #
-    # Format 2 via struct-like object (no .content field)         #
-    # ----------------------------------------------------------- #
+    @testset "content blocks with reasoning and text" begin
+        response = Dict{String,Any}(
+            "content" => Any[
+                Dict{String,Any}("type" => "reasoning", "text" => "Thinking..."),
+                Dict{String,Any}("type" => "text", "text" => "Here's the answer."),
+            ],
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test has_toolcalls == false
+        @test length(tc_list) == 0
+        @test length(assistant_msg.content) == 2
+        @test assistant_msg.content[1].text == "Thinking..."
+        @test assistant_msg.content[2].text == "Here's the answer."
+    end
+
+    @testset "content blocks with text and tool_call (tool_call not in content)" begin
+        response = Dict{String,Any}(
+            "content" => Any[
+                Dict{String,Any}("type" => "text", "text" => "Sure, I'll check."),
+                Dict{String,Any}(
+                    "type" => "tool_call",
+                    "id" => "tc_mix",
+                    "name" => "getWeather",
+                    "arguments" => Dict{String,Any}("city" => "London"),
+                ),
+            ],
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test has_toolcalls == true
+        @test length(tc_list) == 1
+        @test tc_list[1].name == "getWeather"
+        # text block included, tool_call block excluded from content
+        @test length(assistant_msg.content) == 1
+        @test assistant_msg.content[1].text == "Sure, I'll check."
+    end
+
+    # ------------------------------------------------------------------ #
+    # assistantMessage construction                                       #
+    # ------------------------------------------------------------------ #
+
+    @testset "assistantMessage with error_message and errorMessage fallback" begin
+        response = Dict{String,Any}(
+            "error_message" => "rate limit",
+            "content" => Any[Dict{String,Any}("type" => "text", "text" => "fail")],
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test has_toolcalls == false
+        @test assistant_msg.errorMessage == "rate limit"
+    end
+
+    @testset "assistantMessage with usage tracking" begin
+        response = Dict{String,Any}(
+            "content" => Any[Dict{String,Any}("type" => "text", "text" => "hi")],
+            "usage" => llmUsage(100, 50),
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test assistant_msg.usage.inputTokens == 100
+        @test assistant_msg.usage.outputTokens == 50
+    end
+
+    @testset "assistantMessage with invalid usage defaults to zero" begin
+        response = Dict{String,Any}(
+            "content" => Any[Dict{String,Any}("type" => "text", "text" => "hi")],
+            "usage" => "invalid",
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test assistant_msg.usage.inputTokens == 0
+        @test assistant_msg.usage.outputTokens == 0
+    end
+
+    @testset "reasoning_content as textContent" begin
+        response = Dict{String,Any}(
+            "reasoning_content" => textContent("internal thought"),
+            "content" => Any[Dict{String,Any}("type" => "text", "text" => "output")],
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test length(assistant_msg.content) == 2
+        @test assistant_msg.content[1].text == "internal thought"
+        @test assistant_msg.content[2].text == "output"
+    end
+
+    # ------------------------------------------------------------------ #
+    # No tool call cases                                                  #
+    # ------------------------------------------------------------------ #

    @testset "no tool calls found" begin
        response = Dict{String,Any}(
@@ -202,16 +321,20 @@ import YiemAgent.agentCore: _extractToolCalls
                Dict{String,Any}("type" => "text", "text" => "Hello world."),
            ],
        )
-        has_toolcalls, tc_list = _extractToolCalls(response)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == false
        @test length(tc_list) == 0
+        @test assistant_msg.stopReason == "end_turn"
    end

    @testset "empty message" begin
        response = Dict{String,Any}()
-        has_toolcalls, tc_list = _extractToolCalls(response)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == false
        @test length(tc_list) == 0
+        @test assistant_msg.role == "assistant"
+        @test assistant_msg.stopReason == "end_turn"
+        @test length(assistant_msg.content) == 0
    end

    @testset "message with empty tool_calls array" begin
@@ -221,9 +344,21 @@ import YiemAgent.agentCore: _extractToolCalls
                "tool_calls" => Any[],
            ),
        )
-        has_toolcalls, tc_list = _extractToolCalls(response)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == false
        @test length(tc_list) == 0
+        @test assistant_msg.role == "assistant"
+    end
+
+    @testset "message field is not a Dict" begin
+        response = Dict{String,Any}(
+            "message" => "not a dict",
+            "content" => Any[Dict{String,Any}("type" => "text", "text" => "fallback")],
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test has_toolcalls == false
+        @test length(tc_list) == 0
+        @test length(assistant_msg.content) == 1
    end

    @testset "Format 1 takes priority over Format 2" begin
@@ -250,15 +385,15 @@ import YiemAgent.agentCore: _extractToolCalls
                ),
            ],
        )
-        has_toolcalls, tc_list = _extractToolCalls(response)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getWeather"
    end

-    # ----------------------------------------------------------- #
-    # edge cases                                                    #
-    # ----------------------------------------------------------- #
+    # ------------------------------------------------------------------ #
+    # Edge cases                                                          #
+    # ------------------------------------------------------------------ #

    @testset "tool call with null arguments" begin
        response = Dict{String,Any}(
@@ -276,7 +411,7 @@ import YiemAgent.agentCore: _extractToolCalls
                ],
            ),
        )
-        has_toolcalls, tc_list = _extractToolCalls(response)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getTime"
@@ -294,7 +429,7 @@ import YiemAgent.agentCore: _extractToolCalls
                ],
            ),
        )
-        has_toolcalls, tc_list = _extractToolCalls(response)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == ""
@@ -313,7 +448,7 @@ import YiemAgent.agentCore: _extractToolCalls
                ],
            ),
        )
-        has_toolcalls, tc_list = _extractToolCalls(response)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == ""
@@ -333,11 +468,171 @@ import YiemAgent.agentCore: _extractToolCalls
            ),
        ))
        parsed = JSON.parse(json_str)
-        has_toolcalls, tc_list = _extractToolCalls(parsed)
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(parsed)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getWeather"
        @test tc_list[1].arguments["city"] == "Test"
    end

+    @testset "tool_calls block with mixed content types (text + tool_calls)" begin
+        response = Dict{String,Any}(
+            "content" => Any[
+                Dict{String,Any}("type" => "text", "text" => "I'll check both."),
+                Dict{String,Any}(
+                    "type" => "tool_calls",
+                    "tool_calls" => Any[
+                        Dict{String,Any}(
+                            "type" => "function",
+                            "function" => Dict{String,Any}(
+                                "name" => "getWeather",
+                                "arguments" => "{\"city\":\"London\"}",
+                            ),
+                            "id" => "tc_mix1",
+                        ),
+                        Dict{String,Any}(
+                            "type" => "function",
+                            "function" => Dict{String,Any}(
+                                "name" => "getTime",
+                                "arguments" => "{\"timezone\":\"UTC\"}",
+                            ),
+                            "id" => "tc_mix2",
+                        ),
+                    ],
+                ),
+            ],
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test has_toolcalls == true
+        @test length(tc_list) == 2
+        @test tc_list[1].name == "getWeather"
+        @test tc_list[2].name == "getTime"
+        @test length(assistant_msg.content) == 1
+        @test assistant_msg.content[1].text == "I'll check both."
+    end
+
+    @testset "tool_call block without arguments field" begin
+        response = Dict{String,Any}(
+            "content" => Any[
+                Dict{String,Any}(
+                    "type" => "tool_call",
+                    "id" => "tc_noargs",
+                    "name" => "getTime",
+                ),
+            ],
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test has_toolcalls == true
+        @test length(tc_list) == 1
+        @test tc_list[1].arguments == Dict{String,Any}()
+    end
+
+    @testset "tool_calls block with empty tool_calls array" begin
+        response = Dict{String,Any}(
+            "content" => Any[
+                Dict{String,Any}(
+                    "type" => "tool_calls",
+                    "tool_calls" => Any[],
+                ),
+            ],
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test has_toolcalls == false
+        @test length(tc_list) == 0
+    end
+
+    @testset "tool_calls block with non-AbstractDict elements" begin
+        response = Dict{String,Any}(
+            "content" => Any[
+                Dict{String,Any}(
+                    "type" => "tool_calls",
+                    "tool_calls" => Any["not a dict", 42, nothing],
+                ),
+            ],
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test has_toolcalls == false
+        @test length(tc_list) == 0
+    end
+
+    @testset "content field is not a Vector" begin
+        response = Dict{String,Any}(
+            "content" => "not a vector",
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test has_toolcalls == false
+        @test length(tc_list) == 0
+        @test length(assistant_msg.content) == 0
+    end
+
+    @testset "Dict-based response with all metadata fields" begin
+        response = Dict{String,Any}(
+            "api" => "openai",
+            "provider" => "anthropic",
+            "model" => "claude-3-sonnet",
+            "content" => Any[
+                Dict{String,Any}(
+                    "type" => "tool_call",
+                    "id" => "tc_meta",
+                    "name" => "getTime",
+                    "arguments" => Dict{String,Any}("city" => "Seoul"),
+                ),
+            ],
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test has_toolcalls == true
+        @test length(tc_list) == 1
+        @test tc_list[1].name == "getTime"
+        @test tc_list[1].arguments["city"] == "Seoul"
+        @test assistant_msg.api == "openai"
+        @test assistant_msg.provider == "anthropic"
+        @test assistant_msg.model == "claude-3-sonnet"
+    end
+
+    @testset "default role is assistant" begin
+        response = Dict{String,Any}(
+            "content" => Any[Dict{String,Any}("type" => "text", "text" => "no role specified")],
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test assistant_msg.role == "assistant"
+    end
+
+    @testset "tool call with custom role in message format" begin
+        response = Dict{String,Any}(
+            "message" => Dict{String,Any}(
+                "role" => "custom_role",
+                "tool_calls" => Any[
+                    Dict{String,Any}(
+                        "type" => "function",
+                        "function" => Dict{String,Any}(
+                            "name" => "getWeather",
+                            "arguments" => "{}",
+                        ),
+                        "id" => "tc_role",
+                    ),
+                ],
+            ),
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test assistant_msg.role == "custom_role"
+    end
+
+    @testset "image content block handling" begin
+        response = Dict{String,Any}(
+            "content" => Any[
+                Dict{String,Any}(
+                    "type" => "image_url",
+                    "image_url" => Dict("url" => "data:image/png;base64,abc123"),
+                ),
+                Dict{String,Any}("type" => "text", "text" => "What is this?"),
+            ],
+        )
+        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
+        @test has_toolcalls == false
+        @test length(assistant_msg.content) == 2
+        @test assistant_msg.content[1] isa textContent
+        @test assistant_msg.content[1].text == ""
+        @test assistant_msg.content[2].text == "What is this?"
+    end
+
 end