Add support for Drive folder structure

docker/traefik/dynamic.toml

@@ -28,7 +28,7 @@
 
 [http.services.ui.loadBalancer]
   [[http.services.ui.loadBalancer.servers]]
-    url = "http://cube-ui:5173"
+    url = "http://ui:5173"
   [http.services.ui.loadBalancer.healthCheck]
     scheme = "http"
     path = "/health"

internal/embedder/api/transport/chat.go

@@ -14,6 +14,12 @@ import (
 	"github.com/ultravioletrs/cube/internal/embedder/domain"
 )
 
+// maxChatRecordIDs bounds the explicit record scope a single chat request may
+// carry. A larger allowlist bloats the request and the SQL IN (...) clause
+// (Postgres caps bind parameters at 65535, and the plan degrades well before).
+// Clients that want "all records" send no record_ids, which searches unscoped.
+const maxChatRecordIDs = 1000
+
 // MountChat registers the streaming chat endpoint.
 func MountChat(r chi.Router, svc domain.ChatService, conversations domain.ConversationRepository) {
 	r.Post("/api/v1/chat", chatHandler(svc, conversations))
@@ -41,6 +47,10 @@ func chatHandler(svc domain.ChatService, conversations domain.ConversationReposi
 			writeJSON(w, http.StatusBadRequest, errBody("messages is required"))
 			return
 		}
+		if len(req.RecordIDs) > maxChatRecordIDs {
+			writeJSON(w, http.StatusBadRequest, errBody(fmt.Sprintf("record_ids exceeds limit of %d", maxChatRecordIDs)))
+			return
+		}
 
 		// Ensure we have a conversation to save messages into.
 		convID := req.ConversationID

internal/embedder/api/transport/records.go

@@ -6,6 +6,7 @@ package transport
 import (
 	"errors"
 	"net/http"
+	"strings"
 
 	"github.com/go-chi/chi/v5"
 	"github.com/ultravioletrs/cube/internal/embedder/auth"
@@ -46,6 +47,8 @@ type recordResponse struct {
 	ExternalURL         string                `json:"external_url"`
 	ExternalRef         string                `json:"external_ref,omitempty"`
 	MimeType            string                `json:"mime_type,omitempty"`
+	FolderPath          *string               `json:"folder_path,omitempty"`
+	FolderID            *string               `json:"folder_id,omitempty"`
 	Description         string                `json:"description,omitempty"`
 	ChunkCount          *int                  `json:"chunks,omitempty"`
 	IngestTotalChunks   *int                  `json:"ingest_total_chunks,omitempty"`
@@ -78,6 +81,8 @@ func toRecordResponse(rec domain.Record) recordResponse {
 		ExternalURL:         rec.ExternalURL,
 		ExternalRef:         rec.ExternalRef,
 		MimeType:            rec.MimeType,
+		FolderPath:          rec.FolderPath,
+		FolderID:            rec.FolderID,
 		Description:         rec.Description,
 		ChunkCount:          rec.ChunkCount,
 		IngestTotalChunks:   rec.IngestTotalChunks,
@@ -224,6 +229,12 @@ func parseRecordFilter(r *http.Request) domain.RecordFilter {
 		fmt := domain.RecordFormat(s)
 		f.Format = &fmt
 	}
+	if s := strings.TrimSpace(r.URL.Query().Get("q")); s != "" {
+		f.Name = &s
+	}
+	if s := strings.TrimSpace(r.URL.Query().Get("folder")); s != "" {
+		f.FolderPrefix = &s
+	}
 	return f
 }
 

internal/embedder/domain/record.go

@@ -57,6 +57,12 @@ type Record struct {
 	ExternalRef string
 	MimeType    string
 
+	// FolderPath is the human-readable containing-folder path within the source
+	// (e.g. /Docs/2024/Q3); FolderID is the immediate parent folder ID. Both are
+	// populated for folder-tree ingests (Google Drive); nil otherwise.
+	FolderPath *string
+	FolderID   *string
+
 	// Content metadata populated after successful ingestion.
 	Description         string
 	ChunkCount          *int
@@ -99,6 +105,11 @@ type RecordFilter struct {
 	SourceID *string
 	Status   *RecordStatus
 	Format   *RecordFormat
+	// Name is a case-insensitive substring matched against the record name.
+	Name *string
+	// FolderPrefix matches records whose folder_path equals or is nested under
+	// the given path (prefix match), e.g. "/Docs/2024".
+	FolderPrefix *string
 }
 
 // IngestResult holds post-ingestion metadata written back to the record.

internal/embedder/ingest/drive.go

@@ -10,6 +10,7 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"encoding/xml"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -37,6 +38,14 @@ var (
 	driveDownloadURL  = "https://www.googleapis.com/drive/v3/files/%s?alt=media"
 )
 
+var (
+	// ErrDriveNotFound is returned when a Drive item no longer exists (404/410),
+	// so callers can skip stale references instead of failing the whole sync.
+	ErrDriveNotFound = errors.New("drive file not found")
+	// ErrUnsupportedDriveFile is returned when a file's MIME type cannot be ingested.
+	ErrUnsupportedDriveFile = errors.New("drive file has unsupported MIME type")
+)
+
 // SetDriveAPIEndpoints overrides Drive API endpoints and returns a restore function.
 // Intended for tests that need deterministic HTTP fixtures.
 func SetDriveAPIEndpoints(filesURL, exportURLFmt, downloadURLFmt string) func() {
@@ -64,6 +73,13 @@ type DriveFile struct {
 	ModifiedTime string   `json:"modifiedTime"`
 	WebViewLink  string   `json:"webViewLink"`
 	Parents      []string `json:"parents"`
+
+	// FolderPath and FolderID are populated by ListFilesRecursive while walking
+	// the folder tree. FolderPath is the human-readable path of the containing
+	// folder relative to the walked root (e.g. /Docs/2024/Q3); FolderID is the
+	// immediate parent folder ID. Empty for flat (whole-drive) listings.
+	FolderPath string `json:"-"`
+	FolderID   string `json:"-"`
 }
 
 // ImageIngestMode describes which signals should be indexed for an image.
@@ -226,6 +242,46 @@ func (d *DriveReader) ListFiles(ctx context.Context, folderID string) ([]DriveFi
 	return all, nil
 }
 
+// GetFile returns Drive metadata for a single file ID.
+func (d *DriveReader) GetFile(ctx context.Context, fileID string) (DriveFile, error) {
+	id := strings.TrimSpace(fileID)
+	if id == "" {
+		return DriveFile{}, fmt.Errorf("drive file id is required")
+	}
+
+	params := url.Values{
+		"fields":            {"id,name,mimeType,version,modifiedTime,webViewLink,parents"},
+		"supportsAllDrives": {"true"},
+	}
+	reqURL := strings.TrimRight(driveFilesURL, "/") + "/" + url.PathEscape(id) + "?" + params.Encode()
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, http.NoBody)
+	if err != nil {
+		return DriveFile{}, fmt.Errorf("drive get file request: %w", err)
+	}
+	resp, err := d.httpClient.Do(req)
+	if err != nil {
+		return DriveFile{}, fmt.Errorf("drive get file: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusGone {
+		return DriveFile{}, fmt.Errorf("drive get file %s: %w", id, ErrDriveNotFound)
+	}
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		return DriveFile{}, fmt.Errorf("drive get file status %d: %s", resp.StatusCode, body)
+	}
+
+	var file DriveFile
+	if err := json.NewDecoder(resp.Body).Decode(&file); err != nil {
+		return DriveFile{}, fmt.Errorf("drive get file decode: %w", err)
+	}
+	if !supportsDriveFile(file) {
+		return DriveFile{}, fmt.Errorf("drive file %s (%q): %w", id, file.MimeType, ErrUnsupportedDriveFile)
+	}
+	return file, nil
+}
+
 // ListFilesRecursive returns supported files contained in folderID and all of
 // its descendant folders.
 func (d *DriveReader) ListFilesRecursive(ctx context.Context, folderID string) ([]DriveFile, error) {
@@ -234,15 +290,22 @@ func (d *DriveReader) ListFilesRecursive(ctx context.Context, folderID string) (
 		return d.ListFiles(ctx, "")
 	}
 
-	queue := []string{rootID}
+	// BFS the folder tree carrying each folder's path so files can record the
+	// human-readable folder path of their container without extra lookups.
+	type folderNode struct {
+		id   string
+		path string
+	}
+	rootPath := "/" + d.folderName(ctx, rootID)
+	queue := []folderNode{{id: rootID, path: rootPath}}
 	seenFolders := map[string]struct{}{rootID: {}}
 	filesByID := make(map[string]DriveFile)
 
 	for len(queue) > 0 {
 		current := queue[0]
 		queue = queue[1:]
 
-		folders, files, err := d.ListFolderContent(ctx, current)
+		folders, files, err := d.ListFolderContent(ctx, current.id)
 		if err != nil {
 			return nil, err
 		}
@@ -251,9 +314,11 @@ func (d *DriveReader) ListFilesRecursive(ctx context.Context, folderID string) (
 				continue
 			}
 			seenFolders[folder.ID] = struct{}{}
-			queue = append(queue, folder.ID)
+			queue = append(queue, folderNode{id: folder.ID, path: current.path + "/" + folder.Name})
 		}
 		for _, file := range files {
+			file.FolderPath = current.path
+			file.FolderID = current.id
 			filesByID[file.ID] = file
 		}
 	}
@@ -268,6 +333,31 @@ func (d *DriveReader) ListFilesRecursive(ctx context.Context, folderID string) (
 	return out, nil
 }
 
+// folderName fetches a folder's display name. On any error it falls back to the
+// folder ID so path construction stays best-effort and never blocks ingest.
+func (d *DriveReader) folderName(ctx context.Context, folderID string) string {
+	params := url.Values{"fields": {"name"}, "supportsAllDrives": {"true"}}
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, driveFilesURL+"/"+folderID+"?"+params.Encode(), http.NoBody)
+	if err != nil {
+		return folderID
+	}
+	resp, err := d.httpClient.Do(req)
+	if err != nil {
+		return folderID
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return folderID
+	}
+	var meta struct {
+		Name string `json:"name"`
+	}
+	if err := json.NewDecoder(resp.Body).Decode(&meta); err != nil || strings.TrimSpace(meta.Name) == "" {
+		return folderID
+	}
+	return meta.Name
+}
+
 // ListFolderContent returns direct children of a folder split into folders and
 // supported files. Empty folderID resolves to Drive root.
 func (d *DriveReader) ListFolderContent(ctx context.Context, folderID string) ([]DriveFile, []DriveFile, error) {

internal/embedder/ingest/source_providers.go

@@ -19,6 +19,10 @@ type SourceFile struct {
 	MimeType         string
 	SourceVersion    string
 	SourceModifiedAt *time.Time
+	// FolderPath is the human-readable containing-folder path (e.g. /Docs/2024);
+	// FolderID is the immediate parent folder ID. Both optional / best-effort.
+	FolderPath string
+	FolderID   string
 }
 
 // SourceProviderCapabilities describes what integration operations are supported.

internal/embedder/ingest/sources/google/provider.go

@@ -6,6 +6,7 @@ package google
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"strings"
 	"time"
@@ -64,6 +65,10 @@ func (p *sourceProvider) ListFiles(
 
 	out := make([]ingest.SourceFile, 0, len(files))
 	for _, file := range files {
+		folderID := file.FolderID
+		if folderID == "" && len(file.Parents) > 0 {
+			folderID = file.Parents[0]
+		}
 		out = append(out, ingest.SourceFile{
 			ExternalID:       file.ID,
 			Name:             file.Name,
@@ -72,6 +77,8 @@ func (p *sourceProvider) ListFiles(
 			MimeType:         file.MimeType,
 			SourceVersion:    file.Version,
 			SourceModifiedAt: parseRFC3339Ptr(file.ModifiedTime),
+			FolderPath:       file.FolderPath,
+			FolderID:         folderID,
 		})
 	}
 	return out, nil
@@ -156,7 +163,18 @@ func applyDriveSelection(
 	for _, id := range selectedFiles {
 		if file, ok := baseByID[id]; ok {
 			collected[file.ID] = file
+			continue
+		}
+		file, err := reader.GetFile(ctx, id)
+		if err != nil {
+			// Skip a selected file that has gone away or can't be ingested rather
+			// than failing the entire sync over one stale reference.
+			if errors.Is(err, ingest.ErrDriveNotFound) || errors.Is(err, ingest.ErrUnsupportedDriveFile) {
+				continue
+			}
+			return nil, err
 		}
+		collected[file.ID] = file
 	}
 
 	for _, folderID := range selectedFolders {