story/internal/vikingmemory/client.go

// Package vikingmemory integrates Volcano's managed memory collection, not OpenViking.
package vikingmemory

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"data-terminal-cloud/internal/memoryidentity"
	"data-terminal-cloud/internal/outbox"
)

const AssistantID = memoryidentity.AssistantID

type Options struct {
	BaseURL, ResourceID, CollectionName, APIKey, EventType, ProfileType string
	Timeout                                                             time.Duration
}
type Client struct {
	options Options
	http    *http.Client
}
type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
	RoleID  string `json:"role_id"`
	Time    int64  `json:"time"`
}
type Item struct {
	ID          string                     `json:"id"`
	UserID      []string                   `json:"user_id"`
	AssistantID []string                   `json:"assistant_id"`
	SessionID   string                     `json:"session_id"`
	Score       float64                    `json:"score"`
	MemoryInfo  map[string]json.RawMessage `json:"memory_info"`
	Content     string                     `json:"-"`
}

func NewClient(options Options) (*Client, error) {
	if options.BaseURL == "" {
		options.BaseURL = "https://api-knowledgebase.mlp.cn-beijing.volces.com"
	}
	u, e := url.Parse(options.BaseURL)
	if e != nil || u.Scheme != "https" || u.Host == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" {
		return nil, errors.New("Viking requires an HTTPS endpoint")
	}
	if options.ResourceID == "" || strings.TrimSpace(options.APIKey) == "" || options.EventType == "" || options.ProfileType == "" || options.Timeout <= 0 {
		return nil, errors.New("Viking resource, key, event/profile types and timeout are required")
	}
	options.BaseURL = strings.TrimRight(options.BaseURL, "/")
	options.APIKey = strings.TrimSpace(options.APIKey)
	return &Client{options: options, http: &http.Client{Timeout: options.Timeout, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}}, nil
}
func UserID(tenantID, subjectID string) string {
	return memoryidentity.UserID(tenantID, subjectID)
}
func SessionID(conversationID string, batch int) string {
	return "c_" + strings.ReplaceAll(conversationID, "-", "") + fmt.Sprintf("_%d", batch)
}
func (c *Client) call(ctx context.Context, path string, payload map[string]any, target any) error {
	payload["resource_id"] = c.options.ResourceID
	body, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.options.BaseURL+path, bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "Bearer "+c.options.APIKey)
	req.Header.Set("Content-Type", "application/json")
	resp, err := c.http.Do(req)
	if err != nil {
		return outbox.Retryable(errors.New("Viking request unavailable"))
	}
	defer resp.Body.Close()
	var result struct {
		Code *int            `json:"code"`
		Data json.RawMessage `json:"data"`
	}
	err = json.NewDecoder(io.LimitReader(resp.Body, 2*1024*1024)).Decode(&result)
	if resp.StatusCode < 200 || resp.StatusCode >= 300 || err != nil || result.Code == nil || *result.Code != 0 {
		code := -1
		if result.Code != nil {
			code = *result.Code
		}
		failure := fmt.Errorf("Viking %s HTTP %d code %d", path, resp.StatusCode, code)
		// A valid nonzero business response requires classification/operator review.
		// Only an unreadable successful response has an uncertain write outcome.
		if resp.StatusCode == http.StatusTooManyRequests || code == 1000029 {
			return outbox.RetryAfter(failure, vikingRetryDelay(resp.Header.Get("Retry-After"), time.Now()))
		}
		if resp.StatusCode >= 500 || (resp.StatusCode >= 200 && resp.StatusCode < 300 && (err != nil || result.Code == nil)) {
			return outbox.Retryable(failure)
		}
		return failure
	}
	if target != nil {
		if err := json.Unmarshal(result.Data, target); err != nil {
			return errors.New("Viking returned invalid data")
		}
	}
	return nil
}

func vikingRetryDelay(header string, now time.Time) time.Duration {
	// Rate limits without a usable Retry-After get a full cooldown minute.
	// This prevents the short generic retry schedule exhausting the job.
	if seconds, err := strconv.ParseInt(strings.TrimSpace(header), 10, 64); err == nil && seconds > 0 {
		return time.Duration(min(seconds, int64((24*time.Hour)/time.Second))) * time.Second
	}
	if deadline, err := http.ParseTime(header); err == nil && deadline.After(now) {
		return min(deadline.Sub(now), 24*time.Hour)
	}
	return time.Minute
}
func (c *Client) AddSession(ctx context.Context, sessionID, userID string, messages []Message, when int64) error {
	if sessionID == "" || userID == "" || len(messages) == 0 || len(messages) > 100 {
		return errors.New("invalid Viking session")
	}
	return c.call(ctx, "/api/memory/session/add", map[string]any{"session_id": sessionID, "messages": messages, "metadata": map[string]any{"default_user_id": userID, "default_assistant_id": AssistantID, "time": when}, "extract_memory_type": []string{c.options.EventType, c.options.ProfileType}}, nil)
}
func (c *Client) Search(ctx context.Context, userID, query string, limit int) ([]Item, error) {
	if userID == "" || limit < 1 || limit > 5 {
		return nil, errors.New("Viking user and bounded limit are required")
	}
	items := []Item{}
	remaining := 800
	for _, kind := range []struct{ path, rule string }{{"profile", c.options.ProfileType}, {"event", c.options.EventType}} {
		payload := map[string]any{"filter": map[string]any{"user_id": []string{userID}, "assistant_id": []string{AssistantID}, "memory_type": []string{kind.rule}}, "limit": limit}
		if strings.TrimSpace(query) != "" {
			payload["query"] = query
		}
		var data struct {
			Items []Item `json:"result_list"`
		}
		if err := c.call(ctx, "/api/memory/"+kind.path+"/search", payload, &data); err != nil {
			return nil, err
		}
		for _, item := range data.Items {
			if !contains(item.UserID, userID) || !contains(item.AssistantID, AssistantID) {
				continue
			}
			parts := []string{}
			for _, key := range []string{"summary", "user_profile"} {
				raw := item.MemoryInfo[key]
				if len(raw) == 0 {
					continue
				}
				var text string
				if json.Unmarshal(raw, &text) != nil {
					text = string(raw)
				}
				if text != "" && text != "{}" {
					parts = append(parts, text)
				}
			}
			content := []rune(strings.Join(parts, "；"))
			if len(content) > 400 {
				content = content[:400]
			}
			if len(content) > remaining {
				content = content[:remaining]
			}
			if len(content) == 0 {
				continue
			}
			remaining -= len(content)
			item.Content = string(content)
			items = append(items, item)
			if len(items) >= limit || remaining <= 0 {
				return items, nil
			}
		}
	}
	return items, nil
}
func contains(values []string, expected string) bool {
	for _, v := range values {
		if v == expected {
			return true
		}
	}
	return false
}

func (c *Client) DeleteUser(ctx context.Context, userID string) error {
	if userID == "" {
		return errors.New("Viking deletion requires an exact user")
	}
	for _, kind := range []string{"event", "profile"} {
		if err := c.call(ctx, "/api/memory/"+kind+"/batch_delete", map[string]any{"delete_type": "full", "filter": map[string]any{"user_id": []string{userID}, "assistant_id": []string{AssistantID}}}, nil); err != nil {
			return err
		}
	}
	return nil
}