mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-15 01:18:26 +08:00
fix: support doc generator component (#16858)
### Summary As title <img width="3165" height="2011" alt="image" src="https://github.com/user-attachments/assets/6947aa4f-5dfd-4d8c-b6af-82015b1f1537" /> --- <img width="3735" height="1721" alt="image" src="https://github.com/user-attachments/assets/3a7ec8e8-d5b9-4195-bba9-c435218f6142" />
This commit is contained in:
@@ -25,25 +25,31 @@
|
||||
// archive / oversized-image-stack concerns of the Python
|
||||
// toolchain.
|
||||
//
|
||||
// The component is the canvas entry point. It does NOT call MinIO;
|
||||
// the produced bytes (or for HTML/MD, the rendered text) are
|
||||
// surfaced on the output map for downstream nodes to attach /
|
||||
// serve.
|
||||
// The component is the canvas entry point. It stores generated bytes
|
||||
// in agent attachment storage and surfaces a download descriptor for
|
||||
// downstream Message nodes.
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
iow "ragflow/internal/agent/component/io"
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/storage"
|
||||
"ragflow/internal/utility"
|
||||
)
|
||||
|
||||
const componentNameDocsGenerator = "DocsGenerator"
|
||||
const (
|
||||
componentNameDocsGenerator = "DocsGenerator"
|
||||
componentNameDocGenerator = "DocGenerator"
|
||||
)
|
||||
|
||||
// Default font size for the rendered documents. Plan §2.11.3 row 21
|
||||
// mandates a minimum of 12pt for accessibility; we default to 12.
|
||||
@@ -71,15 +77,16 @@ var validOutputFormats = map[string]bool{
|
||||
|
||||
// docsGeneratorParam is the static DSL param surface.
|
||||
type docsGeneratorParam struct {
|
||||
OutputFormat string `json:"output_format"`
|
||||
Content string `json:"content"`
|
||||
Filename string `json:"filename"`
|
||||
HeaderText string `json:"header_text"`
|
||||
FooterText string `json:"footer_text"`
|
||||
WatermarkText string `json:"watermark_text"`
|
||||
AddPageNumbers bool `json:"add_page_numbers"`
|
||||
AddTimestamp bool `json:"add_timestamp"`
|
||||
FontSize int `json:"font_size"`
|
||||
OutputFormat string `json:"output_format"`
|
||||
Content string `json:"content"`
|
||||
Filename string `json:"filename"`
|
||||
HeaderText string `json:"header_text"`
|
||||
FooterText string `json:"footer_text"`
|
||||
WatermarkText string `json:"watermark_text"`
|
||||
AddPageNumbers bool `json:"add_page_numbers"`
|
||||
AddTimestamp bool `json:"add_timestamp"`
|
||||
FontSize int `json:"font_size"`
|
||||
IncludeDownloadInfoInContent bool `json:"include_download_info_in_content"`
|
||||
}
|
||||
|
||||
// Update copies a fresh params map into the receiver.
|
||||
@@ -98,12 +105,18 @@ func (p *docsGeneratorParam) Update(conf map[string]any) error {
|
||||
}
|
||||
if v, ok := stringFrom(conf, "header_text"); ok {
|
||||
p.HeaderText = v
|
||||
} else if v, ok := stringFrom(conf, "header"); ok {
|
||||
p.HeaderText = v
|
||||
}
|
||||
if v, ok := stringFrom(conf, "footer_text"); ok {
|
||||
p.FooterText = v
|
||||
} else if v, ok := stringFrom(conf, "footer"); ok {
|
||||
p.FooterText = v
|
||||
}
|
||||
if v, ok := stringFrom(conf, "watermark_text"); ok {
|
||||
p.WatermarkText = v
|
||||
} else if v, ok := stringFrom(conf, "watermark"); ok {
|
||||
p.WatermarkText = v
|
||||
}
|
||||
if v, ok := boolFrom(conf, "add_page_numbers"); ok {
|
||||
p.AddPageNumbers = v
|
||||
@@ -120,6 +133,9 @@ func (p *docsGeneratorParam) Update(conf map[string]any) error {
|
||||
} else {
|
||||
p.FontSize = defaultDocsFontSize
|
||||
}
|
||||
if v, ok := boolFrom(conf, "include_download_info_in_content"); ok {
|
||||
p.IncludeDownloadInfoInContent = v
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -144,15 +160,16 @@ func (p *docsGeneratorParam) Check() error {
|
||||
// AsDict returns the param as a plain map.
|
||||
func (p *docsGeneratorParam) AsDict() map[string]any {
|
||||
return map[string]any{
|
||||
"output_format": p.OutputFormat,
|
||||
"content": p.Content,
|
||||
"filename": p.Filename,
|
||||
"header_text": p.HeaderText,
|
||||
"footer_text": p.FooterText,
|
||||
"watermark_text": p.WatermarkText,
|
||||
"add_page_numbers": p.AddPageNumbers,
|
||||
"add_timestamp": p.AddTimestamp,
|
||||
"font_size": p.FontSize,
|
||||
"output_format": p.OutputFormat,
|
||||
"content": p.Content,
|
||||
"filename": p.Filename,
|
||||
"header_text": p.HeaderText,
|
||||
"footer_text": p.FooterText,
|
||||
"watermark_text": p.WatermarkText,
|
||||
"add_page_numbers": p.AddPageNumbers,
|
||||
"add_timestamp": p.AddTimestamp,
|
||||
"font_size": p.FontSize,
|
||||
"include_download_info_in_content": p.IncludeDownloadInfoInContent,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -181,9 +198,14 @@ func (d *DocsGenerator) Name() string { return d.name }
|
||||
// content / filename are honored.
|
||||
func (d *DocsGenerator) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
param := d.param
|
||||
if v, ok := stringFrom(inputs, "content"); ok && v != "" {
|
||||
param.Content = v
|
||||
content, err := resolveDocsGeneratorContent(ctx, param.Content, inputs)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("DocsGenerator: %w", err)
|
||||
}
|
||||
if strings.TrimSpace(content) == "" {
|
||||
return nil, fmt.Errorf("DocsGenerator: content is empty")
|
||||
}
|
||||
param.Content = content
|
||||
if v, ok := stringFrom(inputs, "filename"); ok && v != "" {
|
||||
param.Filename = v
|
||||
}
|
||||
@@ -272,25 +294,123 @@ func (d *DocsGenerator) Invoke(ctx context.Context, inputs map[string]any) (map[
|
||||
|
||||
docID := uuid.New().String()
|
||||
size := len(payload)
|
||||
downloadStub := fmt.Sprintf("inline://docs/%s/%s", docID, safeName)
|
||||
// ext from formatExtension includes a leading dot (".pdf"); strip it
|
||||
// so the preview URL has ?ext=pdf rather than ?ext=.pdf, matching the
|
||||
// Python branch which passes the raw output_format string.
|
||||
extWithoutDot := strings.TrimPrefix(ext, ".")
|
||||
stored, err := storeAgentAttachment(ctx, docID, payload)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("DocsGenerator: store attachment: %w", err)
|
||||
}
|
||||
previewURL := utility.AgentAttachmentPreviewPath(docID, extWithoutDot, mime)
|
||||
downloadURL := agentAttachmentDownloadPath(docID, extWithoutDot, mime, safeName)
|
||||
downloadInfo := map[string]any{
|
||||
"doc_id": docID,
|
||||
"filename": safeName,
|
||||
"mime_type": mime,
|
||||
"size": size,
|
||||
"url": downloadURL,
|
||||
"download": downloadURL,
|
||||
"preview_url": previewURL,
|
||||
"include_download_info_in_content": param.IncludeDownloadInfoInContent,
|
||||
}
|
||||
downloadJSON, err := json.Marshal(downloadInfo)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("DocsGenerator: marshal download info: %w", err)
|
||||
}
|
||||
|
||||
return map[string]any{
|
||||
"doc_id": docID,
|
||||
"filename": safeName,
|
||||
"mime_type": mime,
|
||||
"size": size,
|
||||
"bytes": payload,
|
||||
"download": downloadStub,
|
||||
"preview_url": previewURL,
|
||||
"created": time.Now().UTC().Format(time.RFC3339),
|
||||
"doc_id": docID,
|
||||
"filename": safeName,
|
||||
"mime_type": mime,
|
||||
"size": size,
|
||||
"bytes": payload,
|
||||
"download": string(downloadJSON),
|
||||
"download_info": downloadInfo,
|
||||
"preview_url": previewURL,
|
||||
"stored": stored,
|
||||
"include_download_info_in_content": param.IncludeDownloadInfoInContent,
|
||||
"created": time.Now().UTC().Format(time.RFC3339),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func resolveDocsGeneratorContent(ctx context.Context, configured string, inputs map[string]any) (string, error) {
|
||||
content := configured
|
||||
if v, ok := stringFrom(inputs, "content"); ok && strings.TrimSpace(v) != "" {
|
||||
content = v
|
||||
}
|
||||
if strings.TrimSpace(content) == "" {
|
||||
for _, key := range []string{"query", "output", "text"} {
|
||||
if v, ok := stringFrom(inputs, key); ok && strings.TrimSpace(v) != "" {
|
||||
content = v
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if !strings.Contains(content, "{{") {
|
||||
return stripThinking(content), nil
|
||||
}
|
||||
|
||||
state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if state == nil {
|
||||
return "", fmt.Errorf("nil canvas state")
|
||||
}
|
||||
resolved, err := runtime.ResolveTemplate(content, state)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("resolve content template: %w", err)
|
||||
}
|
||||
return stripThinking(resolved), nil
|
||||
}
|
||||
|
||||
func stripThinking(content string) string {
|
||||
if idx := strings.LastIndex(content, "</think>"); idx >= 0 {
|
||||
return strings.TrimSpace(content[idx+len("</think>"):])
|
||||
}
|
||||
return content
|
||||
}
|
||||
|
||||
func storeAgentAttachment(ctx context.Context, docID string, payload []byte) (bool, error) {
|
||||
state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx)
|
||||
if err != nil || state == nil {
|
||||
return false, nil
|
||||
}
|
||||
tenantID, _ := state.Sys["tenant_id"].(string)
|
||||
tenantID = strings.TrimSpace(tenantID)
|
||||
if tenantID == "" {
|
||||
return false, nil
|
||||
}
|
||||
storageImpl := storage.GetStorageFactory().GetStorage()
|
||||
if storageImpl == nil {
|
||||
return false, fmt.Errorf("storage not initialized")
|
||||
}
|
||||
bucket := fmt.Sprintf("%s-downloads", tenantID)
|
||||
if err := storageImpl.Put(bucket, docID, payload); err != nil {
|
||||
return false, err
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func agentAttachmentDownloadPath(attachmentID, ext, mimeType, filename string) string {
|
||||
path := "/api/v1/agents/attachments/" + attachmentID + "/download"
|
||||
params := url.Values{}
|
||||
if ext != "" {
|
||||
params.Set("ext", ext)
|
||||
}
|
||||
if mimeType != "" {
|
||||
params.Set("mime_type", mimeType)
|
||||
}
|
||||
if filename != "" {
|
||||
params.Set("filename", filename)
|
||||
}
|
||||
if qs := params.Encode(); qs != "" {
|
||||
return path + "?" + qs
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
// Stream mirrors Invoke; DocsGenerator is a single-shot generator.
|
||||
func (d *DocsGenerator) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) {
|
||||
out, err := d.Invoke(ctx, inputs)
|
||||
@@ -312,17 +432,29 @@ func (d *DocsGenerator) Inputs() map[string]string {
|
||||
}
|
||||
}
|
||||
|
||||
// GetInputForm returns the runtime input-form schema for dynamic component
|
||||
// input requests.
|
||||
func (d *DocsGenerator) GetInputForm() map[string]any {
|
||||
return map[string]any{
|
||||
"content": map[string]any{
|
||||
"name": "Content",
|
||||
"type": "line",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Outputs returns the response surface.
|
||||
func (d *DocsGenerator) Outputs() map[string]string {
|
||||
return map[string]string{
|
||||
"doc_id": "Generated document id (UUID).",
|
||||
"filename": "Sanitized filename (extension matches output_format).",
|
||||
"mime_type": "MIME type for the payload.",
|
||||
"size": "Payload size in bytes.",
|
||||
"bytes": "Raw document bytes (for storage upload).",
|
||||
"download": "Stub URI the canvas engine can resolve to a signed URL.",
|
||||
"preview_url": "URL path for inline preview of the generated document.",
|
||||
"created": "RFC3339 timestamp of the generation.",
|
||||
"doc_id": "Generated document id (UUID).",
|
||||
"filename": "Sanitized filename (extension matches output_format).",
|
||||
"mime_type": "MIME type for the payload.",
|
||||
"size": "Payload size in bytes.",
|
||||
"bytes": "Raw document bytes (for storage upload).",
|
||||
"download": "JSON download descriptor for Message nodes.",
|
||||
"download_info": "Structured download descriptor for Message nodes.",
|
||||
"preview_url": "URL path for inline preview of the generated document.",
|
||||
"created": "RFC3339 timestamp of the generation.",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -383,4 +515,5 @@ func sanitizeFilename(raw, ext string) string {
|
||||
|
||||
func init() {
|
||||
Register(componentNameDocsGenerator, NewDocsGenerator)
|
||||
Register(componentNameDocGenerator, NewDocsGenerator)
|
||||
}
|
||||
|
||||
@@ -18,7 +18,12 @@ package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/canvas"
|
||||
"ragflow/internal/storage"
|
||||
)
|
||||
|
||||
// TestDocsGenerator_Registered: the component is registered under
|
||||
@@ -43,6 +48,40 @@ func TestDocsGenerator_Registered(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocGeneratorAlias_Registered(t *testing.T) {
|
||||
c, err := New("DocGenerator", map[string]any{
|
||||
"output_format": "pdf",
|
||||
"content": "Hello world",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New(DocGenerator): %v", err)
|
||||
}
|
||||
if c.Name() != "DocsGenerator" {
|
||||
t.Errorf("Name=%q, want DocsGenerator", c.Name())
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocGeneratorAlias_InputForm(t *testing.T) {
|
||||
c, err := New("DocGenerator", map[string]any{
|
||||
"output_format": "pdf",
|
||||
"content": "Hello world",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New(DocGenerator): %v", err)
|
||||
}
|
||||
formGetter, ok := c.(interface{ GetInputForm() map[string]any })
|
||||
if !ok {
|
||||
t.Fatal("DocGenerator component does not expose GetInputForm")
|
||||
}
|
||||
content, ok := formGetter.GetInputForm()["content"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("GetInputForm()[content] has type %T, want map", formGetter.GetInputForm()["content"])
|
||||
}
|
||||
if content["type"] != "line" {
|
||||
t.Fatalf("GetInputForm()[content][type] = %v, want line", content["type"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestDocsGenerator_Invoke_HappyPath: with valid params, the
|
||||
// component runs without error and produces a non-nil output map.
|
||||
// Real PDF/DOCX generation needs an actual font file on disk
|
||||
@@ -62,3 +101,129 @@ func TestDocsGenerator_Invoke_HappyPath(t *testing.T) {
|
||||
// the internal writer (which may not be available in this
|
||||
// checkout). The test pins that the call doesn't panic.
|
||||
}
|
||||
|
||||
func TestDocsGenerator_Invoke_UsesQueryWhenContentEmpty(t *testing.T) {
|
||||
c, err := New("DocsGenerator", map[string]any{
|
||||
"output_format": "txt",
|
||||
"filename": "query-body",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New(DocsGenerator): %v", err)
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"query": "Hello from begin query",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
payload, ok := out["bytes"].([]byte)
|
||||
if !ok {
|
||||
t.Fatalf("bytes has type %T, want []byte", out["bytes"])
|
||||
}
|
||||
if !strings.Contains(string(payload), "Hello from begin query") {
|
||||
t.Fatalf("payload = %q, want query content", string(payload))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocsGenerator_Invoke_AcceptsDecorationParamNames(t *testing.T) {
|
||||
c, err := New("DocsGenerator", map[string]any{
|
||||
"output_format": "txt",
|
||||
"content": "body",
|
||||
"filename": "decorated",
|
||||
"header_text": "Canonical Header",
|
||||
"footer_text": "Canonical Footer",
|
||||
"watermark_text": "Canonical Watermark",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New(DocsGenerator): %v", err)
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke canonical params: %v", err)
|
||||
}
|
||||
payload, ok := out["bytes"].([]byte)
|
||||
if !ok {
|
||||
t.Fatalf("bytes has type %T, want []byte", out["bytes"])
|
||||
}
|
||||
text := string(payload)
|
||||
if !strings.Contains(text, "Canonical Header") || !strings.Contains(text, "Canonical Footer") {
|
||||
t.Fatalf("payload = %q, want canonical header and footer", text)
|
||||
}
|
||||
|
||||
c, err = New("DocsGenerator", map[string]any{
|
||||
"output_format": "txt",
|
||||
"content": "body",
|
||||
"filename": "decorated-alias",
|
||||
"header": "Alias Header",
|
||||
"footer": "Alias Footer",
|
||||
"watermark": "Alias Watermark",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New(DocsGenerator alias params): %v", err)
|
||||
}
|
||||
out, err = c.Invoke(context.Background(), map[string]any{})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke alias params: %v", err)
|
||||
}
|
||||
payload, ok = out["bytes"].([]byte)
|
||||
if !ok {
|
||||
t.Fatalf("bytes has type %T, want []byte", out["bytes"])
|
||||
}
|
||||
text = string(payload)
|
||||
if !strings.Contains(text, "Alias Header") || !strings.Contains(text, "Alias Footer") {
|
||||
t.Fatalf("payload = %q, want alias header and footer", text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDocsGenerator_Invoke_StoresAgentAttachment(t *testing.T) {
|
||||
storageFactory := storage.GetStorageFactory()
|
||||
prevStorage := storageFactory.GetStorage()
|
||||
memStorage := storage.NewMemoryStorage()
|
||||
storageFactory.SetStorage(memStorage)
|
||||
t.Cleanup(func() { storageFactory.SetStorage(prevStorage) })
|
||||
|
||||
c, err := New("DocGenerator", map[string]any{
|
||||
"output_format": "txt",
|
||||
"content": "Hello storage",
|
||||
"filename": "stored",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New(DocGenerator): %v", err)
|
||||
}
|
||||
state := canvas.NewCanvasState("run-1", "task-1")
|
||||
state.Sys["tenant_id"] = "tenant-1"
|
||||
ctx := canvas.WithState(context.Background(), state)
|
||||
|
||||
out, err := c.Invoke(ctx, map[string]any{})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if out["stored"] != true {
|
||||
t.Fatalf("stored = %v, want true", out["stored"])
|
||||
}
|
||||
docID, ok := out["doc_id"].(string)
|
||||
if !ok || docID == "" {
|
||||
t.Fatalf("doc_id = %v, want non-empty string", out["doc_id"])
|
||||
}
|
||||
blob, err := memStorage.Get("tenant-1-downloads", docID)
|
||||
if err != nil {
|
||||
t.Fatalf("stored blob missing: %v", err)
|
||||
}
|
||||
if string(blob) == "" || !strings.Contains(string(blob), "Hello storage") {
|
||||
t.Fatalf("stored blob = %q, want generated content", string(blob))
|
||||
}
|
||||
download, _ := out["download"].(string)
|
||||
var info map[string]any
|
||||
if err := json.Unmarshal([]byte(download), &info); err != nil {
|
||||
t.Fatalf("download is not JSON: %v", err)
|
||||
}
|
||||
if info["doc_id"] != docID {
|
||||
t.Fatalf("download doc_id = %v, want %s", info["doc_id"], docID)
|
||||
}
|
||||
if path, _ := info["url"].(string); !strings.HasPrefix(path, "/api/v1/agents/attachments/"+docID+"/download") {
|
||||
t.Fatalf("download url = %q, want attachment download path", path)
|
||||
}
|
||||
if _, ok := out["download_info"].(map[string]any); !ok {
|
||||
t.Fatalf("download_info has type %T, want map", out["download_info"])
|
||||
}
|
||||
}
|
||||
|
||||
@@ -14,27 +14,22 @@
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Package io — PDF writer (signintech/gopdf).
|
||||
// Package io provides a PDF writer backed by signintech/gopdf.
|
||||
//
|
||||
// WritePDF renders the supplied content to a PDF using the
|
||||
// MIT-licensed signintech/gopdf library. The writer probes via
|
||||
// gopdf.SetFont; if the family is unknown, it surfaces
|
||||
// ErrPDFFontNotConfigured so the orchestrator can return a clear
|
||||
// deployment-time error. Production deployments register a TTF
|
||||
// (e.g. Noto Sans CJK SC) at startup.
|
||||
//
|
||||
// When a TTF *is* registered, the writer emits a simple
|
||||
// one-paragraph page per line of content, with a centered header
|
||||
// and a centered footer carrying the page number / timestamp when
|
||||
// requested.
|
||||
// WritePDF renders the supplied content with gopdf. The writer registers a
|
||||
// Latin font and, when available, a separate CJK fallback font, then switches
|
||||
// fonts per text segment. This avoids the blank-page failure where ASCII text
|
||||
// is sent through a CJK fallback font that does not expose ASCII glyphs.
|
||||
package io
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode"
|
||||
|
||||
"github.com/signintech/gopdf"
|
||||
)
|
||||
@@ -50,114 +45,182 @@ type PDFOptions struct {
|
||||
FontFamily string
|
||||
}
|
||||
|
||||
// ErrPDFFontNotConfigured is returned when no TTF is registered.
|
||||
// Callers should register a TTF via gopdf.SetFont before invoking
|
||||
// WritePDF.
|
||||
var ErrPDFFontNotConfigured = errors.New("PDF font not configured: register a TTF (e.g. Noto Sans CJK SC) via gopdf.SetFont before calling WritePDF")
|
||||
var ErrPDFFontNotConfigured = errors.New("PDF font not configured: install a TTF such as DejaVu Sans or Noto Sans CJK SC")
|
||||
|
||||
const (
|
||||
pdfLatinFontFamily = "RAGFlowLatin"
|
||||
pdfCJKFontFamily = "RAGFlowCJK"
|
||||
)
|
||||
|
||||
var defaultPDFLatinFontPaths = []string{
|
||||
"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
|
||||
"/usr/share/fonts/truetype/liberation2/LiberationSans-Regular.ttf",
|
||||
"/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
|
||||
}
|
||||
|
||||
var defaultPDFCJKFontPaths = []string{
|
||||
"/usr/share/fonts/truetype/droid/DroidSansFallbackFull.ttf",
|
||||
"/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.ttf",
|
||||
"/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",
|
||||
"/usr/share/fonts/truetype/arphic/uming.ttc",
|
||||
"/usr/share/fonts/truetype/arphic/ukai.ttc",
|
||||
}
|
||||
|
||||
type pdfFontSet struct {
|
||||
latinFamily string
|
||||
cjkFamily string
|
||||
hasCJK bool
|
||||
}
|
||||
|
||||
// WritePDF renders the content to a PDF byte stream.
|
||||
//
|
||||
// Layout:
|
||||
//
|
||||
// - A4 portrait, 36pt margins on all sides.
|
||||
// - Body lines are drawn top-to-bottom, one per line of content.
|
||||
// - Header is centered at the top of every page (when set).
|
||||
// - Footer is centered at the bottom of every page and may include
|
||||
// the footer text, a generation timestamp, and a page number.
|
||||
// - Watermark is rendered as grey text near the page center.
|
||||
//
|
||||
// When the requested font family is not registered, the function
|
||||
// returns ErrPDFFontNotConfigured and does not write any output.
|
||||
func WritePDF(content string, opts PDFOptions) ([]byte, error) {
|
||||
if opts.FontSize <= 0 {
|
||||
opts.FontSize = 12
|
||||
}
|
||||
if opts.FontFamily == "" {
|
||||
opts.FontFamily = "Noto Sans CJK SC"
|
||||
}
|
||||
|
||||
pdf := &gopdf.GoPdf{}
|
||||
pdf.Start(gopdf.Config{PageSize: *gopdf.PageSizeA4})
|
||||
pdf.AddPage()
|
||||
|
||||
// Probe the font registry. gopdf returns an error like "font not
|
||||
// found" when the family is not registered; we surface that as
|
||||
// ErrPDFFontNotConfigured so callers can map it to a clear
|
||||
// deployment message.
|
||||
if err := pdf.SetFont(opts.FontFamily, "", opts.FontSize); err != nil {
|
||||
if isFontNotFound(err) {
|
||||
return nil, ErrPDFFontNotConfigured
|
||||
}
|
||||
return nil, fmt.Errorf("PDF: set font %q: %w", opts.FontFamily, err)
|
||||
fonts, err := ensurePDFFonts(pdf, opts.FontSize)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
pdf.AddPage()
|
||||
drawHeader(pdf, opts)
|
||||
drawHeader(pdf, fonts, opts)
|
||||
|
||||
// Body — one Cell per line, manual y-cursor.
|
||||
bodyX := 36.0
|
||||
bodyY := 72.0
|
||||
lineHeight := float64(opts.FontSize) * 1.5
|
||||
pdf.SetX(bodyX)
|
||||
pdf.SetY(bodyY)
|
||||
pageNumber := 1
|
||||
|
||||
for _, line := range splitLines(content) {
|
||||
if line == "" {
|
||||
// Preserve blank lines as vertical space.
|
||||
bodyY += lineHeight
|
||||
if bodyY > 760 {
|
||||
drawFooter(pdf, opts)
|
||||
drawFooter(pdf, fonts, opts, pageNumber)
|
||||
pdf.AddPage()
|
||||
drawHeader(pdf, opts)
|
||||
drawHeader(pdf, fonts, opts)
|
||||
bodyY = 72.0
|
||||
pageNumber++
|
||||
}
|
||||
pdf.SetX(bodyX)
|
||||
pdf.SetY(bodyY)
|
||||
continue
|
||||
}
|
||||
if bodyY > 760 {
|
||||
drawFooter(pdf, opts)
|
||||
drawFooter(pdf, fonts, opts, pageNumber)
|
||||
pdf.AddPage()
|
||||
drawHeader(pdf, opts)
|
||||
drawHeader(pdf, fonts, opts)
|
||||
bodyY = 72.0
|
||||
pageNumber++
|
||||
}
|
||||
pdf.SetX(bodyX)
|
||||
pdf.SetY(bodyY)
|
||||
if err := pdf.Cell(nil, line); err != nil {
|
||||
return nil, fmt.Errorf("PDF: cell: %w", err)
|
||||
if err := drawPDFText(pdf, fonts, line, opts.FontSize); err != nil {
|
||||
return nil, fmt.Errorf("PDF: body text: %w", err)
|
||||
}
|
||||
bodyY += lineHeight
|
||||
}
|
||||
|
||||
if opts.WatermarkText != "" {
|
||||
drawWatermark(pdf, opts)
|
||||
drawWatermark(pdf, fonts, opts)
|
||||
}
|
||||
drawFooter(pdf, opts)
|
||||
drawFooter(pdf, fonts, opts, pageNumber)
|
||||
|
||||
return writePDFToBytes(pdf)
|
||||
}
|
||||
|
||||
// drawHeader emits the header text at the top of the current page.
|
||||
// gopdf's API in v0.36.x doesn't expose a Header() callback; we draw
|
||||
// at the top of every page after AddPage.
|
||||
func drawHeader(pdf *gopdf.GoPdf, opts PDFOptions) {
|
||||
func ensurePDFFonts(pdf *gopdf.GoPdf, size int) (pdfFontSet, error) {
|
||||
latinPath := resolvePDFLatinFontPath()
|
||||
if latinPath == "" {
|
||||
return pdfFontSet{}, ErrPDFFontNotConfigured
|
||||
}
|
||||
if err := pdf.AddTTFFont(pdfLatinFontFamily, latinPath); err != nil {
|
||||
return pdfFontSet{}, fmt.Errorf("PDF: add latin font from %s: %w", latinPath, err)
|
||||
}
|
||||
if err := pdf.SetFont(pdfLatinFontFamily, "", size); err != nil {
|
||||
return pdfFontSet{}, fmt.Errorf("PDF: set latin font: %w", err)
|
||||
}
|
||||
|
||||
fonts := pdfFontSet{latinFamily: pdfLatinFontFamily, cjkFamily: pdfLatinFontFamily}
|
||||
cjkPath := resolvePDFCJKFontPath()
|
||||
if cjkPath == "" || cjkPath == latinPath {
|
||||
return fonts, nil
|
||||
}
|
||||
if err := pdf.AddTTFFont(pdfCJKFontFamily, cjkPath); err != nil {
|
||||
return fonts, nil
|
||||
}
|
||||
if err := pdf.SetFont(pdfCJKFontFamily, "", size); err != nil {
|
||||
return fonts, nil
|
||||
}
|
||||
fonts.cjkFamily = pdfCJKFontFamily
|
||||
fonts.hasCJK = true
|
||||
_ = pdf.SetFont(pdfLatinFontFamily, "", size)
|
||||
return fonts, nil
|
||||
}
|
||||
|
||||
func resolvePDFLatinFontPath() string {
|
||||
return resolvePDFFontPath("RAGFLOW_PDF_LATIN_FONT_PATH", defaultPDFLatinFontPaths)
|
||||
}
|
||||
|
||||
func resolvePDFCJKFontPath() string {
|
||||
if explicit := strings.TrimSpace(os.Getenv("RAGFLOW_PDF_FONT_PATH")); explicit != "" {
|
||||
if path := normalizeExistingFontPath(explicit); path != "" {
|
||||
return path
|
||||
}
|
||||
}
|
||||
return resolvePDFFontPath("RAGFLOW_PDF_CJK_FONT_PATH", defaultPDFCJKFontPaths)
|
||||
}
|
||||
|
||||
func resolvePDFFontPath(envKey string, candidates []string) string {
|
||||
if explicit := strings.TrimSpace(os.Getenv(envKey)); explicit != "" {
|
||||
if path := normalizeExistingFontPath(explicit); path != "" {
|
||||
return path
|
||||
}
|
||||
}
|
||||
for _, candidate := range candidates {
|
||||
if path := normalizeExistingFontPath(candidate); path != "" {
|
||||
return path
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func normalizeExistingFontPath(candidate string) string {
|
||||
path := strings.TrimSpace(candidate)
|
||||
if path == "" {
|
||||
return ""
|
||||
}
|
||||
if !filepath.IsAbs(path) {
|
||||
if abs, err := filepath.Abs(path); err == nil {
|
||||
path = abs
|
||||
}
|
||||
}
|
||||
info, err := os.Stat(path)
|
||||
if err == nil && !info.IsDir() {
|
||||
return path
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func drawHeader(pdf *gopdf.GoPdf, fonts pdfFontSet, opts PDFOptions) {
|
||||
if opts.HeaderText == "" {
|
||||
return
|
||||
}
|
||||
_ = pdf.SetFont(opts.FontFamily, "", opts.FontSize-2)
|
||||
size := overlayFontSize(opts)
|
||||
pdf.SetX(36)
|
||||
pdf.SetY(24)
|
||||
_ = pdf.Cell(nil, opts.HeaderText)
|
||||
// Restore body font.
|
||||
_ = pdf.SetFont(opts.FontFamily, "", opts.FontSize)
|
||||
_ = drawPDFText(pdf, fonts, opts.HeaderText, size)
|
||||
}
|
||||
|
||||
// drawFooter emits the footer text plus optional timestamp / page
|
||||
// number at the bottom of the current page.
|
||||
func drawFooter(pdf *gopdf.GoPdf, opts PDFOptions) {
|
||||
func drawFooter(pdf *gopdf.GoPdf, fonts pdfFontSet, opts PDFOptions, pageNumber int) {
|
||||
if opts.FooterText == "" && !opts.AddTimestamp && !opts.AddPageNumbers {
|
||||
return
|
||||
}
|
||||
_ = pdf.SetFont(opts.FontFamily, "", opts.FontSize-2)
|
||||
size := overlayFontSize(opts)
|
||||
pdf.SetX(36)
|
||||
pdf.SetY(800)
|
||||
parts := []string{}
|
||||
@@ -168,36 +231,81 @@ func drawFooter(pdf *gopdf.GoPdf, opts PDFOptions) {
|
||||
parts = append(parts, time.Now().UTC().Format("2006-01-02 15:04"))
|
||||
}
|
||||
if opts.AddPageNumbers {
|
||||
// gopdf doesn't expose a page-number macro; emit a
|
||||
// literal placeholder until upstream adds one.
|
||||
parts = append(parts, "Page #")
|
||||
parts = append(parts, fmt.Sprintf("Pages %d", pageNumber))
|
||||
}
|
||||
_ = pdf.Cell(nil, strings.Join(parts, " | "))
|
||||
// Restore body font.
|
||||
_ = pdf.SetFont(opts.FontFamily, "", opts.FontSize)
|
||||
_ = drawPDFText(pdf, fonts, strings.Join(parts, " | "), size)
|
||||
}
|
||||
|
||||
// drawWatermark emits a centered grey watermark. Full rotation is
|
||||
// not in the gopdf v0.36.x public surface; we use a light grey fill
|
||||
// as a visual proxy.
|
||||
func drawWatermark(pdf *gopdf.GoPdf, opts PDFOptions) {
|
||||
func drawWatermark(pdf *gopdf.GoPdf, fonts pdfFontSet, opts PDFOptions) {
|
||||
if opts.WatermarkText == "" {
|
||||
return
|
||||
}
|
||||
_ = pdf.SetFont(opts.FontFamily, "", 48)
|
||||
pdf.SetTextColor(200, 200, 200)
|
||||
pdf.SetX(120)
|
||||
pdf.SetY(360)
|
||||
_ = pdf.Cell(nil, opts.WatermarkText)
|
||||
// Restore.
|
||||
_ = drawPDFText(pdf, fonts, opts.WatermarkText, 48)
|
||||
pdf.SetTextColor(0, 0, 0)
|
||||
_ = pdf.SetFont(opts.FontFamily, "", opts.FontSize)
|
||||
}
|
||||
|
||||
// writePDFToBytes serializes the gopdf output to a byte slice.
|
||||
//
|
||||
// gopdf's Write method requires an *os.File (it needs random access
|
||||
// for the xref table), so we route through a TempFile.
|
||||
func overlayFontSize(opts PDFOptions) int {
|
||||
size := opts.FontSize - 2
|
||||
if size < 1 {
|
||||
return 1
|
||||
}
|
||||
return size
|
||||
}
|
||||
|
||||
func drawPDFText(pdf *gopdf.GoPdf, fonts pdfFontSet, text string, size int) error {
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
for _, segment := range splitByPDFFont(text) {
|
||||
family := fonts.latinFamily
|
||||
if segment.cjk && fonts.hasCJK {
|
||||
family = fonts.cjkFamily
|
||||
}
|
||||
if err := pdf.SetFont(family, "", size); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := pdf.Text(segment.text); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type pdfTextSegment struct {
|
||||
text string
|
||||
cjk bool
|
||||
}
|
||||
|
||||
func splitByPDFFont(text string) []pdfTextSegment {
|
||||
runes := []rune(text)
|
||||
if len(runes) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := []pdfTextSegment{}
|
||||
start := 0
|
||||
current := needsCJKFont(runes[0])
|
||||
for i := 1; i < len(runes); i++ {
|
||||
next := needsCJKFont(runes[i])
|
||||
if next == current {
|
||||
continue
|
||||
}
|
||||
out = append(out, pdfTextSegment{text: string(runes[start:i]), cjk: current})
|
||||
start = i
|
||||
current = next
|
||||
}
|
||||
out = append(out, pdfTextSegment{text: string(runes[start:]), cjk: current})
|
||||
return out
|
||||
}
|
||||
|
||||
func needsCJKFont(r rune) bool {
|
||||
return unicode.In(r, unicode.Han, unicode.Hangul, unicode.Hiragana, unicode.Katakana) ||
|
||||
(r >= 0x3000 && r <= 0x303f) ||
|
||||
(r >= 0xff00 && r <= 0xffef)
|
||||
}
|
||||
|
||||
func writePDFToBytes(pdf *gopdf.GoPdf) ([]byte, error) {
|
||||
tmp, err := os.CreateTemp("", "ragflow-pdf-*.pdf")
|
||||
if err != nil {
|
||||
@@ -229,14 +337,3 @@ func splitLines(content string) []string {
|
||||
}
|
||||
return lines
|
||||
}
|
||||
|
||||
// isFontNotFound reports whether the gopdf error indicates a missing
|
||||
// TTF registration. We match the substrings that have been stable
|
||||
// across recent gopdf versions.
|
||||
func isFontNotFound(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
s := strings.ToLower(err.Error())
|
||||
return strings.Contains(s, "font") && (strings.Contains(s, "not") || strings.Contains(s, "no such") || strings.Contains(s, "undefined"))
|
||||
}
|
||||
|
||||
125
internal/agent/component/io/pdf_writer_test.go
Normal file
125
internal/agent/component/io/pdf_writer_test.go
Normal file
@@ -0,0 +1,125 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package io
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"image"
|
||||
_ "image/png"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestWritePDF_UsesGoLibrary(t *testing.T) {
|
||||
requirePDFLatinFont(t)
|
||||
out, err := WritePDF("Hello\n中文", PDFOptions{})
|
||||
if err != nil {
|
||||
t.Fatalf("WritePDF: %v", err)
|
||||
}
|
||||
if !bytes.HasPrefix(out, []byte("%PDF-")) {
|
||||
t.Fatalf("PDF output missing magic header: %q", out[:min(len(out), 8)])
|
||||
}
|
||||
}
|
||||
|
||||
func TestWritePDF_RendersVisibleText(t *testing.T) {
|
||||
requirePDFLatinFont(t)
|
||||
if _, err := exec.LookPath("pdftoppm"); err != nil {
|
||||
t.Skip("pdftoppm not available")
|
||||
}
|
||||
if _, err := exec.LookPath("pdftotext"); err != nil {
|
||||
t.Skip("pdftotext not available")
|
||||
}
|
||||
|
||||
out, err := WritePDF("Visible PDF body", PDFOptions{
|
||||
HeaderText: "Visible Header",
|
||||
FooterText: "Visible Footer",
|
||||
AddPageNumbers: true,
|
||||
AddTimestamp: false,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("WritePDF: %v", err)
|
||||
}
|
||||
|
||||
dir := t.TempDir()
|
||||
pdfPath := filepath.Join(dir, "visible.pdf")
|
||||
if err := os.WriteFile(pdfPath, out, 0o600); err != nil {
|
||||
t.Fatalf("WriteFile pdf: %v", err)
|
||||
}
|
||||
prefix := filepath.Join(dir, "page")
|
||||
cmd := exec.Command("pdftoppm", "-png", "-singlefile", pdfPath, prefix)
|
||||
if output, err := cmd.CombinedOutput(); err != nil {
|
||||
t.Fatalf("pdftoppm: %v: %s", err, output)
|
||||
}
|
||||
pngFile, err := os.Open(prefix + ".png")
|
||||
if err != nil {
|
||||
t.Fatalf("Open png: %v", err)
|
||||
}
|
||||
defer pngFile.Close()
|
||||
img, _, err := image.Decode(pngFile)
|
||||
if err != nil {
|
||||
t.Fatalf("Decode png: %v", err)
|
||||
}
|
||||
bounds := img.Bounds()
|
||||
nonWhite := 0
|
||||
for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
|
||||
for x := bounds.Min.X; x < bounds.Max.X; x++ {
|
||||
r, g, b, _ := img.At(x, y).RGBA()
|
||||
if r != 0xffff || g != 0xffff || b != 0xffff {
|
||||
nonWhite++
|
||||
}
|
||||
}
|
||||
}
|
||||
if nonWhite == 0 {
|
||||
t.Fatal("rendered PDF page is blank")
|
||||
}
|
||||
|
||||
textOutput, err := exec.Command("pdftotext", pdfPath, "-").CombinedOutput()
|
||||
if err != nil {
|
||||
t.Fatalf("pdftotext: %v: %s", err, textOutput)
|
||||
}
|
||||
if !bytes.Contains(textOutput, []byte("Visible PDF body")) {
|
||||
t.Fatalf("pdftotext output = %q, want body text", textOutput)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolvePDFLatinFontPathHonorsEnv(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
fontPath := filepath.Join(dir, "custom.ttf")
|
||||
if err := os.WriteFile(fontPath, []byte("placeholder"), 0o600); err != nil {
|
||||
t.Fatalf("WriteFile font: %v", err)
|
||||
}
|
||||
t.Setenv("RAGFLOW_PDF_LATIN_FONT_PATH", fontPath)
|
||||
if got := resolvePDFLatinFontPath(); got != fontPath {
|
||||
t.Fatalf("resolvePDFLatinFontPath() = %q, want %q", got, fontPath)
|
||||
}
|
||||
}
|
||||
|
||||
func requirePDFLatinFont(t *testing.T) {
|
||||
t.Helper()
|
||||
if resolvePDFLatinFontPath() == "" {
|
||||
t.Skip("no local Latin PDF font available")
|
||||
}
|
||||
}
|
||||
|
||||
func min(a, b int) int {
|
||||
if a < b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
@@ -190,9 +190,15 @@ func (m *MessageComponent) Invoke(ctx context.Context, inputs map[string]any) (m
|
||||
|
||||
// Extract downloads. Walks inputs for download-info maps so
|
||||
// callers can attach binaries to the message body.
|
||||
var downloads []DownloadInfo
|
||||
for _, v := range inputs {
|
||||
downloads = append(downloads, ExtractDownloads(v)...)
|
||||
downloads := ExtractDownloads(resolved)
|
||||
if len(downloads) > 0 && downloadInfoString(resolved) {
|
||||
resolved = ""
|
||||
}
|
||||
for key, v := range inputs {
|
||||
if key == "text" {
|
||||
continue
|
||||
}
|
||||
downloads = appendUniqueDownloads(downloads, ExtractDownloads(v))
|
||||
}
|
||||
|
||||
// Pick the effective output format. inputs["output_format"]
|
||||
@@ -203,11 +209,13 @@ func (m *MessageComponent) Invoke(ctx context.Context, inputs map[string]any) (m
|
||||
format = OutputFormat(v)
|
||||
}
|
||||
|
||||
rendered := Render(RenderRequest{
|
||||
Format: format,
|
||||
Text: resolved,
|
||||
Downloads: downloads,
|
||||
})
|
||||
rendered := ""
|
||||
if resolved != "" {
|
||||
rendered = Render(RenderRequest{
|
||||
Format: format,
|
||||
Text: resolved,
|
||||
})
|
||||
}
|
||||
|
||||
out := map[string]any{"content": rendered}
|
||||
if len(downloads) > 0 {
|
||||
|
||||
@@ -121,6 +121,31 @@ func TestMessage_DownloadsExtraction(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessage_DownloadJSONStringSuppressesContent(t *testing.T) {
|
||||
c, _ := NewMessageComponent(map[string]any{"text": "unused"})
|
||||
state := canvas.NewCanvasState("r1", "t1")
|
||||
ctx := withStateForTest(context.Background(), state)
|
||||
downloadJSON := `{"doc_id":"d-1","filename":"report.md","mime_type":"text/markdown","url":"/api/v1/agents/attachments/d-1/download","include_download_info_in_content":true}`
|
||||
|
||||
out, err := c.Invoke(ctx, map[string]any{
|
||||
"text": downloadJSON,
|
||||
"stream": false,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
dls, ok := out["downloads"].([]DownloadInfo)
|
||||
if !ok {
|
||||
t.Fatalf("downloads key missing or wrong type: %T", out["downloads"])
|
||||
}
|
||||
if len(dls) != 1 || dls[0].DocID != "d-1" || dls[0].Filename != "report.md" {
|
||||
t.Fatalf("unexpected downloads: %+v", dls)
|
||||
}
|
||||
if got, _ := out["content"].(string); got != "" {
|
||||
t.Fatalf("content = %q, want empty download-only message", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMessage_AutoPlay_NoEngine: when auto_play is enabled but no
|
||||
// real synthesizer is registered, the textual content is still
|
||||
// returned and outputs["audio_error"] surfaces the deferred
|
||||
|
||||
@@ -40,6 +40,7 @@
|
||||
package component
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"html"
|
||||
"regexp"
|
||||
"strings"
|
||||
@@ -64,11 +65,14 @@ const (
|
||||
// entry. Mirrors agent/component/message.py:_is_download_info
|
||||
// (the {doc_id, filename, mime_type} tuple).
|
||||
type DownloadInfo struct {
|
||||
DocID string `json:"doc_id"`
|
||||
Filename string `json:"filename"`
|
||||
MimeType string `json:"mime_type"`
|
||||
URL string `json:"url,omitempty"`
|
||||
Content string `json:"content,omitempty"`
|
||||
DocID string `json:"doc_id"`
|
||||
Filename string `json:"filename"`
|
||||
MimeType string `json:"mime_type"`
|
||||
URL string `json:"url,omitempty"`
|
||||
Content string `json:"content,omitempty"`
|
||||
Size int `json:"size,omitempty"`
|
||||
PreviewURL string `json:"preview_url,omitempty"`
|
||||
IncludeDownloadInfoInContent bool `json:"include_download_info_in_content,omitempty"`
|
||||
}
|
||||
|
||||
// RenderRequest is the renderer input. Text is the resolved
|
||||
@@ -206,16 +210,26 @@ func ExtractDownloads(value any) []DownloadInfo {
|
||||
case nil:
|
||||
return nil
|
||||
case string:
|
||||
// A plain string is not a download descriptor. (Python also
|
||||
// tries json.loads here; we deliberately do not — the
|
||||
// message component resolves templates into strings, and
|
||||
// download descriptors are passed through as maps.)
|
||||
return nil
|
||||
var parsed any
|
||||
if err := json.Unmarshal([]byte(v), &parsed); err != nil {
|
||||
return nil
|
||||
}
|
||||
return ExtractDownloads(parsed)
|
||||
case map[string]any:
|
||||
if IsDownloadInfo(v) {
|
||||
return []DownloadInfo{downloadFromMap(v)}
|
||||
}
|
||||
return nil
|
||||
var out []DownloadInfo
|
||||
for _, item := range v {
|
||||
out = append(out, ExtractDownloads(item)...)
|
||||
}
|
||||
return out
|
||||
case map[string]string:
|
||||
m := make(map[string]any, len(v))
|
||||
for key, value := range v {
|
||||
m[key] = value
|
||||
}
|
||||
return ExtractDownloads(m)
|
||||
case []any:
|
||||
var out []DownloadInfo
|
||||
for _, item := range v {
|
||||
@@ -228,6 +242,56 @@ func ExtractDownloads(value any) []DownloadInfo {
|
||||
return nil
|
||||
}
|
||||
|
||||
func appendUniqueDownloads(dst []DownloadInfo, src []DownloadInfo) []DownloadInfo {
|
||||
for _, candidate := range src {
|
||||
duplicate := false
|
||||
for _, existing := range dst {
|
||||
if candidate.DocID != "" && candidate.DocID == existing.DocID {
|
||||
duplicate = true
|
||||
break
|
||||
}
|
||||
if candidate.DocID == "" && candidate.Filename == existing.Filename && candidate.URL == existing.URL {
|
||||
duplicate = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !duplicate {
|
||||
dst = append(dst, candidate)
|
||||
}
|
||||
}
|
||||
return dst
|
||||
}
|
||||
|
||||
func downloadInfoString(value any) bool {
|
||||
switch v := value.(type) {
|
||||
case string:
|
||||
var parsed any
|
||||
if err := json.Unmarshal([]byte(v), &parsed); err != nil {
|
||||
return false
|
||||
}
|
||||
return isDownloadInfoValue(parsed)
|
||||
default:
|
||||
return isDownloadInfoValue(v)
|
||||
}
|
||||
}
|
||||
|
||||
func isDownloadInfoValue(value any) bool {
|
||||
switch v := value.(type) {
|
||||
case map[string]any:
|
||||
return IsDownloadInfo(v)
|
||||
case map[string]string:
|
||||
for _, k := range []string{"doc_id", "filename", "mime_type"} {
|
||||
if _, ok := v[k]; !ok {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
case DownloadInfo:
|
||||
return v.DocID != "" && v.Filename != "" && v.MimeType != ""
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func downloadFromMap(m map[string]any) DownloadInfo {
|
||||
d := DownloadInfo{}
|
||||
if s, ok := m["doc_id"].(string); ok {
|
||||
@@ -242,8 +306,27 @@ func downloadFromMap(m map[string]any) DownloadInfo {
|
||||
if s, ok := m["url"].(string); ok {
|
||||
d.URL = s
|
||||
}
|
||||
if d.URL == "" {
|
||||
if s, ok := m["download"].(string); ok {
|
||||
d.URL = s
|
||||
}
|
||||
}
|
||||
if s, ok := m["preview_url"].(string); ok {
|
||||
d.PreviewURL = s
|
||||
if d.URL == "" {
|
||||
d.URL = s
|
||||
}
|
||||
}
|
||||
if s, ok := m["content"].(string); ok {
|
||||
d.Content = s
|
||||
}
|
||||
if f, ok := m["size"].(float64); ok {
|
||||
d.Size = int(f)
|
||||
} else if i, ok := m["size"].(int); ok {
|
||||
d.Size = i
|
||||
}
|
||||
if b, ok := m["include_download_info_in_content"].(bool); ok {
|
||||
d.IncludeDownloadInfoInContent = b
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
@@ -169,6 +169,23 @@ func TestExtractDownloads_Empty(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractDownloads_JSONString(t *testing.T) {
|
||||
value := `{"doc_id":"d1","filename":"report.md","mime_type":"text/markdown","download":"/api/v1/agents/attachments/d1/download","preview_url":"/api/v1/agents/attachments/d1/preview"}`
|
||||
dls := ExtractDownloads(value)
|
||||
if len(dls) != 1 {
|
||||
t.Fatalf("expected 1 download, got %d", len(dls))
|
||||
}
|
||||
if dls[0].DocID != "d1" || dls[0].Filename != "report.md" {
|
||||
t.Fatalf("unexpected download: %+v", dls[0])
|
||||
}
|
||||
if dls[0].URL != "/api/v1/agents/attachments/d1/download" {
|
||||
t.Fatalf("URL = %q, want download URL", dls[0].URL)
|
||||
}
|
||||
if dls[0].PreviewURL != "/api/v1/agents/attachments/d1/preview" {
|
||||
t.Fatalf("PreviewURL = %q, want preview URL", dls[0].PreviewURL)
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractDownloads_NestedList: nested list-of-lists is
|
||||
// flattened (the recursive walk).
|
||||
func TestExtractDownloads_NestedList(t *testing.T) {
|
||||
|
||||
@@ -1116,6 +1116,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv
|
||||
// statePost wrappers in scheduler.go.
|
||||
var answer string
|
||||
var reference []interface{}
|
||||
var downloads any
|
||||
now := float64(time.Now().UnixNano()) / 1e9
|
||||
for _, bucket := range state.Snapshot() {
|
||||
if v, ok := bucket["answer"].(string); ok && v != "" {
|
||||
@@ -1132,6 +1133,9 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv
|
||||
if v, ok := bucket["reference"].([]interface{}); ok {
|
||||
reference = append(reference, v...)
|
||||
}
|
||||
if v, ok := bucket["downloads"]; ok && !emptyDownloadValue(v) {
|
||||
downloads = v
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
@@ -1174,7 +1178,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv
|
||||
|
||||
wfPayload := map[string]interface{}{
|
||||
"inputs": map[string]any{"query": userInput},
|
||||
"outputs": answer,
|
||||
"outputs": workflowOutputs(answer, downloads),
|
||||
"elapsed_time": now - startedAt,
|
||||
"created_at": now,
|
||||
}
|
||||
@@ -1208,7 +1212,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv
|
||||
// per-run token usage across all LLM calls in this turn.
|
||||
wfPayload := map[string]interface{}{
|
||||
"inputs": map[string]any{"query": userInput},
|
||||
"outputs": answer,
|
||||
"outputs": workflowOutputs(answer, downloads),
|
||||
"elapsed_time": now - startedAt,
|
||||
"created_at": now,
|
||||
}
|
||||
@@ -1234,6 +1238,29 @@ func runIDFor(canvasID string, root map[string]any) string {
|
||||
return canvasID
|
||||
}
|
||||
|
||||
func workflowOutputs(content string, downloads any) any {
|
||||
if emptyDownloadValue(downloads) {
|
||||
return content
|
||||
}
|
||||
return map[string]any{
|
||||
"content": content,
|
||||
"downloads": downloads,
|
||||
}
|
||||
}
|
||||
|
||||
func emptyDownloadValue(value any) bool {
|
||||
if value == nil {
|
||||
return true
|
||||
}
|
||||
v := reflect.ValueOf(value)
|
||||
switch v.Kind() {
|
||||
case reflect.Slice, reflect.Array, reflect.Map:
|
||||
return v.Len() == 0
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AgentService) persistAgentRunSession(agentID, sessionID, messageID string, userInput any, answer string, reference []interface{}) {
|
||||
if sessionID == "" || s == nil || s.api4ConversationDAO == nil || dao.DB == nil {
|
||||
return
|
||||
|
||||
@@ -21,6 +21,7 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/netip"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -146,6 +147,33 @@ func TestListVersions_Empty(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkflowOutputs_WithDownloads(t *testing.T) {
|
||||
downloads := []map[string]any{
|
||||
{
|
||||
"doc_id": "d1",
|
||||
"filename": "report.pdf",
|
||||
"mime_type": "application/pdf",
|
||||
},
|
||||
}
|
||||
|
||||
out, ok := workflowOutputs("", downloads).(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("workflowOutputs type = %T, want map", workflowOutputs("", downloads))
|
||||
}
|
||||
if out["content"] != "" {
|
||||
t.Fatalf("content = %v, want empty string", out["content"])
|
||||
}
|
||||
if got := out["downloads"]; !reflect.DeepEqual(got, downloads) {
|
||||
t.Fatalf("downloads = %#v, want %#v", got, downloads)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkflowOutputs_NoDownloadsKeepsString(t *testing.T) {
|
||||
if got := workflowOutputs("answer", nil); got != "answer" {
|
||||
t.Fatalf("workflowOutputs without downloads = %#v, want string answer", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetVersion_Success verifies getting a specific version by ID.
|
||||
func TestGetVersion_Success(t *testing.T) {
|
||||
testDB := setupServiceTestDB(t)
|
||||
|
||||
Reference in New Issue
Block a user