106 lines
2.7 KiB
Go
106 lines
2.7 KiB
Go
|
// Operator utility for an explicitly requested, attributed chat import.
|
||
|
|
// It uses the same permission, consent, extraction and memoir services as Web.
|
||
|
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"data-terminal-cloud/internal/extraction"
|
||
|
|
"data-terminal-cloud/internal/identity"
|
||
|
|
"data-terminal-cloud/internal/memoir"
|
||
|
|
"data-terminal-cloud/internal/outbox"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"flag"
|
||
|
|
"fmt"
|
||
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
||
|
|
"os"
|
||
|
|
"strings"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
func secret(name string) string {
|
||
|
|
if file := os.Getenv(name + "_FILE"); file != "" {
|
||
|
|
b, e := os.ReadFile(file)
|
||
|
|
if e != nil {
|
||
|
|
panic("secret file unavailable")
|
||
|
|
}
|
||
|
|
return strings.TrimSpace(string(b))
|
||
|
|
}
|
||
|
|
return os.Getenv(name)
|
||
|
|
}
|
||
|
|
func main() {
|
||
|
|
if e := run(); e != nil {
|
||
|
|
fmt.Fprintln(os.Stderr, e)
|
||
|
|
os.Exit(1)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
func run() error {
|
||
|
|
input := flag.String("input", "", "JSON file with issuer, account_subject, elder_id, import_id, title, transcript")
|
||
|
|
mode := flag.String("mode", "import", "import, extract, generate, inspect")
|
||
|
|
flag.Parse()
|
||
|
|
var in struct {
|
||
|
|
Issuer string `json:"issuer"`
|
||
|
|
Subject string `json:"account_subject"`
|
||
|
|
memoir.ImportInput
|
||
|
|
}
|
||
|
|
b, e := os.ReadFile(*input)
|
||
|
|
if e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
if e = json.Unmarshal(b, &in); e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||
|
|
defer cancel()
|
||
|
|
pool, e := pgxpool.New(ctx, secret("DATABASE_URL"))
|
||
|
|
if e != nil {
|
||
|
|
return errors.New("database unavailable")
|
||
|
|
}
|
||
|
|
defer pool.Close()
|
||
|
|
key := secret("ARK_API_KEY")
|
||
|
|
base := "https://ark.cn-beijing.volces.com/api/v3"
|
||
|
|
extractor, e := extraction.NewArkGenerator(base, key, "deepseek-v4-1-flash-260910", 90*time.Second)
|
||
|
|
if e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
editor, e := extraction.NewArkJSONGenerator(base, key, "doubao-seed-evolving", 180*time.Second, 65536)
|
||
|
|
if e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
s := &memoir.Service{Pool: pool, Generator: editor, Illustrator: memoir.NewIllustrator(base, key, "doubao-seedream-5-0-pro-260628"), Model: "doubao-seed-evolving"}
|
||
|
|
p := identity.Principal{Issuer: in.Issuer, Subject: in.Subject}
|
||
|
|
switch *mode {
|
||
|
|
case "import":
|
||
|
|
id, e := s.Import(ctx, p, in.ImportInput)
|
||
|
|
if e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
fmt.Println("import", id)
|
||
|
|
case "extract":
|
||
|
|
processor, e := extraction.NewProcessor(pool, extractor)
|
||
|
|
if e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
if e = processor.HandleImport(ctx, outbox.Event{AggregateID: in.ID}); e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
case "generate":
|
||
|
|
id, e := s.Enqueue(ctx, p, in.ElderID)
|
||
|
|
if e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
fmt.Println("job", id)
|
||
|
|
if e = s.Process(ctx, id); e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
case "inspect":
|
||
|
|
default:
|
||
|
|
return errors.New("unknown mode")
|
||
|
|
}
|
||
|
|
result, e := s.Query(ctx, p, in.ElderID)
|
||
|
|
if e != nil {
|
||
|
|
return e
|
||
|
|
}
|
||
|
|
return json.NewEncoder(os.Stdout).Encode(result)
|
||
|
|
}
|