// Package vikingmemory integrates Volcano's managed memory collection, not OpenViking. package vikingmemory import ( "bytes" "context" "crypto/sha256" "encoding/hex" "encoding/json" "errors" "fmt" "io" "net/http" "net/url" "strings" "time" "data-terminal-cloud/internal/outbox" ) const AssistantID = "platform_companion_v1" 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 { digest := sha256.Sum256([]byte(tenantID + ":" + subjectID)) return "u_" + hex.EncodeToString(digest[:]) } 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) if resp.StatusCode == 429 || resp.StatusCode >= 500 || (resp.StatusCode >= 200 && resp.StatusCode < 300) { 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 (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 }