Agent-System/internal/controlcenter/server/docs_handlers.go

package server

import (
	"context"
	"encoding/json"
	"log/slog"
	"net/http"
	"strconv"
	"strings"
	"time"

	"chat-agent-system/internal/controlcenter/embed"
)

const (
	searchDefaultLimit = 20
	searchMaxLimit     = 100
	snippetRunes       = 300
)

// docListItem is one row of GET /api/v1/docs.
type docListItem struct {
	ID          string    `json:"id"`
	Path        string    `json:"path"`
	Title       string    `json:"title"`
	Taxonomy    string    `json:"taxonomy"`
	Frontmatter any       `json:"frontmatter,omitempty"`
	SyncedAt    time.Time `json:"synced_at"`
	Chunks      int       `json:"chunks"`
}

func (s *Server) handleDocList(w http.ResponseWriter, r *http.Request) {
	taxonomy := strings.TrimSpace(r.URL.Query().Get("taxonomy"))
	q := strings.TrimSpace(r.URL.Query().Get("q"))
	like := "%" + escapeLike(q) + "%"

	rows, err := s.store.Pool.Query(r.Context(), `
		SELECT d.id, d.path, d.title, d.taxonomy, d.synced_at,
		       (SELECT count(*) FROM cc_doc_chunks c WHERE c.doc_id = d.id) AS chunks
		FROM cc_documents d
		WHERE ($1 = '' OR d.taxonomy = $1)
		  AND ($2 = '' OR d.path ILIKE $3 ESCAPE '\' OR d.title ILIKE $3 ESCAPE '\')
		ORDER BY d.path`, taxonomy, q, like)
	if err != nil {
		slog.Error("doc list", "err", err)
		writeError(w, http.StatusInternalServerError, "internal error")
		return
	}
	defer rows.Close()
	items := []docListItem{}
	for rows.Next() {
		var it docListItem
		if err := rows.Scan(&it.ID, &it.Path, &it.Title, &it.Taxonomy, &it.SyncedAt, &it.Chunks); err != nil {
			slog.Error("doc list scan", "err", err)
			writeError(w, http.StatusInternalServerError, "internal error")
			return
		}
		items = append(items, it)
	}
	writeJSON(w, http.StatusOK, map[string]any{"items": items, "count": len(items)})
}

func (s *Server) handleDocGet(w http.ResponseWriter, r *http.Request) {
	var it docListItem
	var fm []byte
	err := s.store.Pool.QueryRow(r.Context(), `
		SELECT d.id, d.path, d.title, d.taxonomy, d.frontmatter, d.synced_at,
		       (SELECT count(*) FROM cc_doc_chunks c WHERE c.doc_id = d.id) AS chunks
		FROM cc_documents d WHERE d.id = $1`, r.PathValue("id")).
		Scan(&it.ID, &it.Path, &it.Title, &it.Taxonomy, &fm, &it.SyncedAt, &it.Chunks)
	if err != nil {
		writeError(w, http.StatusNotFound, "document not found")
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{
		"id": it.ID, "path": it.Path, "title": it.Title, "taxonomy": it.Taxonomy,
		"frontmatter": jsonRaw(fm), "synced_at": it.SyncedAt, "chunks": it.Chunks,
	})
}

// searchHit is one scored chunk with its source resolved.
type searchHit struct {
	chunkID  int64
	docID    *string
	artID    *string
	heading  string
	content  string
	score    float64
	docPath  string
	docTitle string
	taskSlug string
	artStage string
	artTitle string
}

const searchSelect = `
	SELECT c.id, c.doc_id, c.artifact_id, c.heading, c.content,
	       COALESCE(d.path,''), COALESCE(d.title,''),
	       COALESCE(t.slug,''), COALESCE(a.stage,''), COALESCE(a.title,'')`

const searchJoin = `
	FROM cc_doc_chunks c
	LEFT JOIN cc_documents d ON d.id = c.doc_id
	LEFT JOIN cc_stage_artifacts a ON a.id = c.artifact_id
	LEFT JOIN cc_tasks t ON t.id = a.task_id`

