Enhance PDF and Content Processing with New File Handling Features
- Added support for parsing various file types (PDF, DOCX, PPTX) in the content processing pipeline, allowing for more flexible content extraction. - Implemented a new method to convert file attachments to raw text when content parsing is skipped, improving performance for internal calls. - Introduced loading message suppression for image processing to enhance user experience during PDF analysis. - Updated the PDF handler to cache processed text and manage loading messages effectively, ensuring smoother interactions during content retrieval. - Enhanced error handling and logging for PDF processing, improving traceability and debugging capabilities.
This commit is contained in:
Max
parent
983dbd3cee
commit
195391afaa
17 changed files with 2152 additions and 26 deletions
11
.github/workflows/pr-test.yml
vendored
11
.github/workflows/pr-test.yml
vendored
|
|
@ -524,6 +524,17 @@ jobs:
|
|||
- name: Setup Go Tools
|
||||
run: make tools
|
||||
|
||||
- name: Install pdftoppm, mutool, imagemagick
|
||||
run: |
|
||||
sudo apt update
|
||||
sudo apt install -y poppler-utils mupdf-tools imagemagick
|
||||
|
||||
- name: Test pdftoppm, mutool, imagemagick
|
||||
run: |
|
||||
pdftoppm -v
|
||||
mutool -v
|
||||
convert -version
|
||||
|
||||
- name: Setup ENV (SQLite)
|
||||
run: |
|
||||
mkdir -p ${{ github.WORKSPACE }}/../app/db
|
||||
|
|
|
|||
11
.github/workflows/unit-test.yml
vendored
11
.github/workflows/unit-test.yml
vendored
|
|
@ -421,6 +421,17 @@ jobs:
|
|||
- name: Setup Go Tools
|
||||
run: make tools
|
||||
|
||||
- name: Install pdftoppm, mutool, imagemagick
|
||||
run: |
|
||||
sudo apt update
|
||||
sudo apt install -y poppler-utils mupdf-tools imagemagick
|
||||
|
||||
- name: Test pdftoppm, mutool, imagemagick
|
||||
run: |
|
||||
pdftoppm -v
|
||||
mutool -v
|
||||
convert -version
|
||||
|
||||
- name: Setup ENV (SQLite)
|
||||
run: |
|
||||
mkdir -p ${{ github.WORKSPACE }}/../app/db
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import (
|
|||
"fmt"
|
||||
|
||||
"github.com/yaoapp/yao/agent/content"
|
||||
"github.com/yaoapp/yao/agent/content/text"
|
||||
contentTypes "github.com/yaoapp/yao/agent/content/types"
|
||||
"github.com/yaoapp/yao/agent/context"
|
||||
)
|
||||
|
|
@ -13,6 +14,12 @@ import (
|
|||
//
|
||||
// This should be called after BuildRequest and before executing LLM call
|
||||
func (ast *Assistant) BuildContent(ctx *context.Context, messages []context.Message, options *context.CompletionOptions, opts *context.Options) ([]context.Message, error) {
|
||||
// Skip complex content parsing if requested (for internal calls like needsearch)
|
||||
// Still convert file attachments to raw text
|
||||
if opts != nil && opts.Skip != nil && opts.Skip.ContentParsing {
|
||||
return convertFilesToText(ctx, messages), nil
|
||||
}
|
||||
|
||||
// Set AssistantID in context for file info tracking in Space
|
||||
// This ensures hooks can access file information using the correct namespace
|
||||
if ctx.AssistantID == "" {
|
||||
|
|
@ -45,3 +52,102 @@ func (ast *Assistant) BuildContent(ctx *context.Context, messages []context.Mess
|
|||
|
||||
return contentMessages, nil
|
||||
}
|
||||
|
||||
// convertFilesToText converts file attachments in messages to raw text
|
||||
// Used when SkipContentParsing is enabled - simple text extraction without vision/PDF processing
|
||||
func convertFilesToText(ctx *context.Context, messages []context.Message) []context.Message {
|
||||
result := make([]context.Message, 0, len(messages))
|
||||
textHandler := text.New(nil)
|
||||
|
||||
for _, msg := range messages {
|
||||
// Only process user messages
|
||||
if msg.Role != context.RoleUser {
|
||||
result = append(result, msg)
|
||||
continue
|
||||
}
|
||||
|
||||
// Handle content parts
|
||||
parts, ok := msg.Content.([]context.ContentPart)
|
||||
if !ok {
|
||||
// Try []interface{} (from history/JSON)
|
||||
if iparts, ok := msg.Content.([]interface{}); ok {
|
||||
parts = convertInterfaceToParts(iparts)
|
||||
}
|
||||
}
|
||||
|
||||
if len(parts) == 0 {
|
||||
result = append(result, msg)
|
||||
continue
|
||||
}
|
||||
|
||||
// Convert file parts to text
|
||||
newParts := make([]context.ContentPart, 0, len(parts))
|
||||
for _, part := range parts {
|
||||
switch part.Type {
|
||||
case context.ContentFile:
|
||||
// Convert file to raw text
|
||||
if part.File != nil && part.File.URL != "" {
|
||||
textPart, _, err := textHandler.ParseRaw(ctx, part)
|
||||
if err == nil {
|
||||
newParts = append(newParts, textPart)
|
||||
continue
|
||||
}
|
||||
}
|
||||
newParts = append(newParts, part)
|
||||
|
||||
case context.ContentImageURL:
|
||||
// Skip images - cannot convert to text without vision
|
||||
continue
|
||||
|
||||
default:
|
||||
newParts = append(newParts, part)
|
||||
}
|
||||
}
|
||||
|
||||
newMsg := msg
|
||||
newMsg.Content = newParts
|
||||
result = append(result, newMsg)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// convertInterfaceToParts converts []interface{} to []ContentPart for file extraction
|
||||
func convertInterfaceToParts(items []interface{}) []context.ContentPart {
|
||||
parts := make([]context.ContentPart, 0, len(items))
|
||||
for _, item := range items {
|
||||
m, ok := item.(map[string]interface{})
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
typeStr, _ := m["type"].(string)
|
||||
part := context.ContentPart{
|
||||
Type: context.ContentPartType(typeStr),
|
||||
}
|
||||
|
||||
switch typeStr {
|
||||
case "text":
|
||||
if t, ok := m["text"].(string); ok {
|
||||
part.Text = t
|
||||
}
|
||||
case "file":
|
||||
if fileData, ok := m["file"].(map[string]interface{}); ok {
|
||||
part.File = &context.FileAttachment{}
|
||||
if url, ok := fileData["url"].(string); ok {
|
||||
part.File.URL = url
|
||||
}
|
||||
if filename, ok := fileData["filename"].(string); ok {
|
||||
part.File.Filename = filename
|
||||
}
|
||||
}
|
||||
case "image_url":
|
||||
part.Type = context.ContentImageURL
|
||||
default:
|
||||
continue
|
||||
}
|
||||
|
||||
parts = append(parts, part)
|
||||
}
|
||||
return parts
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,9 +2,14 @@ package content
|
|||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/yaoapp/kun/log"
|
||||
"github.com/yaoapp/yao/agent/content/docx"
|
||||
"github.com/yaoapp/yao/agent/content/image"
|
||||
"github.com/yaoapp/yao/agent/content/pdf"
|
||||
"github.com/yaoapp/yao/agent/content/pptx"
|
||||
"github.com/yaoapp/yao/agent/content/text"
|
||||
"github.com/yaoapp/yao/agent/content/types"
|
||||
agentContext "github.com/yaoapp/yao/agent/context"
|
||||
searchTypes "github.com/yaoapp/yao/agent/search/types"
|
||||
|
|
@ -99,7 +104,7 @@ func parseContentPart(ctx *agentContext.Context, content agentContext.ContentPar
|
|||
return content, nil, nil
|
||||
|
||||
case agentContext.ContentFile:
|
||||
return content, nil, nil
|
||||
return parseFileContent(ctx, content, options)
|
||||
|
||||
case agentContext.ContentData:
|
||||
return content, nil, nil
|
||||
|
|
@ -109,6 +114,35 @@ func parseContentPart(ctx *agentContext.Context, content agentContext.ContentPar
|
|||
}
|
||||
}
|
||||
|
||||
// parseFileContent parses file content based on file type
|
||||
func parseFileContent(ctx *agentContext.Context, content agentContext.ContentPart, options *types.Options) (agentContext.ContentPart, []*searchTypes.Reference, error) {
|
||||
if content.File == nil || content.File.URL == "" {
|
||||
return content, nil, nil
|
||||
}
|
||||
|
||||
// Determine file type from filename
|
||||
filename := strings.ToLower(content.File.Filename)
|
||||
|
||||
// Check file type and route to appropriate handler
|
||||
switch {
|
||||
case strings.HasSuffix(filename, ".pdf"):
|
||||
return pdf.New(options).Parse(ctx, content)
|
||||
|
||||
case strings.HasSuffix(filename, ".docx"):
|
||||
return docx.New(options).Parse(ctx, content)
|
||||
|
||||
case strings.HasSuffix(filename, ".pptx"):
|
||||
return pptx.New(options).Parse(ctx, content)
|
||||
|
||||
case text.IsSupportedExtension(filename):
|
||||
return text.New(options).Parse(ctx, content)
|
||||
}
|
||||
|
||||
// For unsupported file types, try to read as text
|
||||
// This allows any file to be converted to text content
|
||||
return text.New(options).ParseRaw(ctx, content)
|
||||
}
|
||||
|
||||
// convertToContentParts converts []interface{} to []ContentPart
|
||||
// This is needed when content is loaded from JSON/history and is []interface{} instead of []ContentPart
|
||||
func convertToContentParts(content []interface{}) ([]agentContext.ContentPart, bool) {
|
||||
|
|
|
|||
143
agent/content/docx/docx.go
Normal file
143
agent/content/docx/docx.go
Normal file
|
|
@ -0,0 +1,143 @@
|
|||
package docx
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/yaoapp/gou/office"
|
||||
"github.com/yaoapp/yao/agent/content/types"
|
||||
agentContext "github.com/yaoapp/yao/agent/context"
|
||||
searchTypes "github.com/yaoapp/yao/agent/search/types"
|
||||
"github.com/yaoapp/yao/attachment"
|
||||
)
|
||||
|
||||
// Docx handles DOCX content
|
||||
type Docx struct {
|
||||
options *types.Options
|
||||
}
|
||||
|
||||
// New creates a new DOCX handler
|
||||
func New(options *types.Options) *Docx {
|
||||
return &Docx{options: options}
|
||||
}
|
||||
|
||||
// Parse parses DOCX content and returns text
|
||||
func (h *Docx) Parse(ctx *agentContext.Context, content agentContext.ContentPart) (agentContext.ContentPart, []*searchTypes.Reference, error) {
|
||||
if content.File == nil || content.File.URL == "" {
|
||||
return content, nil, fmt.Errorf("file content missing URL")
|
||||
}
|
||||
|
||||
url := content.File.URL
|
||||
|
||||
// Check cache first
|
||||
cachedText, found, err := h.readFromCache(ctx, url)
|
||||
if err == nil && found {
|
||||
return agentContext.ContentPart{
|
||||
Type: agentContext.ContentText,
|
||||
Text: cachedText,
|
||||
}, nil, nil
|
||||
}
|
||||
|
||||
// Read DOCX file
|
||||
data, err := h.readFile(ctx, url)
|
||||
if err != nil {
|
||||
return content, nil, fmt.Errorf("failed to read DOCX: %w", err)
|
||||
}
|
||||
|
||||
// Parse DOCX using gou/office
|
||||
parser := office.NewParser()
|
||||
result, err := parser.Parse(data)
|
||||
if err != nil {
|
||||
return content, nil, fmt.Errorf("failed to parse DOCX: %w", err)
|
||||
}
|
||||
|
||||
text := result.Markdown
|
||||
if text == "" {
|
||||
return content, nil, fmt.Errorf("no text content extracted from DOCX")
|
||||
}
|
||||
|
||||
// Cache the result
|
||||
if err := h.saveToCache(ctx, url, text); err != nil {
|
||||
// Log warning but don't fail
|
||||
fmt.Printf("Warning: failed to cache DOCX text: %v\n", err)
|
||||
}
|
||||
|
||||
return agentContext.ContentPart{
|
||||
Type: agentContext.ContentText,
|
||||
Text: text,
|
||||
}, nil, nil
|
||||
}
|
||||
|
||||
// readFile reads DOCX content from various sources
|
||||
func (h *Docx) readFile(ctx *agentContext.Context, url string) ([]byte, error) {
|
||||
if strings.HasPrefix(url, "__") {
|
||||
return h.readFromUploader(ctx, url)
|
||||
}
|
||||
|
||||
if strings.HasPrefix(url, "http://") || strings.HasPrefix(url, "https://") {
|
||||
return nil, fmt.Errorf("HTTP URL fetch not implemented yet: %s", url)
|
||||
}
|
||||
|
||||
// Try to read as local file path
|
||||
if _, err := os.Stat(url); err == nil {
|
||||
return os.ReadFile(url)
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("unsupported DOCX source: %s", url)
|
||||
}
|
||||
|
||||
// readFromUploader reads DOCX content from file uploader
|
||||
func (h *Docx) readFromUploader(ctx *agentContext.Context, wrapper string) ([]byte, error) {
|
||||
uploaderName, fileID, ok := attachment.Parse(wrapper)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid uploader wrapper format: %s", wrapper)
|
||||
}
|
||||
|
||||
manager, exists := attachment.Managers[uploaderName]
|
||||
if !exists {
|
||||
return nil, fmt.Errorf("uploader '%s' not found", uploaderName)
|
||||
}
|
||||
|
||||
data, err := manager.Read(ctx.Context, fileID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read file: %w", err)
|
||||
}
|
||||
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// readFromCache reads cached text content for a DOCX
|
||||
func (h *Docx) readFromCache(ctx *agentContext.Context, url string) (string, bool, error) {
|
||||
uploaderName, fileID, isWrapper := attachment.Parse(url)
|
||||
if !isWrapper {
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
manager, exists := attachment.Managers[uploaderName]
|
||||
if !exists {
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
text, err := manager.GetText(ctx.Context, fileID, false)
|
||||
if err == nil && text != "" {
|
||||
return text, true, nil
|
||||
}
|
||||
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
// saveToCache saves processed text to cache
|
||||
func (h *Docx) saveToCache(ctx *agentContext.Context, url string, text string) error {
|
||||
uploaderName, fileID, isWrapper := attachment.Parse(url)
|
||||
if !isWrapper {
|
||||
return nil
|
||||
}
|
||||
|
||||
manager, exists := attachment.Managers[uploaderName]
|
||||
if !exists {
|
||||
return nil
|
||||
}
|
||||
|
||||
return manager.SaveText(ctx.Context, fileID, text)
|
||||
}
|
||||
132
agent/content/docx/docx_test.go
Normal file
132
agent/content/docx/docx_test.go
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
package docx_test
|
||||
|
||||
import (
|
||||
stdContext "context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/yaoapp/gou/connector/openai"
|
||||
"github.com/yaoapp/yao/agent/content/docx"
|
||||
contentTypes "github.com/yaoapp/yao/agent/content/types"
|
||||
agentContext "github.com/yaoapp/yao/agent/context"
|
||||
"github.com/yaoapp/yao/agent/output/message"
|
||||
"github.com/yaoapp/yao/agent/testutils"
|
||||
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
|
||||
)
|
||||
|
||||
const testFilesDir = "assistants/tests/vision-helper/tests"
|
||||
|
||||
func newTestContext() *agentContext.Context {
|
||||
authorized := &oauthTypes.AuthorizedInfo{
|
||||
Subject: "test-user",
|
||||
ClientID: "test-client-id",
|
||||
UserID: "test-user-123",
|
||||
}
|
||||
ctx := agentContext.New(stdContext.Background(), authorized, "test-chat")
|
||||
ctx.AssistantID = "test-assistant"
|
||||
ctx.Locale = "en-us"
|
||||
ctx.IDGenerator = message.NewIDGenerator()
|
||||
return ctx
|
||||
}
|
||||
|
||||
func newTestOptions() *contentTypes.Options {
|
||||
return &contentTypes.Options{
|
||||
Capabilities: &openai.Capabilities{},
|
||||
}
|
||||
}
|
||||
|
||||
func getTestFilePath(filename string) string {
|
||||
yaoRoot := os.Getenv("YAO_TEST_APPLICATION")
|
||||
if yaoRoot == "" {
|
||||
yaoRoot = os.Getenv("YAO_ROOT")
|
||||
}
|
||||
return filepath.Join(yaoRoot, testFilesDir, filename)
|
||||
}
|
||||
|
||||
// TestParseWithMissingURL tests parsing DOCX with missing URL
|
||||
func TestParseWithMissingURL(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
options := newTestOptions()
|
||||
ctx := newTestContext()
|
||||
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: nil,
|
||||
}
|
||||
|
||||
handler := docx.New(options)
|
||||
_, _, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "missing URL")
|
||||
}
|
||||
|
||||
// TestParseWithLocalDocx tests parsing a local DOCX file
|
||||
func TestParseWithLocalDocx(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
docxPath := getTestFilePath("docx.docx")
|
||||
if _, err := os.Stat(docxPath); os.IsNotExist(err) {
|
||||
t.Skipf("Test DOCX file not found: %s", docxPath)
|
||||
}
|
||||
|
||||
options := newTestOptions()
|
||||
ctx := newTestContext()
|
||||
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: docxPath,
|
||||
Filename: "docx.docx",
|
||||
},
|
||||
}
|
||||
|
||||
handler := docx.New(options)
|
||||
result, refs, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Nil(t, refs)
|
||||
assert.Equal(t, agentContext.ContentText, result.Type)
|
||||
assert.NotEmpty(t, result.Text)
|
||||
t.Logf("DOCX parse result (first 500 chars): %.500s...", result.Text)
|
||||
}
|
||||
|
||||
// TestParseWithNonExistentFile tests parsing DOCX with non-existent file
|
||||
func TestParseWithNonExistentFile(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
options := newTestOptions()
|
||||
ctx := newTestContext()
|
||||
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: "/non/existent/path/test.docx",
|
||||
Filename: "test.docx",
|
||||
},
|
||||
}
|
||||
|
||||
handler := docx.New(options)
|
||||
_, _, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "unsupported DOCX source")
|
||||
}
|
||||
|
|
@ -361,7 +361,13 @@ func (h *Image) callMCPVisionTool(ctx *agentContext.Context, serverID string, co
|
|||
}
|
||||
|
||||
// sendLoading sends a loading message and returns the message ID
|
||||
// Returns empty string if SilentLoading is enabled
|
||||
func (h *Image) sendLoading(ctx *agentContext.Context, msg string) string {
|
||||
// Skip loading message if SilentLoading is enabled (called from parent handler like PDF)
|
||||
if h.options != nil && h.options.SilentLoading {
|
||||
return ""
|
||||
}
|
||||
|
||||
loadingMsg := &message.Message{
|
||||
Type: message.TypeLoading,
|
||||
Props: map[string]interface{}{
|
||||
|
|
|
|||
|
|
@ -2,48 +2,375 @@ package pdf
|
|||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
goupdf "github.com/yaoapp/gou/pdf"
|
||||
"github.com/yaoapp/yao/agent/content/image"
|
||||
"github.com/yaoapp/yao/agent/content/types"
|
||||
agentContext "github.com/yaoapp/yao/agent/context"
|
||||
"github.com/yaoapp/yao/agent/i18n"
|
||||
"github.com/yaoapp/yao/agent/output/message"
|
||||
searchTypes "github.com/yaoapp/yao/agent/search/types"
|
||||
"github.com/yaoapp/yao/attachment"
|
||||
kbTypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
// PDF handlePDF content
|
||||
// PDF handles PDF content
|
||||
type PDF struct {
|
||||
options *types.Options
|
||||
}
|
||||
|
||||
// New creates a new image handler
|
||||
// New creates a new PDF handler
|
||||
func New(options *types.Options) *PDF {
|
||||
return &PDF{options: options}
|
||||
}
|
||||
|
||||
// Parse parses pdf content
|
||||
// Parse parses PDF content by converting to images and processing each page
|
||||
// Returns multiple ContentPart (one text part per page) combined into a single text part
|
||||
func (h *PDF) Parse(ctx *agentContext.Context, content agentContext.ContentPart) (agentContext.ContentPart, []*searchTypes.Reference, error) {
|
||||
return content, nil, nil
|
||||
if content.File == nil || content.File.URL == "" {
|
||||
return content, nil, fmt.Errorf("file content missing URL")
|
||||
}
|
||||
|
||||
url := content.File.URL
|
||||
|
||||
// Check cache first
|
||||
cachedText, found, err := h.readFromCache(ctx, url)
|
||||
if err == nil && found {
|
||||
return agentContext.ContentPart{
|
||||
Type: agentContext.ContentText,
|
||||
Text: cachedText,
|
||||
}, nil, nil
|
||||
}
|
||||
|
||||
// Convert PDF to images and process each page
|
||||
return h.asImages(ctx, content)
|
||||
}
|
||||
|
||||
// asImages converts pdf content to images then parse as image content
|
||||
// ParseMulti parses PDF content and returns multiple ContentParts (one per page)
|
||||
// This is useful when you need separate parts for each page
|
||||
func (h *PDF) ParseMulti(ctx *agentContext.Context, content agentContext.ContentPart) ([]agentContext.ContentPart, []*searchTypes.Reference, error) {
|
||||
if content.File == nil || content.File.URL == "" {
|
||||
return nil, nil, fmt.Errorf("file content missing URL")
|
||||
}
|
||||
|
||||
url := content.File.URL
|
||||
|
||||
// Check cache first - if cached, return as single text part
|
||||
cachedText, found, err := h.readFromCache(ctx, url)
|
||||
if err == nil && found {
|
||||
return []agentContext.ContentPart{
|
||||
{
|
||||
Type: agentContext.ContentText,
|
||||
Text: cachedText,
|
||||
},
|
||||
}, nil, nil
|
||||
}
|
||||
|
||||
// Convert PDF to images and process each page
|
||||
return h.asImagesMulti(ctx, content)
|
||||
}
|
||||
|
||||
// asImages converts PDF to images and processes each page, returning combined result
|
||||
func (h *PDF) asImages(ctx *agentContext.Context, content agentContext.ContentPart) (agentContext.ContentPart, []*searchTypes.Reference, error) {
|
||||
return agentContext.ContentPart{}, nil, nil
|
||||
parts, refs, err := h.asImagesMulti(ctx, content)
|
||||
if err != nil {
|
||||
return content, nil, err
|
||||
}
|
||||
|
||||
if len(parts) == 0 {
|
||||
return content, nil, fmt.Errorf("no pages extracted from PDF")
|
||||
}
|
||||
|
||||
// Check if any parts are text (vision agent was used) or image_url (model supports vision)
|
||||
hasTextParts := false
|
||||
hasImageParts := false
|
||||
for _, part := range parts {
|
||||
if part.Type == agentContext.ContentText {
|
||||
hasTextParts = true
|
||||
} else if part.Type == agentContext.ContentImageURL {
|
||||
hasImageParts = true
|
||||
}
|
||||
}
|
||||
|
||||
// If all parts are image_url (model supports vision), return the first image
|
||||
// The caller should use ParseMulti to get all images
|
||||
if hasImageParts && !hasTextParts {
|
||||
return parts[0], refs, nil
|
||||
}
|
||||
|
||||
// Combine all text parts into one
|
||||
var combinedText strings.Builder
|
||||
pageNum := 0
|
||||
for _, part := range parts {
|
||||
if part.Type == agentContext.ContentText && part.Text != "" {
|
||||
pageNum++
|
||||
if pageNum > 1 {
|
||||
combinedText.WriteString("\n\n---\n\n") // Page separator
|
||||
}
|
||||
combinedText.WriteString(fmt.Sprintf("## Page %d\n\n", pageNum))
|
||||
combinedText.WriteString(part.Text)
|
||||
}
|
||||
}
|
||||
|
||||
result := agentContext.ContentPart{
|
||||
Type: agentContext.ContentText,
|
||||
Text: combinedText.String(),
|
||||
}
|
||||
|
||||
// Cache the combined result
|
||||
if content.File != nil && content.File.URL != "" && combinedText.Len() > 0 {
|
||||
h.saveToCache(ctx, content.File.URL, combinedText.String())
|
||||
}
|
||||
|
||||
return result, refs, nil
|
||||
}
|
||||
|
||||
// base64 encodes image content to base64 ( for PDF support )
|
||||
func (h *PDF) base64(ctx *agentContext.Context, content agentContext.ContentPart) (agentContext.ContentPart, []*searchTypes.Reference, error) {
|
||||
return agentContext.ContentPart{}, nil, nil
|
||||
// asImagesMulti converts PDF to images and processes each page separately
|
||||
func (h *PDF) asImagesMulti(ctx *agentContext.Context, content agentContext.ContentPart) ([]agentContext.ContentPart, []*searchTypes.Reference, error) {
|
||||
if content.File == nil || content.File.URL == "" {
|
||||
return nil, nil, fmt.Errorf("file content missing URL")
|
||||
}
|
||||
|
||||
url := content.File.URL
|
||||
|
||||
// Read PDF file
|
||||
pdfData, err := h.readPDF(ctx, url)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("failed to read PDF: %w", err)
|
||||
}
|
||||
|
||||
// Create temporary file for PDF
|
||||
tempDir := os.TempDir()
|
||||
pdfPath := filepath.Join(tempDir, fmt.Sprintf("pdf_%d.pdf", time.Now().UnixNano()))
|
||||
if err := os.WriteFile(pdfPath, pdfData, 0644); err != nil {
|
||||
return nil, nil, fmt.Errorf("failed to write temp PDF: %w", err)
|
||||
}
|
||||
defer os.Remove(pdfPath)
|
||||
|
||||
// Get PDF processor with global config
|
||||
processor, err := h.getPDFProcessor()
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("failed to create PDF processor: %w", err)
|
||||
}
|
||||
|
||||
// Create output directory for images
|
||||
imagesDir := filepath.Join(tempDir, fmt.Sprintf("pdf_images_%d", time.Now().UnixNano()))
|
||||
if err := os.MkdirAll(imagesDir, 0755); err != nil {
|
||||
return nil, nil, fmt.Errorf("failed to create images directory: %w", err)
|
||||
}
|
||||
defer os.RemoveAll(imagesDir)
|
||||
|
||||
// Convert PDF to images
|
||||
convertConfig := goupdf.ConvertConfig{
|
||||
OutputDir: imagesDir,
|
||||
OutputPrefix: "page",
|
||||
Format: "png",
|
||||
DPI: 150,
|
||||
Quality: 90,
|
||||
PageRange: "all",
|
||||
}
|
||||
|
||||
imageFiles, err := processor.Convert(ctx.Context, pdfPath, convertConfig)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("failed to convert PDF to images: %w", err)
|
||||
}
|
||||
|
||||
if len(imageFiles) == 0 {
|
||||
return nil, nil, fmt.Errorf("no pages extracted from PDF")
|
||||
}
|
||||
|
||||
// Process each image using the image handler (with SilentLoading to suppress image loading messages)
|
||||
imageOptions := *h.options // Copy options
|
||||
imageOptions.SilentLoading = true
|
||||
imageHandler := image.New(&imageOptions)
|
||||
var parts []agentContext.ContentPart
|
||||
var allRefs []*searchTypes.Reference
|
||||
|
||||
for i, imageFile := range imageFiles {
|
||||
// Send loading message for this page
|
||||
loadingMsg := fmt.Sprintf(i18n.T(ctx.Locale, "content.pdf.analyzing_page"), i+1, len(imageFiles))
|
||||
loadingID := h.sendLoading(ctx, loadingMsg)
|
||||
|
||||
// Read image file
|
||||
imageData, err := os.ReadFile(imageFile)
|
||||
if err != nil {
|
||||
h.sendLoadingDone(ctx, loadingID)
|
||||
continue
|
||||
}
|
||||
|
||||
// Convert to base64 data URI
|
||||
base64Data := image.EncodeToBase64DataURI(imageData, "image/png")
|
||||
|
||||
// Create image content part
|
||||
imagePart := agentContext.ContentPart{
|
||||
Type: agentContext.ContentImageURL,
|
||||
ImageURL: &agentContext.ImageURL{
|
||||
URL: base64Data,
|
||||
Detail: agentContext.DetailAuto,
|
||||
},
|
||||
}
|
||||
|
||||
// Parse image using image handler
|
||||
parsedPart, refs, err := imageHandler.Parse(ctx, imagePart)
|
||||
|
||||
// Mark loading as done
|
||||
h.sendLoadingDone(ctx, loadingID)
|
||||
|
||||
if err != nil {
|
||||
// If parsing fails, skip this page
|
||||
continue
|
||||
}
|
||||
|
||||
parts = append(parts, parsedPart)
|
||||
if refs != nil {
|
||||
allRefs = append(allRefs, refs...)
|
||||
}
|
||||
}
|
||||
|
||||
if len(parts) == 0 {
|
||||
return nil, nil, fmt.Errorf("failed to process any PDF pages")
|
||||
}
|
||||
|
||||
return parts, allRefs, nil
|
||||
}
|
||||
|
||||
// read image content from file uploader __uploader://fileid return agentContext.ContentPart
|
||||
func (h *PDF) read(ctx *agentContext.Context, content agentContext.ContentPart) (agentContext.ContentPart, []*searchTypes.Reference, error) {
|
||||
return agentContext.ContentPart{}, nil, fmt.Errorf("not implemented")
|
||||
// readPDF reads PDF content from various sources
|
||||
func (h *PDF) readPDF(ctx *agentContext.Context, url string) ([]byte, error) {
|
||||
if strings.HasPrefix(url, "__") {
|
||||
// Uploader wrapper format: __uploader://fileid
|
||||
return h.readFromUploader(ctx, url)
|
||||
}
|
||||
|
||||
if strings.HasPrefix(url, "http://") || strings.HasPrefix(url, "https://") {
|
||||
return nil, fmt.Errorf("HTTP URL fetch not implemented yet: %s", url)
|
||||
}
|
||||
|
||||
// Try to read as local file path
|
||||
if _, err := os.Stat(url); err == nil {
|
||||
return os.ReadFile(url)
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("unsupported PDF source: %s", url)
|
||||
}
|
||||
|
||||
// read image content from file uploader __uploader://fileid return agentContext.ContentPart
|
||||
func (h *PDF) readFromCache(ctx *agentContext.Context, content agentContext.ContentPart) (agentContext.ContentPart, []*searchTypes.Reference, error) {
|
||||
return agentContext.ContentPart{}, nil, fmt.Errorf("not implemented")
|
||||
// readFromUploader reads PDF content from file uploader
|
||||
func (h *PDF) readFromUploader(ctx *agentContext.Context, wrapper string) ([]byte, error) {
|
||||
uploaderName, fileID, ok := attachment.Parse(wrapper)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid uploader wrapper format: %s", wrapper)
|
||||
}
|
||||
|
||||
manager, exists := attachment.Managers[uploaderName]
|
||||
if !exists {
|
||||
return nil, fmt.Errorf("uploader '%s' not found", uploaderName)
|
||||
}
|
||||
|
||||
data, err := manager.Read(ctx.Context, fileID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read file: %w", err)
|
||||
}
|
||||
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// read image content from file uploader __uploader://fileid return base64 encoded string
|
||||
func (h *PDF) readFromUploader(ctx *agentContext.Context, content agentContext.ContentPart) (string, error) {
|
||||
return "", fmt.Errorf("not implemented")
|
||||
// readFromCache reads cached text content for a PDF
|
||||
func (h *PDF) readFromCache(ctx *agentContext.Context, url string) (string, bool, error) {
|
||||
uploaderName, fileID, isWrapper := attachment.Parse(url)
|
||||
if !isWrapper {
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
manager, exists := attachment.Managers[uploaderName]
|
||||
if !exists {
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
text, err := manager.GetText(ctx.Context, fileID, false)
|
||||
if err == nil && text != "" {
|
||||
return text, true, nil
|
||||
}
|
||||
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
// saveToCache saves processed text to cache
|
||||
func (h *PDF) saveToCache(ctx *agentContext.Context, url string, text string) error {
|
||||
uploaderName, fileID, isWrapper := attachment.Parse(url)
|
||||
if !isWrapper {
|
||||
return nil
|
||||
}
|
||||
|
||||
manager, exists := attachment.Managers[uploaderName]
|
||||
if !exists {
|
||||
return nil
|
||||
}
|
||||
|
||||
return manager.SaveText(ctx.Context, fileID, text)
|
||||
}
|
||||
|
||||
// getPDFProcessor creates a PDF processor using global KB config
|
||||
func (h *PDF) getPDFProcessor() (*goupdf.PDF, error) {
|
||||
globalPDF := kbTypes.GetGlobalPDF()
|
||||
|
||||
opts := goupdf.Options{
|
||||
ConvertTool: goupdf.ToolPdftoppm, // default
|
||||
ToolPath: "",
|
||||
}
|
||||
|
||||
if globalPDF != nil {
|
||||
if globalPDF.ConvertTool != "" {
|
||||
switch globalPDF.ConvertTool {
|
||||
case "pdftoppm":
|
||||
opts.ConvertTool = goupdf.ToolPdftoppm
|
||||
case "mutool":
|
||||
opts.ConvertTool = goupdf.ToolMutool
|
||||
case "imagemagick", "convert":
|
||||
opts.ConvertTool = goupdf.ToolImageMagick
|
||||
}
|
||||
}
|
||||
if globalPDF.ToolPath != "" {
|
||||
opts.ToolPath = globalPDF.ToolPath
|
||||
}
|
||||
}
|
||||
|
||||
return goupdf.New(opts), nil
|
||||
}
|
||||
|
||||
// sendLoading sends a loading message and returns the message ID
|
||||
func (h *PDF) sendLoading(ctx *agentContext.Context, msg string) string {
|
||||
loadingMsg := &message.Message{
|
||||
Type: message.TypeLoading,
|
||||
Props: map[string]interface{}{
|
||||
"message": msg,
|
||||
},
|
||||
}
|
||||
|
||||
msgID, err := ctx.SendStream(loadingMsg)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return msgID
|
||||
}
|
||||
|
||||
// sendLoadingDone marks the loading message as done
|
||||
func (h *PDF) sendLoadingDone(ctx *agentContext.Context, loadingID string) {
|
||||
if loadingID == "" {
|
||||
return
|
||||
}
|
||||
|
||||
doneMsg := &message.Message{
|
||||
MessageID: loadingID,
|
||||
Delta: true,
|
||||
DeltaAction: message.DeltaReplace,
|
||||
Type: message.TypeLoading,
|
||||
Props: map[string]interface{}{
|
||||
"done": true,
|
||||
},
|
||||
}
|
||||
|
||||
ctx.Send(doneMsg)
|
||||
}
|
||||
|
|
|
|||
362
agent/content/pdf/pdf_test.go
Normal file
362
agent/content/pdf/pdf_test.go
Normal file
|
|
@ -0,0 +1,362 @@
|
|||
package pdf_test
|
||||
|
||||
import (
|
||||
stdContext "context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/yaoapp/gou/connector/openai"
|
||||
"github.com/yaoapp/yao/agent/content/pdf"
|
||||
contentTypes "github.com/yaoapp/yao/agent/content/types"
|
||||
agentContext "github.com/yaoapp/yao/agent/context"
|
||||
"github.com/yaoapp/yao/agent/output/message"
|
||||
"github.com/yaoapp/yao/agent/testutils"
|
||||
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
|
||||
)
|
||||
|
||||
// Test files directory (relative to yao-dev-app)
|
||||
const testFilesDir = "assistants/tests/vision-helper/tests"
|
||||
|
||||
// newTestContext creates a Context for testing with commonly used fields pre-populated
|
||||
func newTestContext(capabilities *openai.Capabilities) *agentContext.Context {
|
||||
authorized := &oauthTypes.AuthorizedInfo{
|
||||
Subject: "test-user",
|
||||
ClientID: "test-client-id",
|
||||
UserID: "test-user-123",
|
||||
TeamID: "test-team-456",
|
||||
TenantID: "test-tenant-789",
|
||||
}
|
||||
|
||||
ctx := agentContext.New(stdContext.Background(), authorized, "test-chat")
|
||||
ctx.AssistantID = "test-assistant"
|
||||
ctx.Locale = "en-us"
|
||||
ctx.Theme = "light"
|
||||
ctx.Client = agentContext.Client{
|
||||
Type: "web",
|
||||
UserAgent: "TestAgent/1.0",
|
||||
IP: "127.0.0.1",
|
||||
}
|
||||
ctx.Referer = agentContext.RefererAPI
|
||||
ctx.Accept = agentContext.AcceptWebCUI
|
||||
ctx.Route = ""
|
||||
ctx.Metadata = make(map[string]interface{})
|
||||
ctx.Capabilities = capabilities
|
||||
ctx.IDGenerator = message.NewIDGenerator()
|
||||
return ctx
|
||||
}
|
||||
|
||||
// newTestOptions creates test options with the given capabilities
|
||||
func newTestOptions(capabilities *openai.Capabilities, completionOptions *agentContext.CompletionOptions) *contentTypes.Options {
|
||||
return &contentTypes.Options{
|
||||
Capabilities: capabilities,
|
||||
CompletionOptions: completionOptions,
|
||||
}
|
||||
}
|
||||
|
||||
// getTestFilePath returns the full path to a test file
|
||||
func getTestFilePath(filename string) string {
|
||||
yaoRoot := os.Getenv("YAO_TEST_APPLICATION")
|
||||
if yaoRoot == "" {
|
||||
yaoRoot = os.Getenv("YAO_ROOT")
|
||||
}
|
||||
return filepath.Join(yaoRoot, testFilesDir, filename)
|
||||
}
|
||||
|
||||
// TestParseWithMissingURL tests parsing PDF with missing URL
|
||||
func TestParseWithMissingURL(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
capabilities := &openai.Capabilities{
|
||||
Vision: "openai",
|
||||
}
|
||||
|
||||
options := newTestOptions(capabilities, nil)
|
||||
ctx := newTestContext(capabilities)
|
||||
|
||||
// Create content with nil File
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: nil,
|
||||
}
|
||||
|
||||
handler := pdf.New(options)
|
||||
_, _, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "missing URL")
|
||||
}
|
||||
|
||||
// TestParseWithEmptyURL tests parsing PDF with empty URL
|
||||
func TestParseWithEmptyURL(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
capabilities := &openai.Capabilities{
|
||||
Vision: "openai",
|
||||
}
|
||||
|
||||
options := newTestOptions(capabilities, nil)
|
||||
ctx := newTestContext(capabilities)
|
||||
|
||||
// Create content with empty URL
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: "",
|
||||
Filename: "test.pdf",
|
||||
},
|
||||
}
|
||||
|
||||
handler := pdf.New(options)
|
||||
_, _, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "missing URL")
|
||||
}
|
||||
|
||||
// TestParseWithLocalPDFAndVisionSupport tests parsing a local PDF file when model supports vision
|
||||
func TestParseWithLocalPDFAndVisionSupport(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Check if test file exists
|
||||
pdfPath := getTestFilePath("test.pdf")
|
||||
if _, err := os.Stat(pdfPath); os.IsNotExist(err) {
|
||||
t.Skipf("Test PDF file not found: %s", pdfPath)
|
||||
}
|
||||
|
||||
// Create capabilities with vision support
|
||||
capabilities := &openai.Capabilities{
|
||||
Vision: "openai",
|
||||
}
|
||||
|
||||
options := newTestOptions(capabilities, nil)
|
||||
ctx := newTestContext(capabilities)
|
||||
|
||||
// Create content with local file path
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: pdfPath,
|
||||
Filename: "test.pdf",
|
||||
},
|
||||
}
|
||||
|
||||
handler := pdf.New(options)
|
||||
result, refs, err := handler.Parse(ctx, content)
|
||||
|
||||
// Should succeed - PDF converted to images
|
||||
assert.NoError(t, err)
|
||||
assert.Nil(t, refs)
|
||||
|
||||
// When model supports vision, Parse returns the first image_url part
|
||||
// Use ParseMulti to get all pages as separate image_url parts
|
||||
assert.Equal(t, agentContext.ContentImageURL, result.Type)
|
||||
assert.NotNil(t, result.ImageURL)
|
||||
assert.NotEmpty(t, result.ImageURL.URL)
|
||||
t.Logf("PDF parse result type: %s, URL prefix: %s...", result.Type, result.ImageURL.URL[:50])
|
||||
}
|
||||
|
||||
// TestParseWithLocalPDFAndVisionAgent tests parsing a local PDF file using vision agent
|
||||
func TestParseWithLocalPDFAndVisionAgent(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Check if test file exists
|
||||
pdfPath := getTestFilePath("test.pdf")
|
||||
if _, err := os.Stat(pdfPath); os.IsNotExist(err) {
|
||||
t.Skipf("Test PDF file not found: %s", pdfPath)
|
||||
}
|
||||
|
||||
// Create capabilities WITHOUT vision support
|
||||
capabilities := &openai.Capabilities{
|
||||
Vision: nil,
|
||||
}
|
||||
|
||||
// Configure to use vision agent
|
||||
completionOptions := &agentContext.CompletionOptions{
|
||||
Uses: &agentContext.Uses{
|
||||
Vision: "tests.vision-test", // Use our test vision agent
|
||||
},
|
||||
}
|
||||
|
||||
options := newTestOptions(capabilities, completionOptions)
|
||||
ctx := newTestContext(capabilities)
|
||||
|
||||
// Create content with local file path
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: pdfPath,
|
||||
Filename: "test.pdf",
|
||||
},
|
||||
}
|
||||
|
||||
handler := pdf.New(options)
|
||||
result, refs, err := handler.Parse(ctx, content)
|
||||
|
||||
// Should succeed - PDF converted to images and processed by vision agent
|
||||
assert.NoError(t, err)
|
||||
assert.Nil(t, refs)
|
||||
assert.Equal(t, agentContext.ContentText, result.Type)
|
||||
assert.NotEmpty(t, result.Text)
|
||||
t.Logf("PDF parse result (via vision agent): %s", result.Text)
|
||||
}
|
||||
|
||||
// TestParseMultiWithLocalPDF tests ParseMulti which returns separate parts for each page
|
||||
func TestParseMultiWithLocalPDF(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Check if test file exists
|
||||
pdfPath := getTestFilePath("test.pdf")
|
||||
if _, err := os.Stat(pdfPath); os.IsNotExist(err) {
|
||||
t.Skipf("Test PDF file not found: %s", pdfPath)
|
||||
}
|
||||
|
||||
// Create capabilities with vision support
|
||||
capabilities := &openai.Capabilities{
|
||||
Vision: "openai",
|
||||
}
|
||||
|
||||
options := newTestOptions(capabilities, nil)
|
||||
ctx := newTestContext(capabilities)
|
||||
|
||||
// Create content with local file path
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: pdfPath,
|
||||
Filename: "test.pdf",
|
||||
},
|
||||
}
|
||||
|
||||
handler := pdf.New(options)
|
||||
parts, refs, err := handler.ParseMulti(ctx, content)
|
||||
|
||||
// Should succeed and return at least one part (one per page)
|
||||
assert.NoError(t, err)
|
||||
assert.Nil(t, refs)
|
||||
assert.NotEmpty(t, parts)
|
||||
t.Logf("PDF ParseMulti returned %d parts", len(parts))
|
||||
|
||||
// When model supports vision, each part should be image_url type
|
||||
for i, part := range parts {
|
||||
assert.Equal(t, agentContext.ContentImageURL, part.Type)
|
||||
assert.NotNil(t, part.ImageURL)
|
||||
t.Logf(" Part %d: type=%s, has URL=%v", i+1, part.Type, part.ImageURL != nil && part.ImageURL.URL != "")
|
||||
}
|
||||
}
|
||||
|
||||
// TestParseWithUnsupportedSource tests parsing PDF with unsupported source
|
||||
func TestParseWithUnsupportedSource(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
capabilities := &openai.Capabilities{
|
||||
Vision: "openai",
|
||||
}
|
||||
|
||||
options := newTestOptions(capabilities, nil)
|
||||
ctx := newTestContext(capabilities)
|
||||
|
||||
// Create content with HTTP URL (not implemented)
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: "https://example.com/test.pdf",
|
||||
Filename: "test.pdf",
|
||||
},
|
||||
}
|
||||
|
||||
handler := pdf.New(options)
|
||||
_, _, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "HTTP URL fetch not implemented")
|
||||
}
|
||||
|
||||
// TestParseWithNonExistentFile tests parsing PDF with non-existent file
|
||||
func TestParseWithNonExistentFile(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
capabilities := &openai.Capabilities{
|
||||
Vision: "openai",
|
||||
}
|
||||
|
||||
options := newTestOptions(capabilities, nil)
|
||||
ctx := newTestContext(capabilities)
|
||||
|
||||
// Create content with non-existent file
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: "/non/existent/path/test.pdf",
|
||||
Filename: "test.pdf",
|
||||
},
|
||||
}
|
||||
|
||||
handler := pdf.New(options)
|
||||
_, _, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "unsupported PDF source")
|
||||
}
|
||||
|
||||
// TestSilentLoadingOption tests that SilentLoading option is respected
|
||||
func TestSilentLoadingOption(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
capabilities := &openai.Capabilities{
|
||||
Vision: "openai",
|
||||
}
|
||||
|
||||
// Create options with SilentLoading enabled
|
||||
options := &contentTypes.Options{
|
||||
Capabilities: capabilities,
|
||||
SilentLoading: true,
|
||||
}
|
||||
|
||||
// This test just verifies the option can be set
|
||||
// The actual behavior is tested in the image handler tests
|
||||
handler := pdf.New(options)
|
||||
assert.NotNil(t, handler)
|
||||
assert.True(t, options.SilentLoading)
|
||||
}
|
||||
143
agent/content/pptx/pptx.go
Normal file
143
agent/content/pptx/pptx.go
Normal file
|
|
@ -0,0 +1,143 @@
|
|||
package pptx
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/yaoapp/gou/office"
|
||||
"github.com/yaoapp/yao/agent/content/types"
|
||||
agentContext "github.com/yaoapp/yao/agent/context"
|
||||
searchTypes "github.com/yaoapp/yao/agent/search/types"
|
||||
"github.com/yaoapp/yao/attachment"
|
||||
)
|
||||
|
||||
// Pptx handles PPTX content
|
||||
type Pptx struct {
|
||||
options *types.Options
|
||||
}
|
||||
|
||||
// New creates a new PPTX handler
|
||||
func New(options *types.Options) *Pptx {
|
||||
return &Pptx{options: options}
|
||||
}
|
||||
|
||||
// Parse parses PPTX content and returns text
|
||||
func (h *Pptx) Parse(ctx *agentContext.Context, content agentContext.ContentPart) (agentContext.ContentPart, []*searchTypes.Reference, error) {
|
||||
if content.File == nil || content.File.URL == "" {
|
||||
return content, nil, fmt.Errorf("file content missing URL")
|
||||
}
|
||||
|
||||
url := content.File.URL
|
||||
|
||||
// Check cache first
|
||||
cachedText, found, err := h.readFromCache(ctx, url)
|
||||
if err == nil && found {
|
||||
return agentContext.ContentPart{
|
||||
Type: agentContext.ContentText,
|
||||
Text: cachedText,
|
||||
}, nil, nil
|
||||
}
|
||||
|
||||
// Read PPTX file
|
||||
data, err := h.readFile(ctx, url)
|
||||
if err != nil {
|
||||
return content, nil, fmt.Errorf("failed to read PPTX: %w", err)
|
||||
}
|
||||
|
||||
// Parse PPTX using gou/office
|
||||
parser := office.NewParser()
|
||||
result, err := parser.Parse(data)
|
||||
if err != nil {
|
||||
return content, nil, fmt.Errorf("failed to parse PPTX: %w", err)
|
||||
}
|
||||
|
||||
text := result.Markdown
|
||||
if text == "" {
|
||||
return content, nil, fmt.Errorf("no text content extracted from PPTX")
|
||||
}
|
||||
|
||||
// Cache the result
|
||||
if err := h.saveToCache(ctx, url, text); err != nil {
|
||||
// Log warning but don't fail
|
||||
fmt.Printf("Warning: failed to cache PPTX text: %v\n", err)
|
||||
}
|
||||
|
||||
return agentContext.ContentPart{
|
||||
Type: agentContext.ContentText,
|
||||
Text: text,
|
||||
}, nil, nil
|
||||
}
|
||||
|
||||
// readFile reads PPTX content from various sources
|
||||
func (h *Pptx) readFile(ctx *agentContext.Context, url string) ([]byte, error) {
|
||||
if strings.HasPrefix(url, "__") {
|
||||
return h.readFromUploader(ctx, url)
|
||||
}
|
||||
|
||||
if strings.HasPrefix(url, "http://") || strings.HasPrefix(url, "https://") {
|
||||
return nil, fmt.Errorf("HTTP URL fetch not implemented yet: %s", url)
|
||||
}
|
||||
|
||||
// Try to read as local file path
|
||||
if _, err := os.Stat(url); err == nil {
|
||||
return os.ReadFile(url)
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("unsupported PPTX source: %s", url)
|
||||
}
|
||||
|
||||
// readFromUploader reads PPTX content from file uploader
|
||||
func (h *Pptx) readFromUploader(ctx *agentContext.Context, wrapper string) ([]byte, error) {
|
||||
uploaderName, fileID, ok := attachment.Parse(wrapper)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid uploader wrapper format: %s", wrapper)
|
||||
}
|
||||
|
||||
manager, exists := attachment.Managers[uploaderName]
|
||||
if !exists {
|
||||
return nil, fmt.Errorf("uploader '%s' not found", uploaderName)
|
||||
}
|
||||
|
||||
data, err := manager.Read(ctx.Context, fileID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read file: %w", err)
|
||||
}
|
||||
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// readFromCache reads cached text content for a PPTX
|
||||
func (h *Pptx) readFromCache(ctx *agentContext.Context, url string) (string, bool, error) {
|
||||
uploaderName, fileID, isWrapper := attachment.Parse(url)
|
||||
if !isWrapper {
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
manager, exists := attachment.Managers[uploaderName]
|
||||
if !exists {
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
text, err := manager.GetText(ctx.Context, fileID, false)
|
||||
if err == nil && text != "" {
|
||||
return text, true, nil
|
||||
}
|
||||
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
// saveToCache saves processed text to cache
|
||||
func (h *Pptx) saveToCache(ctx *agentContext.Context, url string, text string) error {
|
||||
uploaderName, fileID, isWrapper := attachment.Parse(url)
|
||||
if !isWrapper {
|
||||
return nil
|
||||
}
|
||||
|
||||
manager, exists := attachment.Managers[uploaderName]
|
||||
if !exists {
|
||||
return nil
|
||||
}
|
||||
|
||||
return manager.SaveText(ctx.Context, fileID, text)
|
||||
}
|
||||
132
agent/content/pptx/pptx_test.go
Normal file
132
agent/content/pptx/pptx_test.go
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
package pptx_test
|
||||
|
||||
import (
|
||||
stdContext "context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/yaoapp/gou/connector/openai"
|
||||
"github.com/yaoapp/yao/agent/content/pptx"
|
||||
contentTypes "github.com/yaoapp/yao/agent/content/types"
|
||||
agentContext "github.com/yaoapp/yao/agent/context"
|
||||
"github.com/yaoapp/yao/agent/output/message"
|
||||
"github.com/yaoapp/yao/agent/testutils"
|
||||
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
|
||||
)
|
||||
|
||||
const testFilesDir = "assistants/tests/vision-helper/tests"
|
||||
|
||||
func newTestContext() *agentContext.Context {
|
||||
authorized := &oauthTypes.AuthorizedInfo{
|
||||
Subject: "test-user",
|
||||
ClientID: "test-client-id",
|
||||
UserID: "test-user-123",
|
||||
}
|
||||
ctx := agentContext.New(stdContext.Background(), authorized, "test-chat")
|
||||
ctx.AssistantID = "test-assistant"
|
||||
ctx.Locale = "en-us"
|
||||
ctx.IDGenerator = message.NewIDGenerator()
|
||||
return ctx
|
||||
}
|
||||
|
||||
func newTestOptions() *contentTypes.Options {
|
||||
return &contentTypes.Options{
|
||||
Capabilities: &openai.Capabilities{},
|
||||
}
|
||||
}
|
||||
|
||||
func getTestFilePath(filename string) string {
|
||||
yaoRoot := os.Getenv("YAO_TEST_APPLICATION")
|
||||
if yaoRoot == "" {
|
||||
yaoRoot = os.Getenv("YAO_ROOT")
|
||||
}
|
||||
return filepath.Join(yaoRoot, testFilesDir, filename)
|
||||
}
|
||||
|
||||
// TestParseWithMissingURL tests parsing PPTX with missing URL
|
||||
func TestParseWithMissingURL(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
options := newTestOptions()
|
||||
ctx := newTestContext()
|
||||
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: nil,
|
||||
}
|
||||
|
||||
handler := pptx.New(options)
|
||||
_, _, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "missing URL")
|
||||
}
|
||||
|
||||
// TestParseWithLocalPptx tests parsing a local PPTX file
|
||||
func TestParseWithLocalPptx(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
pptxPath := getTestFilePath("pptx.pptx")
|
||||
if _, err := os.Stat(pptxPath); os.IsNotExist(err) {
|
||||
t.Skipf("Test PPTX file not found: %s", pptxPath)
|
||||
}
|
||||
|
||||
options := newTestOptions()
|
||||
ctx := newTestContext()
|
||||
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: pptxPath,
|
||||
Filename: "pptx.pptx",
|
||||
},
|
||||
}
|
||||
|
||||
handler := pptx.New(options)
|
||||
result, refs, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Nil(t, refs)
|
||||
assert.Equal(t, agentContext.ContentText, result.Type)
|
||||
assert.NotEmpty(t, result.Text)
|
||||
t.Logf("PPTX parse result (first 500 chars): %.500s...", result.Text)
|
||||
}
|
||||
|
||||
// TestParseWithNonExistentFile tests parsing PPTX with non-existent file
|
||||
func TestParseWithNonExistentFile(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
options := newTestOptions()
|
||||
ctx := newTestContext()
|
||||
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: "/non/existent/path/test.pptx",
|
||||
Filename: "test.pptx",
|
||||
},
|
||||
}
|
||||
|
||||
handler := pptx.New(options)
|
||||
_, _, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "unsupported PPTX source")
|
||||
}
|
||||
352
agent/content/text/text.go
Normal file
352
agent/content/text/text.go
Normal file
|
|
@ -0,0 +1,352 @@
|
|||
package text
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/yaoapp/yao/agent/content/types"
|
||||
agentContext "github.com/yaoapp/yao/agent/context"
|
||||
searchTypes "github.com/yaoapp/yao/agent/search/types"
|
||||
"github.com/yaoapp/yao/attachment"
|
||||
)
|
||||
|
||||
// SupportedExtensions text file extensions
|
||||
var SupportedExtensions = map[string]bool{
|
||||
// Markdown
|
||||
".md": true,
|
||||
".markdown": true,
|
||||
// Plain text
|
||||
".txt": true,
|
||||
// Code files
|
||||
".go": true,
|
||||
".ts": true,
|
||||
".tsx": true,
|
||||
".js": true,
|
||||
".jsx": true,
|
||||
".py": true,
|
||||
".java": true,
|
||||
".c": true,
|
||||
".cpp": true,
|
||||
".h": true,
|
||||
".hpp": true,
|
||||
".rs": true,
|
||||
".rb": true,
|
||||
".php": true,
|
||||
".swift": true,
|
||||
".kt": true,
|
||||
".scala": true,
|
||||
".sh": true,
|
||||
".bash": true,
|
||||
".zsh": true,
|
||||
".fish": true,
|
||||
".ps1": true,
|
||||
".bat": true,
|
||||
".cmd": true,
|
||||
".sql": true,
|
||||
".r": true,
|
||||
".lua": true,
|
||||
".perl": true,
|
||||
".pl": true,
|
||||
".groovy": true,
|
||||
".dart": true,
|
||||
".elm": true,
|
||||
".ex": true,
|
||||
".exs": true,
|
||||
".erl": true,
|
||||
".hs": true,
|
||||
".clj": true,
|
||||
".lisp": true,
|
||||
".vim": true,
|
||||
// Config files
|
||||
".json": true,
|
||||
".jsonc": true,
|
||||
".yaml": true,
|
||||
".yml": true,
|
||||
".toml": true,
|
||||
".ini": true,
|
||||
".conf": true,
|
||||
".cfg": true,
|
||||
".env": true,
|
||||
".yao": true,
|
||||
// Web files
|
||||
".html": true,
|
||||
".htm": true,
|
||||
".css": true,
|
||||
".scss": true,
|
||||
".sass": true,
|
||||
".less": true,
|
||||
".xml": true,
|
||||
".svg": true,
|
||||
// Documentation
|
||||
".rst": true,
|
||||
".tex": true,
|
||||
".latex": true,
|
||||
".org": true,
|
||||
".adoc": true,
|
||||
// Data files
|
||||
".csv": true,
|
||||
".tsv": true,
|
||||
// Log files
|
||||
".log": true,
|
||||
}
|
||||
|
||||
// Text handles text file content
|
||||
type Text struct {
|
||||
options *types.Options
|
||||
}
|
||||
|
||||
// New creates a new text handler
|
||||
func New(options *types.Options) *Text {
|
||||
return &Text{options: options}
|
||||
}
|
||||
|
||||
// IsSupportedExtension checks if a file extension is supported
|
||||
func IsSupportedExtension(filename string) bool {
|
||||
ext := strings.ToLower(filepath.Ext(filename))
|
||||
return SupportedExtensions[ext]
|
||||
}
|
||||
|
||||
// Parse parses text file content and returns text
|
||||
func (h *Text) Parse(ctx *agentContext.Context, content agentContext.ContentPart) (agentContext.ContentPart, []*searchTypes.Reference, error) {
|
||||
if content.File == nil || content.File.URL == "" {
|
||||
return content, nil, fmt.Errorf("file content missing URL")
|
||||
}
|
||||
|
||||
url := content.File.URL
|
||||
filename := content.File.Filename
|
||||
|
||||
// Check cache first
|
||||
cachedText, found, err := h.readFromCache(ctx, url)
|
||||
if err == nil && found {
|
||||
return agentContext.ContentPart{
|
||||
Type: agentContext.ContentText,
|
||||
Text: cachedText,
|
||||
}, nil, nil
|
||||
}
|
||||
|
||||
// Read text file
|
||||
data, err := h.readFile(ctx, url)
|
||||
if err != nil {
|
||||
return content, nil, fmt.Errorf("failed to read text file: %w", err)
|
||||
}
|
||||
|
||||
// Convert to string
|
||||
text := string(data)
|
||||
|
||||
// Add file type context if it's a code file
|
||||
ext := strings.ToLower(filepath.Ext(filename))
|
||||
if isCodeFile(ext) {
|
||||
// Wrap in markdown code block with language hint
|
||||
lang := getLanguageFromExt(ext)
|
||||
text = fmt.Sprintf("```%s\n%s\n```", lang, text)
|
||||
}
|
||||
|
||||
// Cache the result
|
||||
if err := h.saveToCache(ctx, url, text); err != nil {
|
||||
// Log warning but don't fail
|
||||
fmt.Printf("Warning: failed to cache text: %v\n", err)
|
||||
}
|
||||
|
||||
return agentContext.ContentPart{
|
||||
Type: agentContext.ContentText,
|
||||
Text: text,
|
||||
}, nil, nil
|
||||
}
|
||||
|
||||
// ParseRaw parses any file as raw text content without code block wrapping
|
||||
// This is used as a fallback for unsupported file types
|
||||
func (h *Text) ParseRaw(ctx *agentContext.Context, content agentContext.ContentPart) (agentContext.ContentPart, []*searchTypes.Reference, error) {
|
||||
if content.File == nil || content.File.URL == "" {
|
||||
return content, nil, fmt.Errorf("file content missing URL")
|
||||
}
|
||||
|
||||
url := content.File.URL
|
||||
filename := content.File.Filename
|
||||
|
||||
// Check cache first
|
||||
cachedText, found, err := h.readFromCache(ctx, url)
|
||||
if err == nil && found {
|
||||
return agentContext.ContentPart{
|
||||
Type: agentContext.ContentText,
|
||||
Text: cachedText,
|
||||
}, nil, nil
|
||||
}
|
||||
|
||||
// Read file
|
||||
data, err := h.readFile(ctx, url)
|
||||
if err != nil {
|
||||
return content, nil, fmt.Errorf("failed to read file: %w", err)
|
||||
}
|
||||
|
||||
// Convert to string directly (no code block wrapping)
|
||||
text := string(data)
|
||||
|
||||
// Add filename as context
|
||||
if filename != "" {
|
||||
text = fmt.Sprintf("File: %s\n\n%s", filename, text)
|
||||
}
|
||||
|
||||
// Cache the result
|
||||
if err := h.saveToCache(ctx, url, text); err != nil {
|
||||
// Log warning but don't fail
|
||||
fmt.Printf("Warning: failed to cache text: %v\n", err)
|
||||
}
|
||||
|
||||
return agentContext.ContentPart{
|
||||
Type: agentContext.ContentText,
|
||||
Text: text,
|
||||
}, nil, nil
|
||||
}
|
||||
|
||||
// readFile reads text content from various sources
|
||||
func (h *Text) readFile(ctx *agentContext.Context, url string) ([]byte, error) {
|
||||
if strings.HasPrefix(url, "__") {
|
||||
return h.readFromUploader(ctx, url)
|
||||
}
|
||||
|
||||
if strings.HasPrefix(url, "http://") || strings.HasPrefix(url, "https://") {
|
||||
return nil, fmt.Errorf("HTTP URL fetch not implemented yet: %s", url)
|
||||
}
|
||||
|
||||
// Try to read as local file path
|
||||
if _, err := os.Stat(url); err == nil {
|
||||
return os.ReadFile(url)
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("unsupported text file source: %s", url)
|
||||
}
|
||||
|
||||
// readFromUploader reads text content from file uploader
|
||||
func (h *Text) readFromUploader(ctx *agentContext.Context, wrapper string) ([]byte, error) {
|
||||
uploaderName, fileID, ok := attachment.Parse(wrapper)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid uploader wrapper format: %s", wrapper)
|
||||
}
|
||||
|
||||
manager, exists := attachment.Managers[uploaderName]
|
||||
if !exists {
|
||||
return nil, fmt.Errorf("uploader '%s' not found", uploaderName)
|
||||
}
|
||||
|
||||
data, err := manager.Read(ctx.Context, fileID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read file: %w", err)
|
||||
}
|
||||
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// readFromCache reads cached text content
|
||||
func (h *Text) readFromCache(ctx *agentContext.Context, url string) (string, bool, error) {
|
||||
uploaderName, fileID, isWrapper := attachment.Parse(url)
|
||||
if !isWrapper {
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
manager, exists := attachment.Managers[uploaderName]
|
||||
if !exists {
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
text, err := manager.GetText(ctx.Context, fileID, false)
|
||||
if err == nil && text != "" {
|
||||
return text, true, nil
|
||||
}
|
||||
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
// saveToCache saves processed text to cache
|
||||
func (h *Text) saveToCache(ctx *agentContext.Context, url string, text string) error {
|
||||
uploaderName, fileID, isWrapper := attachment.Parse(url)
|
||||
if !isWrapper {
|
||||
return nil
|
||||
}
|
||||
|
||||
manager, exists := attachment.Managers[uploaderName]
|
||||
if !exists {
|
||||
return nil
|
||||
}
|
||||
|
||||
return manager.SaveText(ctx.Context, fileID, text)
|
||||
}
|
||||
|
||||
// isCodeFile checks if the extension represents a code file
|
||||
func isCodeFile(ext string) bool {
|
||||
codeExts := map[string]bool{
|
||||
".go": true, ".ts": true, ".tsx": true, ".js": true, ".jsx": true,
|
||||
".py": true, ".java": true, ".c": true, ".cpp": true, ".h": true,
|
||||
".hpp": true, ".rs": true, ".rb": true, ".php": true, ".swift": true,
|
||||
".kt": true, ".scala": true, ".sh": true, ".bash": true, ".zsh": true,
|
||||
".sql": true, ".r": true, ".lua": true, ".perl": true, ".pl": true,
|
||||
".groovy": true, ".dart": true, ".elm": true, ".ex": true, ".exs": true,
|
||||
".erl": true, ".hs": true, ".clj": true, ".lisp": true, ".vim": true,
|
||||
}
|
||||
return codeExts[ext]
|
||||
}
|
||||
|
||||
// getLanguageFromExt returns the language name for markdown code block
|
||||
func getLanguageFromExt(ext string) string {
|
||||
langMap := map[string]string{
|
||||
".go": "go",
|
||||
".ts": "typescript",
|
||||
".tsx": "tsx",
|
||||
".js": "javascript",
|
||||
".jsx": "jsx",
|
||||
".py": "python",
|
||||
".java": "java",
|
||||
".c": "c",
|
||||
".cpp": "cpp",
|
||||
".h": "c",
|
||||
".hpp": "cpp",
|
||||
".rs": "rust",
|
||||
".rb": "ruby",
|
||||
".php": "php",
|
||||
".swift": "swift",
|
||||
".kt": "kotlin",
|
||||
".scala": "scala",
|
||||
".sh": "bash",
|
||||
".bash": "bash",
|
||||
".zsh": "zsh",
|
||||
".fish": "fish",
|
||||
".ps1": "powershell",
|
||||
".bat": "batch",
|
||||
".cmd": "batch",
|
||||
".sql": "sql",
|
||||
".r": "r",
|
||||
".lua": "lua",
|
||||
".perl": "perl",
|
||||
".pl": "perl",
|
||||
".groovy": "groovy",
|
||||
".dart": "dart",
|
||||
".elm": "elm",
|
||||
".ex": "elixir",
|
||||
".exs": "elixir",
|
||||
".erl": "erlang",
|
||||
".hs": "haskell",
|
||||
".clj": "clojure",
|
||||
".lisp": "lisp",
|
||||
".vim": "vim",
|
||||
".json": "json",
|
||||
".jsonc": "jsonc",
|
||||
".yaml": "yaml",
|
||||
".yml": "yaml",
|
||||
".toml": "toml",
|
||||
".xml": "xml",
|
||||
".html": "html",
|
||||
".htm": "html",
|
||||
".css": "css",
|
||||
".scss": "scss",
|
||||
".sass": "sass",
|
||||
".less": "less",
|
||||
".svg": "svg",
|
||||
".yao": "json",
|
||||
}
|
||||
|
||||
if lang, ok := langMap[ext]; ok {
|
||||
return lang
|
||||
}
|
||||
return ""
|
||||
}
|
||||
340
agent/content/text/text_test.go
Normal file
340
agent/content/text/text_test.go
Normal file
|
|
@ -0,0 +1,340 @@
|
|||
package text_test
|
||||
|
||||
import (
|
||||
stdContext "context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/yaoapp/gou/connector/openai"
|
||||
"github.com/yaoapp/yao/agent/content/text"
|
||||
contentTypes "github.com/yaoapp/yao/agent/content/types"
|
||||
agentContext "github.com/yaoapp/yao/agent/context"
|
||||
"github.com/yaoapp/yao/agent/output/message"
|
||||
"github.com/yaoapp/yao/agent/testutils"
|
||||
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
|
||||
)
|
||||
|
||||
const testFilesDir = "assistants/tests/vision-helper/tests"
|
||||
|
||||
func newTestContext() *agentContext.Context {
|
||||
authorized := &oauthTypes.AuthorizedInfo{
|
||||
Subject: "test-user",
|
||||
ClientID: "test-client-id",
|
||||
UserID: "test-user-123",
|
||||
}
|
||||
ctx := agentContext.New(stdContext.Background(), authorized, "test-chat")
|
||||
ctx.AssistantID = "test-assistant"
|
||||
ctx.Locale = "en-us"
|
||||
ctx.IDGenerator = message.NewIDGenerator()
|
||||
return ctx
|
||||
}
|
||||
|
||||
func newTestOptions() *contentTypes.Options {
|
||||
return &contentTypes.Options{
|
||||
Capabilities: &openai.Capabilities{},
|
||||
}
|
||||
}
|
||||
|
||||
func getTestFilePath(filename string) string {
|
||||
yaoRoot := os.Getenv("YAO_TEST_APPLICATION")
|
||||
if yaoRoot == "" {
|
||||
yaoRoot = os.Getenv("YAO_ROOT")
|
||||
}
|
||||
return filepath.Join(yaoRoot, testFilesDir, filename)
|
||||
}
|
||||
|
||||
// TestIsSupportedExtension tests the IsSupportedExtension function
|
||||
func TestIsSupportedExtension(t *testing.T) {
|
||||
// Supported extensions
|
||||
assert.True(t, text.IsSupportedExtension("test.md"))
|
||||
assert.True(t, text.IsSupportedExtension("test.txt"))
|
||||
assert.True(t, text.IsSupportedExtension("test.go"))
|
||||
assert.True(t, text.IsSupportedExtension("test.ts"))
|
||||
assert.True(t, text.IsSupportedExtension("test.json"))
|
||||
assert.True(t, text.IsSupportedExtension("test.jsonc"))
|
||||
assert.True(t, text.IsSupportedExtension("test.yao"))
|
||||
assert.True(t, text.IsSupportedExtension("test.yaml"))
|
||||
assert.True(t, text.IsSupportedExtension("test.yml"))
|
||||
assert.True(t, text.IsSupportedExtension("test.py"))
|
||||
assert.True(t, text.IsSupportedExtension("test.js"))
|
||||
assert.True(t, text.IsSupportedExtension("test.css"))
|
||||
assert.True(t, text.IsSupportedExtension("test.html"))
|
||||
|
||||
// Unsupported extensions
|
||||
assert.False(t, text.IsSupportedExtension("test.docx"))
|
||||
assert.False(t, text.IsSupportedExtension("test.pptx"))
|
||||
assert.False(t, text.IsSupportedExtension("test.pdf"))
|
||||
assert.False(t, text.IsSupportedExtension("test.png"))
|
||||
assert.False(t, text.IsSupportedExtension("test.jpg"))
|
||||
assert.False(t, text.IsSupportedExtension("test.exe"))
|
||||
assert.False(t, text.IsSupportedExtension("test.zip"))
|
||||
}
|
||||
|
||||
// TestParseWithMissingURL tests parsing text with missing URL
|
||||
func TestParseWithMissingURL(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
options := newTestOptions()
|
||||
ctx := newTestContext()
|
||||
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: nil,
|
||||
}
|
||||
|
||||
handler := text.New(options)
|
||||
_, _, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "missing URL")
|
||||
}
|
||||
|
||||
// TestParseWithLocalTextFile tests parsing a local text file
|
||||
func TestParseWithLocalTextFile(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
txtPath := getTestFilePath("text.txt")
|
||||
if _, err := os.Stat(txtPath); os.IsNotExist(err) {
|
||||
t.Skipf("Test text file not found: %s", txtPath)
|
||||
}
|
||||
|
||||
options := newTestOptions()
|
||||
ctx := newTestContext()
|
||||
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: txtPath,
|
||||
Filename: "text.txt",
|
||||
},
|
||||
}
|
||||
|
||||
handler := text.New(options)
|
||||
result, refs, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Nil(t, refs)
|
||||
assert.Equal(t, agentContext.ContentText, result.Type)
|
||||
assert.NotEmpty(t, result.Text)
|
||||
t.Logf("Text parse result: %s", result.Text)
|
||||
}
|
||||
|
||||
// TestParseWithLocalMarkdownFile tests parsing a local markdown file
|
||||
func TestParseWithLocalMarkdownFile(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
mdPath := getTestFilePath("test.md")
|
||||
if _, err := os.Stat(mdPath); os.IsNotExist(err) {
|
||||
t.Skipf("Test markdown file not found: %s", mdPath)
|
||||
}
|
||||
|
||||
options := newTestOptions()
|
||||
ctx := newTestContext()
|
||||
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: mdPath,
|
||||
Filename: "test.md",
|
||||
},
|
||||
}
|
||||
|
||||
handler := text.New(options)
|
||||
result, refs, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Nil(t, refs)
|
||||
assert.Equal(t, agentContext.ContentText, result.Type)
|
||||
assert.NotEmpty(t, result.Text)
|
||||
t.Logf("Markdown parse result: %s", result.Text)
|
||||
}
|
||||
|
||||
// TestParseWithLocalCodeFile tests parsing a local code file (TypeScript)
|
||||
func TestParseWithLocalCodeFile(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
tsPath := getTestFilePath("code.ts")
|
||||
if _, err := os.Stat(tsPath); os.IsNotExist(err) {
|
||||
t.Skipf("Test TypeScript file not found: %s", tsPath)
|
||||
}
|
||||
|
||||
options := newTestOptions()
|
||||
ctx := newTestContext()
|
||||
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: tsPath,
|
||||
Filename: "code.ts",
|
||||
},
|
||||
}
|
||||
|
||||
handler := text.New(options)
|
||||
result, refs, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Nil(t, refs)
|
||||
assert.Equal(t, agentContext.ContentText, result.Type)
|
||||
assert.NotEmpty(t, result.Text)
|
||||
// Code files should be wrapped in markdown code blocks
|
||||
assert.True(t, strings.HasPrefix(result.Text, "```typescript"))
|
||||
assert.True(t, strings.HasSuffix(strings.TrimSpace(result.Text), "```"))
|
||||
t.Logf("Code parse result (first 500 chars): %.500s...", result.Text)
|
||||
}
|
||||
|
||||
// TestParseWithLocalYaoFile tests parsing a local .yao file
|
||||
func TestParseWithLocalYaoFile(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
yaoPath := getTestFilePath("hero.mod.yao")
|
||||
if _, err := os.Stat(yaoPath); os.IsNotExist(err) {
|
||||
t.Skipf("Test .yao file not found: %s", yaoPath)
|
||||
}
|
||||
|
||||
options := newTestOptions()
|
||||
ctx := newTestContext()
|
||||
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: yaoPath,
|
||||
Filename: "hero.mod.yao",
|
||||
},
|
||||
}
|
||||
|
||||
handler := text.New(options)
|
||||
result, refs, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Nil(t, refs)
|
||||
assert.Equal(t, agentContext.ContentText, result.Type)
|
||||
assert.NotEmpty(t, result.Text)
|
||||
t.Logf("Yao file parse result: %s", result.Text)
|
||||
}
|
||||
|
||||
// TestParseWithLocalJsonFile tests parsing a local JSON file
|
||||
func TestParseWithLocalJsonFile(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
jsonPath := getTestFilePath("test.json")
|
||||
if _, err := os.Stat(jsonPath); os.IsNotExist(err) {
|
||||
t.Skipf("Test JSON file not found: %s", jsonPath)
|
||||
}
|
||||
|
||||
options := newTestOptions()
|
||||
ctx := newTestContext()
|
||||
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: jsonPath,
|
||||
Filename: "test.json",
|
||||
},
|
||||
}
|
||||
|
||||
handler := text.New(options)
|
||||
result, refs, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Nil(t, refs)
|
||||
assert.Equal(t, agentContext.ContentText, result.Type)
|
||||
assert.NotEmpty(t, result.Text)
|
||||
t.Logf("JSON parse result: %s", result.Text)
|
||||
}
|
||||
|
||||
// TestParseWithNonExistentFile tests parsing text with non-existent file
|
||||
func TestParseWithNonExistentFile(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
options := newTestOptions()
|
||||
ctx := newTestContext()
|
||||
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: "/non/existent/path/test.txt",
|
||||
Filename: "test.txt",
|
||||
},
|
||||
}
|
||||
|
||||
handler := text.New(options)
|
||||
_, _, err := handler.Parse(ctx, content)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "unsupported text file source")
|
||||
}
|
||||
|
||||
// TestParseRawWithLocalFile tests ParseRaw with a local file
|
||||
func TestParseRawWithLocalFile(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
txtPath := getTestFilePath("text.txt")
|
||||
if _, err := os.Stat(txtPath); os.IsNotExist(err) {
|
||||
t.Skipf("Test text file not found: %s", txtPath)
|
||||
}
|
||||
|
||||
options := newTestOptions()
|
||||
ctx := newTestContext()
|
||||
|
||||
content := agentContext.ContentPart{
|
||||
Type: agentContext.ContentFile,
|
||||
File: &agentContext.FileAttachment{
|
||||
URL: txtPath,
|
||||
Filename: "text.txt",
|
||||
},
|
||||
}
|
||||
|
||||
handler := text.New(options)
|
||||
result, refs, err := handler.ParseRaw(ctx, content)
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Nil(t, refs)
|
||||
assert.Equal(t, agentContext.ContentText, result.Type)
|
||||
assert.NotEmpty(t, result.Text)
|
||||
// ParseRaw should include filename as context
|
||||
assert.True(t, strings.HasPrefix(result.Text, "File: text.txt"))
|
||||
t.Logf("ParseRaw result: %s", result.Text)
|
||||
}
|
||||
|
|
@ -20,4 +20,7 @@ type Options struct {
|
|||
|
||||
// StreamOptions, Current stream options instance
|
||||
StreamOptions *agentContext.StreamOptions
|
||||
|
||||
// SilentLoading, if true, suppress loading messages (used when called from parent handler)
|
||||
SilentLoading bool
|
||||
}
|
||||
|
|
|
|||
|
|
@ -192,11 +192,12 @@ type AssistantInfo struct {
|
|||
|
||||
// Skip configuration for what to skip in this request
|
||||
type Skip struct {
|
||||
History bool `json:"history"` // Skip saving chat history (for internal calls like title/prompt generation)
|
||||
Trace bool `json:"trace"` // Skip trace logging
|
||||
Output bool `json:"output"` // Skip output to client (for internal A2A calls that only need response data)
|
||||
Keyword bool `json:"keyword"` // Skip keyword extraction for web search (use raw query directly)
|
||||
Search bool `json:"search"` // Skip auto search (for internal calls like needsearch intent detection)
|
||||
History bool `json:"history"` // Skip saving chat history (for internal calls like title/prompt generation)
|
||||
Trace bool `json:"trace"` // Skip trace logging
|
||||
Output bool `json:"output"` // Skip output to client (for internal A2A calls that only need response data)
|
||||
Keyword bool `json:"keyword"` // Skip keyword extraction for web search (use raw query directly)
|
||||
Search bool `json:"search"` // Skip auto search (for internal calls like needsearch intent detection)
|
||||
ContentParsing bool `json:"content_parsing"` // Skip content parsing (vision, PDF, docx, etc.), convert files to raw text directly
|
||||
}
|
||||
|
||||
// MessageMetadata stores metadata for sent messages
|
||||
|
|
|
|||
|
|
@ -102,6 +102,9 @@ func init() {
|
|||
// Content: content/image/image.go - Image processing messages
|
||||
"content.image.analyzing": "Analyzing image...",
|
||||
|
||||
// Content: content/pdf/pdf.go - PDF processing messages
|
||||
"content.pdf.analyzing_page": "Analyzing PDF page %d/%d...",
|
||||
|
||||
// Search: assistant/search.go - Output messages
|
||||
"search.loading": "Searching...",
|
||||
"search.success": "Found %d references",
|
||||
|
|
@ -198,6 +201,9 @@ func init() {
|
|||
// Content: content/image/image.go - Image processing messages
|
||||
"content.image.analyzing": "正在分析图片...",
|
||||
|
||||
// Content: content/pdf/pdf.go - PDF processing messages
|
||||
"content.pdf.analyzing_page": "正在分析 PDF 第 %d/%d 页...",
|
||||
|
||||
// Search: assistant/search.go - Output messages
|
||||
"search.loading": "正在搜索...",
|
||||
"search.success": "找到 %d 条参考资料",
|
||||
|
|
@ -322,6 +328,9 @@ func init() {
|
|||
// Content: content/image/image.go - Image processing messages
|
||||
"content.image.analyzing": "正在分析图片...",
|
||||
|
||||
// Content: content/pdf/pdf.go - PDF processing messages
|
||||
"content.pdf.analyzing_page": "正在分析 PDF 第 %d/%d 页...",
|
||||
|
||||
// Search: assistant/search.go - Output messages
|
||||
"search.loading": "正在搜索...",
|
||||
"search.success": "找到 %d 条参考资料",
|
||||
|
|
|
|||
|
|
@ -124,9 +124,23 @@ func (s *Searcher) parallelAll(ctx *context.Context, reqs []*types.Request) ([]*
|
|||
wg.Add(1)
|
||||
go func(idx int, r *types.Request) {
|
||||
defer wg.Done()
|
||||
result, _ := s.Search(ctx, r)
|
||||
defer func() {
|
||||
if err := recover(); err != nil {
|
||||
mu.Lock()
|
||||
results[idx] = &types.Result{Error: "search panic recovered"}
|
||||
mu.Unlock()
|
||||
}
|
||||
}()
|
||||
|
||||
result, err := s.Search(ctx, r)
|
||||
mu.Lock()
|
||||
results[idx] = result
|
||||
if err != nil {
|
||||
results[idx] = &types.Result{Error: err.Error()}
|
||||
} else if result == nil {
|
||||
results[idx] = &types.Result{Error: "empty result"}
|
||||
} else {
|
||||
results[idx] = result
|
||||
}
|
||||
mu.Unlock()
|
||||
}(i, req)
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue