Merge branch 'main' into hosted-gpu

Adityav369 · Adityav369 · commit c71a8317d670 · 2025-12-10T19:49:34.000-08:00
diff --git a/core/database/folder_bootstrap.py b/core/database/folder_bootstrap.py
@@ -19,6 +19,7 @@ async def bootstrap_folder_hierarchy(engine: AsyncEngine, logger) -> None:
         ("folders", "parent_id"),
         ("folders", "depth"),
         ("documents", "folder_path"),
+        ("documents", "folder_id"),
         ("graphs", "folder_path"),
     }
     required_indexes = {
@@ -28,7 +29,9 @@ async def bootstrap_folder_hierarchy(engine: AsyncEngine, logger) -> None:
         "uq_folders_app_full_path",
         "uq_folders_owner_full_path",
         "idx_doc_folder_path",
+        "idx_doc_folder_id",
         "idx_documents_app_folder_path",
+        "idx_documents_app_folder_id",
         "idx_graph_folder_path",
         "idx_graphs_app_folder_path",
     }
@@ -41,7 +44,7 @@ async def bootstrap_folder_hierarchy(engine: AsyncEngine, logger) -> None:
                     SELECT table_name, column_name
                     FROM information_schema.columns
                     WHERE table_name IN ('folders', 'documents', 'graphs')