func (s *Server) semanticHits(ctx context.Context, qvec []float32, limit int) ([]searchHit, error) {
	lit := embed.VectorLiteral(qvec)
	rows, err := s.store.Pool.Query(ctx, searchSelect+`,
	       1 - (c.embedding <=> $1::vector) AS score`+searchJoin+`
	WHERE c.embedding IS NOT NULL
	ORDER BY c.embedding <=> $1::vector
	LIMIT $2`, lit, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanHits(rows)
}

func (s *Server) keywordHits(ctx context.Context, q string, limit int) ([]searchHit, error) {
	like := "%" + escapeLike(q) + "%"
	rows, err := s.store.Pool.Query(ctx, searchSelect+`,
	       (CASE WHEN c.tsv @@ plainto_tsquery('simple', $1)
	             THEN ts_rank(c.tsv, plainto_tsquery('simple', $1)) ELSE 0 END
	      + CASE WHEN c.content ILIKE $2 ESCAPE '\' THEN 0.5 ELSE 0 END
	      + CASE WHEN c.heading ILIKE $2 ESCAPE '\' THEN 0.3 ELSE 0 END) AS score`+searchJoin+`
	WHERE c.tsv @@ plainto_tsquery('simple', $1)
	   OR c.content ILIKE $2 ESCAPE '\'
	   OR c.heading ILIKE $2 ESCAPE '\'
	ORDER BY score DESC
	LIMIT $3`, q, like, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanHits(rows)
}

func scanHits(rows interface {
	Next() bool
	Scan(...any) error
	Err() error
}) ([]searchHit, error) {
	var out []searchHit
	for rows.Next() {
		var h searchHit
		if err := rows.Scan(&h.chunkID, &h.docID, &h.artID, &h.heading, &h.content,
			&h.docPath, &h.docTitle, &h.taskSlug, &h.artStage, &h.artTitle, &h.score); err != nil {
			return nil, err
		}
		out = append(out, h)
	}
	return out, rows.Err()
}

func (s *Server) handleDocSearch(w http.ResponseWriter, r *http.Request) {
	q := strings.TrimSpace(r.URL.Query().Get("q"))
	if q == "" {
		writeError(w, http.StatusBadRequest, "q is required")
		return
	}
	mode := r.URL.Query().Get("mode")
	if mode == "" {
		mode = "hybrid"
	}
	if mode != "semantic" && mode != "keyword" && mode != "hybrid" {
		writeError(w, http.StatusBadRequest, "mode must be semantic|keyword|hybrid")
		return
	}
	limit := searchDefaultLimit
	if raw := r.URL.Query().Get("limit"); raw != "" {
		n, err := strconv.Atoi(raw)
		if err != nil || n < 1 {
			writeError(w, http.StatusBadRequest, "limit must be a positive integer")
			return
		}
		limit = min(n, searchMaxLimit)
	}

	ctx := r.Context()
	degraded := false
	var vecHits []searchHit
	if mode != "keyword" {
		if s.emb == nil {
			degraded = true
		} else if vecs, err := s.emb.Embed(ctx, []string{q}); err != nil || len(vecs) == 0 {
			slog.Warn("search: query embedding failed, degrading to keyword", "err", err)
			degraded = true
		} else if hits, err := s.semanticHits(ctx, vecs[0], limit); err != nil {
			slog.Error("search: semantic query", "err", err)
			writeError(w, http.StatusInternalServerError, "internal error")
			return
		} else {
			vecHits = hits
		}
	}

	var hits []searchHit
	switch {
	case mode == "semantic" && !degraded:
		hits = vecHits
	case mode == "hybrid" && !degraded:
		kwHits, err := s.keywordHits(ctx, q, limit)
		if err != nil {
			slog.Error("search: keyword query", "err", err)
			writeError(w, http.StatusInternalServerError, "internal error")
			return
		}
		hits = fuseHits(vecHits, kwHits, limit)
	default: // keyword, or degraded semantic/hybrid
		kwHits, err := s.keywordHits(ctx, q, limit)
		if err != nil {
			slog.Error("search: keyword query", "err", err)
			writeError(w, http.StatusInternalServerError, "internal error")
			return
		}
		hits = kwHits
	}

	results := make([]map[string]any, 0, len(hits))
	for _, h := range hits {
		source := map[string]any{"type": "doc", "path": h.docPath, "title": h.docTitle}
		if h.artID != nil {
			source = map[string]any{"type": "artifact", "task_slug": h.taskSlug, "stage": h.artStage, "title": h.artTitle}
		}
		results = append(results, map[string]any{
			"score":       h.score,
			"snippet":     snippet(h.content),
			"heading":     h.heading,
			"source":      source,
			"doc_id":      h.docID,
			"artifact_id": h.artID,
		})
	}
	writeJSON(w, http.StatusOK, map[string]any{
		"q": q, "mode": mode, "degraded": degraded, "count": len(results), "results": results,
	})
}

// fuseHits merges semantic (weight 0.6) and keyword (weight 0.4) hits after
// per-list max-normalization.
func fuseHits(vecHits, kwHits []searchHit, limit int) []searchHit {
	byID := map[int64]*searchHit{}
	norm := func(hits []searchHit, weight float64) {
		maxScore := 0.0
		for _, h := range hits {
			if h.score > maxScore {
				maxScore = h.score
			}
		}
		if maxScore <= 0 {
			maxScore = 1
		}
		for _, h := range hits {
			score := weight * h.score / maxScore
			if ex, ok := byID[h.chunkID]; ok {
				ex.score += score
			} else {
				cpy := h
				cpy.score = score
				byID[h.chunkID] = &cpy
			}
		}
	}
	norm(vecHits, 0.6)
	norm(kwHits, 0.4)

	out := make([]searchHit, 0, len(byID))
	for _, h := range byID {
		out = append(out, *h)
	}
	for i := 1; i < len(out); i++ { // insertion sort by score desc, hits are few
		for j := i; j > 0 && out[j].score > out[j-1].score; j-- {
			out[j], out[j-1] = out[j-1], out[j]
		}
	}
	if len(out) > limit {
		out = out[:limit]
	}
	return out
}

// snippet returns the first ~snippetRunes runes of content.
func snippet(content string) string {
	r := []rune(strings.TrimSpace(content))
	if len(r) > snippetRunes {
		r = r[:snippetRunes]
	}
	return string(r)
}

func escapeLike(s string) string {
	r := strings.NewReplacer(`\`, `\\`, `%`, `\%`, `_`, `\_`)
	return r.Replace(s)
}

// jsonRaw returns frontmatter bytes as an any for the JSON encoder.
func jsonRaw(b []byte) any {
	var v any
	if len(b) > 0 {
		if err := json.Unmarshal(b, &v); err == nil {
			return v
		}
	}
	return map[string]any{}
}