Agent-System-Backend/internal/runmgr/material_retrieve.go

package runmgr

import (
	"crypto/sha256"
	"encoding/hex"
	"strings"
	"unicode"
	"unicode/utf8"

	core "agent-system-core"
)

// materialRetrievalResult keeps ordinary retrieval authoritative. A no-hit
// material child may fall back only to the attachments frozen for this task.
func (c *runCtx) materialRetrievalResult(q core.RetrievalQuery, hits []core.RetrievalHit, err error) ([]core.RetrievalHit, error) {
	if err != nil || len(hits) != 0 || c.subAgent != materialSubagentID {
		return hits, err
	}
	return c.materialAttachmentFallback(q)
}

func (c *runCtx) materialAttachmentFallback(q core.RetrievalQuery) ([]core.RetrievalHit, error) {
	if c.task.cardID == "" || c.run.Owner == "" || c.run.ProjectID == "" ||
		q.ProjectID != c.run.ProjectID || strings.TrimSpace(q.Text) == "" {
		return nil, nil
	}
	if len(q.Sources) > 0 {
		allowed := false
		for _, source := range q.Sources {
			allowed = allowed || source == "upload"
		}
		if !allowed {
			return nil, nil
		}
	}
	var matched *core.ParsedDocument
	for i := range c.attachments {
		name := strings.TrimSpace(c.attachments[i].Meta.Filename)
		if exactFilenameIndex(q.Text, name) < 0 {
			continue
		}
		if matched != nil {
			return nil, nil // duplicate or multiple filenames: never guess
		}
		matched = &c.attachments[i]
	}
	if matched == nil || matched.Meta.Origin == nil {
		return nil, nil
	}
	origin := matched.Meta.Origin
	if origin.ArtifactID == "" || origin.ContentHash == "" || origin.ParserRevision == "" {
		return nil, nil
	}
	card, err := c.mgr.st.GetCard(c.run.Owner, c.task.cardID)
	if err != nil {
		return nil, err
	}
	s := card.Task
	if s == nil || s.Terminal() || s.TaskID != c.task.taskID ||
		s.AttemptID != c.task.attemptID || s.RunID != c.run.ID ||
		s.OwnerInstanceID != c.task.instanceID || s.ExecutionEpoch != c.task.epoch {
		return nil, nil
	}
	bound := 0
	for _, input := range s.Inputs {
		if input.ArtifactID == string(origin.ArtifactID) && input.ContentHash == origin.ContentHash &&
			input.ParserRevision == origin.ParserRevision {
			bound++
		}
	}
	if bound != 1 {
		return nil, nil
	}
	a, err := c.mgr.st.GetArtifact(c.run.Owner, string(origin.ArtifactID))
	if err != nil {
		return nil, err
	}
	if a.ProjectID != c.run.ProjectID || a.Name != matched.Meta.Filename {
		return nil, nil
	}
	data, err := c.mgr.st.ArtifactContent(c.run.Owner, string(origin.ArtifactID))
	if err != nil {
		return nil, err
	}
	sum := sha256.Sum256(data)
	if hex.EncodeToString(sum[:]) != origin.ContentHash {
		return nil, nil
	}
	after, err := c.mgr.st.GetArtifact(c.run.Owner, string(origin.ArtifactID))
	if err != nil || !after.UpdatedAt.Equal(a.UpdatedAt) {
		return nil, err
	}
	full := strings.TrimSpace(matched.PlainText())
	if full == "" {
		return nil, nil
	}
	limit := c.prof.Context.MaxToolResultChars
	if limit <= 0 {
		limit = int(c.eff.GetNumber("context.maxToolResultChars", 20000))
	}
	if limit <= 0 || limit > 4000 {
		limit = 4000
	}
	doc := core.KnowledgeDoc{Owner: c.run.Owner, ProjectID: c.run.ProjectID,
		Source: "upload", Title: matched.Meta.Filename, Ref: string(origin.ArtifactID),
		Chunk:     materialExcerpt(full, q.Text, matched.Meta.Filename, c.attachments, origin.ArtifactID, c.eff, limit),
		CreatedAt: a.CreatedAt}
	if c.mgr.knowledgeResourceAllow != nil && !c.mgr.knowledgeResourceAllow(string(c.run.Owner), doc) {
		return nil, nil
	}
	return []core.RetrievalHit{{Doc: doc, Score: 1, Kind: "fulltext"}}, nil
}

// A filename must stand as a whole query item, including when it has spaces.
// Nearby letters, digits, dots and hyphens may belong to another filename.
func exactFilenameIndex(query, name string) int {
	if name == "" {
		return -1
	}
	boundary := func(r rune) bool {
		return unicode.IsSpace(r) || strings.ContainsRune(`"'“”‘’「」『』()（）[]【】,，;；:：`, r)
	}
	for from := 0; from < len(query); {
		at := strings.Index(query[from:], name)
		if at < 0 {
			return -1
		}
		at += from
		before := at == 0
		if !before {
			r, _ := utf8.DecodeLastRuneInString(query[:at])
			before = boundary(r)
		}
		end := at + len(name)
		after := end == len(query)
		if !after {
			r, _ := utf8.DecodeRuneInString(query[end:])
			after = boundary(r)
		}
		if before && after {
			return at
		}
		from = at + 1
	}
	return -1
}

func materialExcerpt(full, query, filename string, docs []core.ParsedDocument, id core.ID, eff core.EffectiveSettings, limit int) string {
	runes := []rune(full)
	start := 0
	matchedTerm := false
	withoutName := query
	if at := exactFilenameIndex(query, filename); at >= 0 {
		withoutName = query[:at] + query[at+len(filename):]
	}
	terms := strings.Fields(strings.TrimSpace(withoutName))
	for _, term := range terms {
		if at := strings.Index(full, term); at >= 0 {
			matchedTerm = true
			start = len([]rune(full[:at])) - 80
			if start < 0 {
				start = 0
			}
			break
		}
	}
	if !matchedTerm {
		perDoc := int(eff.GetNumber(materialPerDocKey, materialPerDocDef))
		if perDoc <= 0 {
			perDoc = materialPerDocDef
		}
		left := materialInlineTotal
		for _, d := range docs {
			count := len([]rune(strings.TrimSpace(d.PlainText())))
			provided := min(perDoc, left, count)
			if d.Meta.Origin != nil && d.Meta.Origin.ArtifactID == id {
				if provided < len(runes) {
					start = provided
				}
				break
			}
			left -= provided
		}
	}
	end := min(start+limit, len(runes))
	return string(runes[start:end])
}