-                    AND column_name IN ('full_path', 'parent_id', 'depth', 'folder_path')
+                    AND column_name IN ('full_path', 'parent_id', 'depth', 'folder_path', 'folder_id')
                     """
                 )
             )
@@ -60,7 +63,9 @@ async def bootstrap_folder_hierarchy(engine: AsyncEngine, logger) -> None:
                         'uq_folders_app_full_path',
                         'uq_folders_owner_full_path',
                         'idx_doc_folder_path',
+                        'idx_doc_folder_id',
                         'idx_documents_app_folder_path',
+                        'idx_documents_app_folder_id',
                         'idx_graph_folder_path',
                         'idx_graphs_app_folder_path'
                     )
@@ -102,6 +107,7 @@ async def bootstrap_folder_hierarchy(engine: AsyncEngine, logger) -> None:
         "ALTER TABLE folders ADD COLUMN IF NOT EXISTS parent_id TEXT",
         "ALTER TABLE folders ADD COLUMN IF NOT EXISTS depth INTEGER",
         "ALTER TABLE documents ADD COLUMN IF NOT EXISTS folder_path TEXT",
+        "ALTER TABLE documents ADD COLUMN IF NOT EXISTS folder_id TEXT",
         "ALTER TABLE graphs ADD COLUMN IF NOT EXISTS folder_path TEXT",
     ]
 
@@ -118,7 +124,9 @@ async def bootstrap_folder_hierarchy(engine: AsyncEngine, logger) -> None:
             "ON folders (owner_id, full_path) WHERE app_id IS NULL"
         ),
         "CREATE INDEX IF NOT EXISTS idx_doc_folder_path ON documents (folder_path)",
+        "CREATE INDEX IF NOT EXISTS idx_doc_folder_id ON documents (folder_id)",
         "CREATE INDEX IF NOT EXISTS idx_documents_app_folder_path ON documents (app_id, folder_path)",
+        "CREATE INDEX IF NOT EXISTS idx_documents_app_folder_id ON documents (app_id, folder_id)",
         "CREATE INDEX IF NOT EXISTS idx_graph_folder_path ON graphs (folder_path)",
         "CREATE INDEX IF NOT EXISTS idx_graphs_app_folder_path ON graphs (app_id, folder_path)",
     ]
diff --git a/core/database/postgres_database.py b/core/database/postgres_database.py
@@ -49,6 +49,7 @@ class DocumentModel(Base):
     app_id = Column(String)
     folder_name = Column(String)
     folder_path = Column(String)
+    folder_id = Column(String)
     end_user_id = Column(String)
 
     # Create indexes
@@ -59,12 +60,14 @@ class DocumentModel(Base):
         Index("idx_doc_app_id", "app_id"),
         Index("idx_doc_folder_name", "folder_name"),
         Index("idx_doc_folder_path", "folder_path"),
+        Index("idx_doc_folder_id", "folder_id"),
         Index("idx_doc_end_user_id", "end_user_id"),
         Index("idx_doc_owner_id", "owner_id"),
         # Composite indexes for common query patterns
         Index("idx_documents_owner_app", "owner_id", "app_id"),
         Index("idx_documents_app_folder", "app_id", "folder_name"),
         Index("idx_documents_app_folder_path", "app_id", "folder_path"),
+        Index("idx_documents_app_folder_id", "app_id", "folder_id"),
         Index("idx_documents_app_end_user", "app_id", "end_user_id"),
     )
 
@@ -368,8 +371,10 @@ async def store_document(self, document: Document, auth: AuthContext) -> bool:
             normalized_metadata, normalized_types = normalize_metadata(metadata, metadata_type_hints)
             doc_dict["doc_metadata"] = normalized_metadata
             doc_dict["metadata_types"] = normalized_types
-            # Mirror folder_name into doc_metadata for convenience in downstream filters (allow clearing)
-            doc_dict["doc_metadata"]["folder_name"] = doc_dict.get("folder_name")
+            # Mirror folder path into doc_metadata for convenience in downstream filters (allow clearing)
+            path_for_metadata = doc_dict.get("folder_path") or doc_dict.get("folder_name")
+            doc_dict["doc_metadata"]["folder_name"] = path_for_metadata
+            doc_dict["folder_id"] = doc_dict.get("folder_id")
 
             # Keep folder_path in sync with folder_name for backward compatibility
             folder_name_value = doc_dict.get("folder_name")
@@ -838,7 +843,13 @@ async def update_document(self, document_id: str, updates: Dict[str, Any], auth:
 
                     # Keep doc_metadata.folder_name in sync with the flattened column (support clearing)
                     if "doc_metadata" in updates:
-                        folder_value = folder_value_for_metadata if "folder_name" in updates else doc_model.folder_name
+                        folder_value = updates.get("folder_path")
+                        if folder_value is None:
+                            folder_value = (
+                                folder_value_for_metadata if "folder_name" in updates else doc_model.folder_path
+                            )
+                        if folder_value is None:
+                            folder_value = doc_model.folder_name
                         try:
                             if isinstance(updates["doc_metadata"], dict):
                                 updates["doc_metadata"]["folder_name"] = folder_value
@@ -1256,6 +1267,7 @@ def _document_model_to_dict(self, doc_model) -> Dict[str, Any]:
             # Include flattened fields
             "folder_name": doc_model.folder_name,
             "folder_path": doc_model.folder_path,
+            "folder_id": doc_model.folder_id,
             "app_id": doc_model.app_id,
             "end_user_id": doc_model.end_user_id,
         }
diff --git a/core/models/documents.py b/core/models/documents.py
@@ -54,6 +54,7 @@ class Document(BaseModel):
     end_user_id: Optional[str] = None
     app_id: Optional[str] = None
     folder_path: Optional[str] = None
+    folder_id: Optional[str] = None
 
     # Ensure storage_info values are strings to maintain backward compatibility
     @field_validator("storage_info", mode="before")
diff --git a/core/models/request.py b/core/models/request.py
@@ -101,19 +101,6 @@ class FolderDetailsRequest(BaseModel):
     )
 
 
-class FolderTreeRequest(BaseModel):
-    """Request model for retrieving a hierarchical folder tree with documents."""
-
-    folder_path: Optional[str] = Field(
-        default=None,
-        description="Base folder path to return. Use '/' or null for the full hierarchy.",
-    )
-    document_fields: Optional[List[str]] = Field(
-        default=None,
-        description="Optional list of fields to include for documents in each folder node (dot notation supported).",
-    )
-
-
 class SearchDocumentsRequest(BaseModel):
     """Request model for searching documents by name"""
 
diff --git a/core/models/responses.py b/core/models/responses.py
@@ -131,21 +131,6 @@ class FolderDetailsResponse(BaseModel):
     folders: List[FolderDetails]
 
 
-class FolderTreeNode(BaseModel):
-    """Nested folder tree entry including contained documents."""
-
-    id: Optional[str] = None
-    name: Optional[str] = None
-    full_path: Optional[str] = None
-    description: Optional[str] = None
-    depth: Optional[int] = None
-    documents: List[Dict[str, Any]] = Field(default_factory=list)
-    children: List["FolderTreeNode"] = Field(default_factory=list)
-
-
-FolderTreeNode.model_rebuild()
-
-
 class RequeueIngestionResult(BaseModel):
     """Result information for an individual requeued ingestion job."""
 
diff --git a/core/routes/folders.py b/core/routes/folders.py
@@ -9,15 +9,14 @@
 from core.database.postgres_database import InvalidMetadataFilterError
 from core.models.auth import AuthContext
 from core.models.folders import Folder, FolderCreate, FolderSummary
-from core.models.request import FolderDetailsRequest, FolderTreeRequest
+from core.models.request import FolderDetailsRequest
 from core.models.responses import (
     DocumentAddToFolderResponse,
     DocumentDeleteResponse,
     FolderDeleteResponse,
     FolderDetails,
     FolderDetailsResponse,
     FolderDocumentInfo,
-    FolderTreeNode,
 )
 from core.routes.utils import project_document_fields
 from core.services.telemetry import TelemetryService
@@ -264,176 +263,6 @@ async def list_folder_summaries(auth: AuthContext = Depends(verify_token)) -> Li
         raise HTTPException(status_code=500, detail=str(exc))
 
 
-@router.post("/tree", response_model=FolderTreeNode)
-async def get_folder_tree(
-    request: FolderTreeRequest,
-    auth: AuthContext = Depends(verify_token),
-) -> FolderTreeNode:
-    """
-    Return a hierarchical folder tree (with documents) rooted at ``folder_path``.
-
-    When ``folder_path`` is null or ``/``, the entire accessible hierarchy is returned.
-    """
-
-    try:
-        folder_path = request.folder_path
-        document_fields = request.document_fields
-        normalized_path: Optional[str] = None
-        if folder_path is not None:
-            if isinstance(folder_path, str) and folder_path.lower() == "null":
-                folder_path = None
-            else:
-                try:
-                    normalized_path = normalize_folder_path(folder_path)
-                except ValueError as exc:
-                    raise HTTPException(status_code=400, detail=str(exc))
-                if normalized_path == "/":
-                    normalized_path = None
-
-        base_path = normalized_path or "/"
-
-        base_folder: Optional[Folder] = None
-        if normalized_path:
-            base_folder = await document_service.db.get_folder_by_full_path(normalized_path, auth)
-            if not base_folder:
-                raise HTTPException(status_code=404, detail=f"Folder {folder_path} not found")
-
-        def _canonical_folder_path(folder: Folder) -> Optional[str]:
-            if folder.full_path:
-                try:
-                    return normalize_folder_path(folder.full_path)
-                except ValueError:
-                    return None
-            if folder.name:
-                try:
-                    return normalize_folder_path(folder.name)
-                except ValueError:
-                    return None
-            return None
-
-        def _parent_path(path: str) -> Optional[str]:
-            if not path or path == "/":
-                return None
-            segments = [part for part in path.strip("/").split("/") if part]
-            if len(segments) <= 1:
-                return "/"
-            return "/" + "/".join(segments[:-1])
-
-        def _attach_child(parent: FolderTreeNode, child: FolderTreeNode) -> None:
-            if not any(existing.full_path == child.full_path for existing in parent.children):
-                parent.children.append(child)
-
-        def _make_node(path: str, folder: Optional[Folder]) -> FolderTreeNode:
-            name = folder.name if folder else ("/" if path == "/" else (path.strip("/").split("/")[-1] or "/"))
-            depth = folder.depth if folder else (0 if path == "/" else None)
-            return FolderTreeNode(
-                id=folder.id if folder else None,
-                name=name,
-                full_path=path,
-                description=folder.description if folder else None,
-                depth=depth,
-                documents=[],
-                children=[],
-            )
-
-        all_folders = await document_service.db.list_folders(auth)
-        folders_with_paths: List[tuple[str, Folder]] = []
-        for folder in all_folders:
-            path = _canonical_folder_path(folder)
-            if path:
-                folders_with_paths.append((path, folder))
-
-        if normalized_path:
-            scoped = []
-            prefix = normalized_path.rstrip("/") + "/"
-            for path, folder in folders_with_paths:
-                if path == normalized_path or path.startswith(prefix):
-                    scoped.append((path, folder))
-            folders_with_paths = scoped
-            if base_folder:
-                base_folder_path = _canonical_folder_path(base_folder)
-                if base_folder_path and all(path != base_folder_path for path, _ in folders_with_paths):
-                    folders_with_paths.append((base_folder_path, base_folder))
-
-        nodes_by_path: Dict[str, FolderTreeNode] = {
-            path: _make_node(path, folder) for path, folder in folders_with_paths
-        }
-
-        root_node = nodes_by_path.get(base_path)
-        if not root_node:
-            root_node = _make_node(base_path, base_folder)
-            nodes_by_path[base_path] = root_node
-
-        for path in sorted(nodes_by_path.keys(), key=lambda p: (p.count("/"), p)):
-            if path == base_path:
-                continue
-            node = nodes_by_path[path]
-            parent_path = _parent_path(path)
-            parent_node = nodes_by_path.get(parent_path)
-            if not parent_node:
-                parent_node = root_node
-            _attach_child(parent_node, node)
-
-        doc_system_filters = {"folder_path_prefix": base_path} if normalized_path else None
-        document_result = await document_service.db.list_documents_flexible(
-            auth=auth,
-            skip=0,
-            limit=None,
-            system_filters=doc_system_filters,
-            include_total_count=False,
-            include_status_counts=False,
-            include_folder_counts=False,
-            return_documents=True,
-            sort_by="filename",
-            sort_direction="asc",
-        )
-
-        documents = document_result.get("documents", []) or []
-        for document in documents:
-            if hasattr(document, "model_dump"):
-                doc_dict = document.model_dump(mode="json")
-            elif hasattr(document, "dict"):
-                doc_dict = document.dict()
-            else:
-                doc_dict = dict(document)
-
-            doc_path_raw = doc_dict.get("folder_path")
-            try:
-                doc_path = normalize_folder_path(doc_path_raw) if doc_path_raw is not None else None
-            except ValueError:
-                doc_path = doc_path_raw
-
-            target_path = doc_path or base_path
-            target_node = nodes_by_path.get(target_path)
-            if not target_node:
-                target_node = _make_node(target_path, None)
-                nodes_by_path[target_path] = target_node
-                parent_path = _parent_path(target_path or "/")
-                parent_node = nodes_by_path.get(parent_path) or root_node
-                _attach_child(parent_node, target_node)
-
-            projected_doc = project_document_fields(doc_dict, document_fields)
-            if doc_path_raw is not None and "folder_path" not in projected_doc:
-                projected_doc["folder_path"] = doc_path_raw
-
-            target_node.documents.append(projected_doc)
-
-        def _sort_tree(node: FolderTreeNode) -> None:
-            node.children.sort(key=lambda child: (child.name or "", child.full_path or ""))
-            for child in node.children:
-                _sort_tree(child)
-            node.documents.sort(key=lambda doc: str(doc.get("filename") or doc.get("external_id") or ""))
-
-        _sort_tree(root_node)
-        return root_node
-
-    except HTTPException:
-        raise
-    except Exception as exc:  # noqa: BLE001
-        logger.error("Error building folder tree: %s", exc)
-        raise HTTPException(status_code=500, detail=str(exc))
-
-
 @router.post("/{folder_id_or_name:path}/documents/{document_id}", response_model=DocumentAddToFolderResponse)
 @telemetry.track(operation_type="add_document_to_folder", metadata_resolver=telemetry.add_document_to_folder_metadata)
 async def add_document_to_folder(
diff --git a/core/routes/ingest.py b/core/routes/ingest.py
diff --git a/core/services/ingestion_service.py b/core/services/ingestion_service.py
diff --git a/core/workers/ingestion_worker.py b/core/workers/ingestion_worker.py
