Refactor created_at stamping and tracing into streaming_response

Agent-Logs-Url: https://github.com/Prism-Shadow/AgentHub/sessions/57d3d757-2462-486d-b417-46e62dbe6cb8

Co-authored-by: hiyouga <16256802+hiyouga@users.noreply.github.com>

Author: Copilot

src_py/agenthub/base_client.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from typing import Any, AsyncIterator, cast
+from typing import Any, AsyncIterator
 import time
 
 from .types import ContentItem, FinishReason, UniConfig, UniEvent, UniMessage, UsageMetadata
@@ -123,15 +123,13 @@ def concat_uni_events_to_uni_message(self, events: list[UniEvent]) -> UniMessage
             finish_reason = event.get("finish_reason")  # finish_reason is taken from the last event
             created_at = event.get("created_at")  # created_at is taken from the last event
 
-        result: UniMessage = {
+        return {
             "role": "assistant",
             "content_items": content_items,
             "usage_metadata": usage_metadata,
             "finish_reason": finish_reason,
+            "created_at": created_at,
         }
-        if created_at is not None:
-            result["created_at"] = created_at
-        return result
 
     @abstractmethod
     async def _streaming_response_internal(
@@ -173,14 +171,29 @@ async def streaming_response(
         Yields:
             Universal events from the streaming response
         """
+        # Stamp any messages that don't yet have a created_at timestamp
+        for msg in messages:
+            if "created_at" not in msg:
+                msg["created_at"] = int(time.time() * 1000)
+
         last_event: UniEvent | None = None
+        events = []
         async for event in self._streaming_response_internal(messages, config):
             event["created_at"] = int(time.time() * 1000)
             last_event = event
+            events.append(event)
             yield event
 
         self._validate_last_event(last_event)
 
+        # Save history to file if trace_id is specified
+        if config.get("trace_id") and events:
+            from .integration.tracer import Tracer
+
+            assistant_message = self.concat_uni_events_to_uni_message(events)
+            tracer = Tracer()
+            tracer.save_history(self._model, messages + [assistant_message], config["trace_id"], config)
+
     async def streaming_response_stateful(
         self,
         message: UniMessage,
@@ -200,10 +213,6 @@ async def streaming_response_stateful(
         Yields:
             Universal events from the streaming response
         """
-        # Stamp input message with current time if not provided
-        if "created_at" not in message:
-            message = cast(UniMessage, {**message, "created_at": int(time.time() * 1000)})
-
         # Build a temporary messages list for inference without mutating history yet
         temp_messages = self._history + [message]
 
@@ -214,18 +223,12 @@ async def streaming_response_stateful(
             yield event
 
         # Only update history after successful inference
+        # temp_messages[-1] is the user message, now stamped with created_at by streaming_response
         if events:
             assistant_message = self.concat_uni_events_to_uni_message(events)
-            self._history.append(message)
+            self._history.append(temp_messages[-1])
             self._history.append(assistant_message)
 
-        # Save history to file if trace_id is specified
-        if config.get("trace_id"):
-            from .integration.tracer import Tracer
-
-            tracer = Tracer()
-            tracer.save_history(self._model, self._history, config["trace_id"], config)
-
     @staticmethod
     def _validate_last_event(last_event: UniEvent | None) -> None:
         """Validate that the last event has usage_metadata and finish_reason.

src_ts/src/baseClient.ts

@@ -127,10 +127,8 @@ export abstract class LLMClient {
       content_items: contentItems,
       usage_metadata: usageMetadata,
       finish_reason: finishReason,
+      created_at: createdAt ?? undefined,
     };
-    if (createdAt !== null) {
-      result.created_at = createdAt;
-    }
     return result;
   }
 
@@ -162,13 +160,37 @@ export abstract class LLMClient {
     messages: UniMessage[];
     config: UniConfig;
   }): AsyncGenerator<UniEvent> {
+    const { messages, config } = options;
+
+    // Stamp any messages that don't yet have a created_at timestamp
+    for (const msg of messages) {
+      if (msg.created_at == null) {
+        msg.created_at = Date.now();
+      }
+    }
+
     let lastEvent: UniEvent | null = null;
+    const events: UniEvent[] = [];
     for await (const event of this._streamingResponseInternal(options)) {
       event.created_at = Date.now();
       lastEvent = event;
+      events.push(event);
       yield event;
     }
     LLMClient._validateLastEvent(lastEvent);
+
+    // Save history to file if trace_id is specified
+    if (config.trace_id && events.length > 0) {
+      const { Tracer } = await import("./integration/tracer");
+      const assistantMessage = this.concatUniEventsToUniMessage(events);
+      const tracer = new Tracer();
+      tracer.saveHistory(
+        this._model,
+        [...messages, assistantMessage],
+        config.trace_id,
+        config,
+      );
+    }
   }
 
   /**
@@ -186,13 +208,7 @@ export abstract class LLMClient {
     message: UniMessage;
     config: UniConfig;
   }): AsyncGenerator<UniEvent> {
-    let { message } = options;
-    const { config } = options;
-
-    // Stamp input message with current time if not provided (shallow copy to avoid mutating caller's object)
-    if (message.created_at == null) {
-      message = { ...message, created_at: Date.now() };
-    }
+    const { message, config } = options;
 
     const tempMessages = [...this._history, message];
 
@@ -205,17 +221,12 @@ export abstract class LLMClient {
       yield event;
     }
 
+    // tempMessages[-1] is the user message, now stamped with created_at by streamingResponse
     if (events.length > 0) {
       const assistantMessage = this.concatUniEventsToUniMessage(events);
-      this._history.push(message);
+      this._history.push(tempMessages[tempMessages.length - 1]);
       this._history.push(assistantMessage);
     }
-
-    if (config.trace_id) {
-      const { Tracer } = await import("./integration/tracer");
-      const tracer = new Tracer();
-      tracer.saveHistory(this._model, this._history, config.trace_id, config);
-    }
   }
 
   /**