Init

internal/ai/runtime/internal/impl/callbacks/agent_trace_handler.go

@@ -0,0 +1,85 @@
+package callbacks
+
+import (
+	"context"
+	"encoding/json"
+	"strings"
+	"time"
+
+	einotool "github.com/cloudwego/eino/components/tool"
+
+	"github.com/cloudwego/eino/adk"
+)
+
+type ToolMetadata struct {
+	ToolCode   string
+	ServerCode string
+	ToolName   string
+}
+
+type RuntimeTraceHandler struct {
+	*adk.BaseChatModelAgentMiddleware
+	collector      *RuntimeTraceCollector
+	toolMetadataBy map[string]ToolMetadata
+}
+
+func NewRuntimeTraceHandler(collector *RuntimeTraceCollector, toolMetadataBy map[string]ToolMetadata) *RuntimeTraceHandler {
+	return &RuntimeTraceHandler{
+		BaseChatModelAgentMiddleware: &adk.BaseChatModelAgentMiddleware{},
+		collector:                    collector,
+		toolMetadataBy:               toolMetadataBy,
+	}
+}
+
+func (h *RuntimeTraceHandler) WrapInvokableToolCall(_ context.Context, endpoint adk.InvokableToolCallEndpoint, tCtx *adk.ToolContext) (adk.InvokableToolCallEndpoint, error) {
+	return func(ctx context.Context, argumentsInJSON string, opts ...einotool.Option) (string, error) {
+		startedAt := time.Now()
+		result, err := endpoint(ctx, argumentsInJSON, opts...)
+		item := ToolTraceItem{
+			ResultPreview: previewToolText(result, 300),
+			LatencyMs:     time.Since(startedAt).Milliseconds(),
+			Status:        "ok",
+		}
+		if tCtx != nil {
+			item.ToolName = strings.TrimSpace(tCtx.Name)
+			if metadata, ok := h.toolMetadataBy[item.ToolName]; ok {
+				item.ToolCode = metadata.ToolCode
+				item.ServerCode = metadata.ServerCode
+				item.ToolName = metadata.ToolName
+			}
+		}
+		if arguments := parseToolArguments(argumentsInJSON); len(arguments) > 0 {
+			item.Arguments = arguments
+		}
+		if err != nil {
+			item.Status = "error"
+			item.ErrorMessage = err.Error()
+		}
+		h.collector.AddToolItem(item)
+		return result, err
+	}, nil
+}
+
+func parseToolArguments(argumentsInJSON string) map[string]any {
+	argumentsInJSON = strings.TrimSpace(argumentsInJSON)
+	if argumentsInJSON == "" {
+		return nil
+	}
+	ret := make(map[string]any)
+	if err := json.Unmarshal([]byte(argumentsInJSON), &ret); err != nil {
+		return nil
+	}
+	return ret
+}
+
+func previewToolText(text string, limit int) string {
+	if limit <= 0 {
+		return ""
+	}
+	text = strings.TrimSpace(text)
+	runes := []rune(text)
+	if len(runes) <= limit {
+		return text
+	}
+	return string(runes[:limit]) + "..."
+}

internal/ai/runtime/internal/impl/callbacks/runlog_callback.go

@@ -0,0 +1,41 @@
+package callbacks
+
+import (
+	"encoding/json"
+	"sync"
+)
+
+type RuntimeTraceCollector struct {
+	mu   sync.Mutex
+	Data RuntimeTraceData
+}
+
+func NewRuntimeTraceCollector() *RuntimeTraceCollector {
+	ret := &RuntimeTraceCollector{}
+	ret.Data.Version = "v1"
+	ret.Data.Status = "started"
+	return ret
+}
+
+func (c *RuntimeTraceCollector) Marshal() string {
+	if c == nil {
+		return ""
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	buf, err := json.Marshal(c.Data)
+	if err != nil {
+		return ""
+	}
+	return string(buf)
+}
+
+func (c *RuntimeTraceCollector) AddToolItem(item ToolTraceItem) {
+	if c == nil {
+		return
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.Data.Tools.Count++
+	c.Data.Tools.Items = append(c.Data.Tools.Items, item)
+}

internal/ai/runtime/internal/impl/callbacks/trace_callback.go

@@ -0,0 +1,63 @@
+package callbacks
+
+type ToolTraceItem struct {
+	ToolCode      string         `json:"toolCode"`
+	ServerCode    string         `json:"serverCode"`
+	ToolName      string         `json:"toolName"`
+	Arguments     map[string]any `json:"arguments,omitempty"`
+	ResultPreview string         `json:"resultPreview,omitempty"`
+	LatencyMs     int64          `json:"latencyMs,omitempty"`
+	Status        string         `json:"status,omitempty"`
+	ErrorMessage  string         `json:"errorMessage,omitempty"`
+}
+
+type RetrieverTraceItem struct {
+	Query           string  `json:"query,omitempty"`
+	KnowledgeBaseID int64   `json:"knowledgeBaseId,omitempty"`
+	DocumentID      int64   `json:"documentId,omitempty"`
+	DocumentTitle   string  `json:"documentTitle,omitempty"`
+	Score           float64 `json:"score,omitempty"`
+	LatencyMs       int64   `json:"latencyMs,omitempty"`
+}
+
+type RuntimeTraceData struct {
+	Version   string `json:"version"`
+	Status    string `json:"status"`
+	RunID     string `json:"runId,omitempty"`
+	Interrupt struct {
+		CheckPointID string                  `json:"checkPointId,omitempty"`
+		Items        []InterruptTraceContext `json:"items,omitempty"`
+	} `json:"interrupt"`
+	Model struct {
+		Provider string `json:"provider,omitempty"`
+		Name     string `json:"name,omitempty"`
+	} `json:"model"`
+	Input struct {
+		HistoryMessageCount       int      `json:"historyMessageCount,omitempty"`
+		KnowledgeBaseIDs          []int64  `json:"knowledgeBaseIds,omitempty"`
+		ToolCodes                 []string `json:"toolCodes,omitempty"`
+		CurrentUserMessagePreview string   `json:"currentUserMessagePreview,omitempty"`
+	} `json:"input"`
+	Retriever struct {
+		Count int                  `json:"count,omitempty"`
+		Items []RetrieverTraceItem `json:"items,omitempty"`
+	} `json:"retriever"`
+	Tools struct {
+		Count int             `json:"count,omitempty"`
+		Items []ToolTraceItem `json:"items,omitempty"`
+	} `json:"tools"`
+	Output struct {
+		ReplyText    string `json:"replyText,omitempty"`
+		FinishReason string `json:"finishReason,omitempty"`
+	} `json:"output"`
+	Error struct {
+		Message string `json:"message,omitempty"`
+		Stage   string `json:"stage,omitempty"`
+	} `json:"error"`
+}
+
+type InterruptTraceContext struct {
+	Type        string `json:"type,omitempty"`
+	ID          string `json:"id"`
+	InfoPreview string `json:"infoPreview,omitempty"`
+}