Folder extraction (#92)

core/api.py

@@ -13,7 +13,7 @@ import logging
 import arq
 from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
 from core.limits_utils import check_and_increment_limits
-from core.models.request import GenerateUriRequest, RetrieveRequest, CompletionQueryRequest, IngestTextRequest, CreateGraphRequest, UpdateGraphRequest, BatchIngestResponse
+from core.models.request import GenerateUriRequest, RetrieveRequest, CompletionQueryRequest, IngestTextRequest, CreateGraphRequest, UpdateGraphRequest, BatchIngestResponse, SetFolderRuleRequest
 from core.models.completion import ChunkSource, CompletionResponse
 from core.models.documents import Document, DocumentResult, ChunkResult
 from core.models.graph import Graph
@@ -32,6 +32,7 @@ from core.storage.local_storage import LocalStorage
 from core.reranker.flag_reranker import FlagReranker
 from core.cache.llama_cache_factory import LlamaCacheFactory
 import tomli
+from pydantic import BaseModel
 
 # Initialize FastAPI app
 app = FastAPI(title="Morphik API")
@@ -1998,3 +1999,230 @@ async def generate_cloud_uri(
     except Exception as e:
         logger.error(f"Error generating cloud URI: {e}")
         raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/folders/{folder_id}/set_rule")
+async def set_folder_rule(
+    folder_id: str,
+    request: SetFolderRuleRequest,
+    auth: AuthContext = Depends(verify_token),
+    apply_to_existing: bool = True,
+):
+    """
+    Set extraction rules for a folder.
+    
+    Args:
+        folder_id: ID of the folder to set rules for
+        request: SetFolderRuleRequest containing metadata extraction rules
+        auth: Authentication context
+        apply_to_existing: Whether to apply rules to existing documents in the folder
+        
+    Returns:
+        Success status with processing results
+    """
+    # Import text here to ensure it's available in this function's scope
+    from sqlalchemy import text
+    try:
+        async with telemetry.track_operation(
+            operation_type="set_folder_rule",
+            user_id=auth.entity_id,
+            metadata={
+                "folder_id": folder_id,
+                "rule_count": len(request.rules),
+                "apply_to_existing": apply_to_existing,
+            },
+        ):
+            # Log detailed information about the rules
+            logger.debug(f"Setting rules for folder {folder_id}")
+            logger.debug(f"Number of rules: {len(request.rules)}")
+            
+            for i, rule in enumerate(request.rules):
+                logger.debug(f"\nRule {i + 1}:")
+                logger.debug(f"Type: {rule.type}")
+                logger.debug("Schema:")
+                for field_name, field_config in rule.schema.items():
+                    logger.debug(f"  Field: {field_name}")
+                    logger.debug(f"    Type: {field_config.get('type', 'unknown')}")
+                    logger.debug(f"    Description: {field_config.get('description', 'No description')}")
+                    if 'schema' in field_config:
+                        logger.debug(f"    Has JSON schema: Yes")
+                        logger.debug(f"    Schema: {field_config['schema']}")
+            
+            # Get the folder
+            folder = await document_service.db.get_folder(folder_id, auth)
+            if not folder:
+                raise HTTPException(status_code=404, detail=f"Folder {folder_id} not found")
+                
+            # Check if user has write access to the folder
+            if not document_service.db._check_folder_access(folder, auth, "write"):
+                raise HTTPException(status_code=403, detail="You don't have write access to this folder")
+                
+            # Update folder with rules
+            # Convert rules to dicts for JSON serialization
+            rules_dicts = [rule.model_dump() for rule in request.rules]
+            
+            # Update the folder in the database
+            async with document_service.db.async_session() as session:
+                # Execute update query
+                await session.execute(
+                    text(
+                        """
+                        UPDATE folders 
+                        SET rules = :rules
+                        WHERE id = :folder_id
+                        """
+                    ),
+                    {"folder_id": folder_id, "rules": json.dumps(rules_dicts)},
+                )
+                await session.commit()
+                
+            logger.info(f"Successfully updated folder {folder_id} with {len(request.rules)} rules")
+            
+            # Get updated folder
+            updated_folder = await document_service.db.get_folder(folder_id, auth)
+            
+            # If apply_to_existing is True, apply these rules to all existing documents in the folder
+            processing_results = {"processed": 0, "errors": []}
+            
+            if apply_to_existing and folder.document_ids:
+                logger.info(f"Applying rules to {len(folder.document_ids)} existing documents in folder")
+                
+                # Import rules processor
+                from core.services.rules_processor import RulesProcessor
+                rules_processor = RulesProcessor()
+                
+                # Get all documents in the folder
+                documents = await document_service.db.get_documents_by_id(folder.document_ids, auth)
+                
+                # Process each document
+                for doc in documents:
+                    try:
+                        # Get document content
+                        logger.info(f"Processing document {doc.external_id}")
+                        
+                        # For each document, apply the rules from the folder
+                        doc_content = None
+                        
+                        # Get content from system_metadata if available
+                        if doc.system_metadata and "content" in doc.system_metadata:
+                            doc_content = doc.system_metadata["content"]
+                            logger.info(f"Retrieved content from system_metadata for document {doc.external_id}")
+                        
+                        # If we still have no content, log error and continue
+                        if not doc_content:
+                            error_msg = f"No content found in system_metadata for document {doc.external_id}"
+                            logger.error(error_msg)
+                            processing_results["errors"].append({
+                                "document_id": doc.external_id,
+                                "error": error_msg
+                            })
+                            continue
+                            
+                        # Process document with rules
+                        try:
+                            # Convert request rules to actual rule models and apply them
+                            from core.models.rules import MetadataExtractionRule
+                            
+                            for rule_request in request.rules:
+                                if rule_request.type == "metadata_extraction":
+                                    # Create the actual rule model
+                                    rule = MetadataExtractionRule(
+                                        type=rule_request.type,
+                                        schema=rule_request.schema
+                                    )
+                                    
+                                    # Apply the rule with retries
+                                    max_retries = 3
+                                    base_delay = 1  # seconds
+                                    extracted_metadata = None
+                                    last_error = None
+
+                                    for retry_count in range(max_retries):
+                                        try:
+                                            if retry_count > 0:
+                                                # Exponential backoff
+                                                delay = base_delay * (2 ** (retry_count - 1))
+                                                logger.info(f"Retry {retry_count}/{max_retries} after {delay}s delay")
+                                                await asyncio.sleep(delay)
+                                            
+                                            extracted_metadata, _ = await rule.apply(doc_content)
+                                            logger.info(f"Successfully extracted metadata on attempt {retry_count + 1}: {extracted_metadata}")
+                                            break  # Success, exit retry loop
+                                            
+                                        except Exception as rule_apply_error:
+                                            last_error = rule_apply_error
+                                            logger.warning(f"Metadata extraction attempt {retry_count + 1} failed: {rule_apply_error}")
+                                            if retry_count == max_retries - 1:  # Last attempt
+                                                logger.error(f"All {max_retries} metadata extraction attempts failed")
+                                                processing_results["errors"].append({
+                                                    "document_id": doc.external_id,
+                                                    "error": f"Failed to extract metadata after {max_retries} attempts: {str(last_error)}"
+                                                })
+                                                continue  # Skip to next document
+                                    
+                                    # Update document metadata if extraction succeeded
+                                    if extracted_metadata:
+                                        # Merge new metadata with existing
+                                        doc.metadata.update(extracted_metadata)
+                                        
+                                        # Create an updates dict that only updates metadata
+                                        # We need to create system_metadata with all preserved fields
+                                        # Note: In the database, metadata is stored as 'doc_metadata', not 'metadata'
+                                        updates = {
+                                            "doc_metadata": doc.metadata,  # Use doc_metadata for the database
+                                            "system_metadata": {}  # Will be merged with existing in update_document
+                                        }
+                                        
+                                        # Explicitly preserve the content field in system_metadata
+                                        if "content" in doc.system_metadata:
+                                            updates["system_metadata"]["content"] = doc.system_metadata["content"]
+                                        
+                                        # Log the updates we're making
+                                        logger.info(f"Updating document {doc.external_id} with metadata: {extracted_metadata}")
+                                        logger.info(f"Full metadata being updated: {doc.metadata}")
+                                        logger.info(f"Update object being sent to database: {updates}")
+                                        logger.info(f"Preserving content in system_metadata: {'content' in doc.system_metadata}")
+                                        
+                                        # Update document in database
+                                        success = await document_service.db.update_document(
+                                            doc.external_id,
+                                            updates,
+                                            auth
+                                        )
+                                        
+                                        if success:
+                                            logger.info(f"Updated metadata for document {doc.external_id}")
+                                            processing_results["processed"] += 1
+                                        else:
+                                            logger.error(f"Failed to update metadata for document {doc.external_id}")
+                                            processing_results["errors"].append({
+                                                "document_id": doc.external_id,
+                                                "error": "Failed to update document metadata"
+                                            })
+                        except Exception as rule_error:
+                            logger.error(f"Error processing rules for document {doc.external_id}: {rule_error}")
+                            processing_results["errors"].append({
+                                "document_id": doc.external_id,
+                                "error": f"Error processing rules: {str(rule_error)}"
+                            })
+                            
+                    except Exception as doc_error:
+                        logger.error(f"Error processing document {doc.external_id}: {doc_error}")
+                        processing_results["errors"].append({
+                            "document_id": doc.external_id,
+                            "error": str(doc_error)
+                        })
+            
+            return {
+                "status": "success",
+                "message": "Rules set successfully",
+                "folder_id": folder_id,
+                "rules": updated_folder.rules,
+                "processing_results": processing_results
+            }
+    except HTTPException as e:
+        # Re-raise HTTP exceptions
+        raise
+    except Exception as e:
+        logger.error(f"Error setting folder rules: {e}")
+        raise HTTPException(status_code=500, detail=str(e))

core/database/postgres_database.py

@@ -81,6 +81,7 @@ class FolderModel(Base):
     document_ids = Column(JSONB, default=list)
     system_metadata = Column(JSONB, default=dict)
     access_control = Column(JSONB, default=dict)
+    rules = Column(JSONB, default=list)
 
     # Create indexes
     __table_args__ = (
@@ -220,6 +221,28 @@ class PostgresDatabase(BaseDatabase):
                     )
                 )
                 
+                # Add rules column to folders table if it doesn't exist
+                result = await conn.execute(
+                    text(
+                        """
+                    SELECT column_name 
+                    FROM information_schema.columns 
+                    WHERE table_name = 'folders' AND column_name = 'rules'
+                    """
+                    )
+                )
+                if not result.first():
+                    # Add rules column to folders table
+                    await conn.execute(
+                        text(
+                            """
+                        ALTER TABLE folders 
+                        ADD COLUMN IF NOT EXISTS rules JSONB DEFAULT '[]'::jsonb
+                        """
+                        )
+                    )
+                    logger.info("Added rules column to folders table")
+                
                 # Create indexes for folders table
                 await conn.execute(text("CREATE INDEX IF NOT EXISTS idx_folder_name ON folders (name);"))
                 await conn.execute(text("CREATE INDEX IF NOT EXISTS idx_folder_owner ON folders USING gin (owner);"))
@@ -557,14 +580,17 @@ class PostgresDatabase(BaseDatabase):
             # Update system metadata
             updates.setdefault("system_metadata", {})
             
-            # Preserve folder_name and end_user_id if not explicitly overridden
+            # Merge with existing system_metadata instead of just preserving specific fields
             if existing_doc.system_metadata:
-                if "folder_name" in existing_doc.system_metadata and "folder_name" not in updates["system_metadata"]:
+                # Start with existing system_metadata
-                    updates["system_metadata"]["folder_name"] = existing_doc.system_metadata["folder_name"]
+                merged_system_metadata = dict(existing_doc.system_metadata)
-                
+                # Update with new values
-                if "end_user_id" in existing_doc.system_metadata and "end_user_id" not in updates["system_metadata"]:
+                merged_system_metadata.update(updates["system_metadata"])
-                    updates["system_metadata"]["end_user_id"] = existing_doc.system_metadata["end_user_id"]
+                # Replace with merged result
+                updates["system_metadata"] = merged_system_metadata
+                logger.debug(f"Merged system_metadata during document update, preserving existing fields")
             
+            # Always update the updated_at timestamp
             updates["system_metadata"]["updated_at"] = datetime.now(UTC)
 
             # Serialize datetime objects to ISO format strings
@@ -577,9 +603,21 @@ class PostgresDatabase(BaseDatabase):
                 doc_model = result.scalar_one_or_none()
 
                 if doc_model:
+                    # Log what we're updating
+                    logger.info(f"Document update: updating fields {list(updates.keys())}")
+                    
+                    # Special handling for metadata/doc_metadata conversion
+                    if "metadata" in updates and "doc_metadata" not in updates:
+                        logger.info("Converting 'metadata' to 'doc_metadata' for database update")
+                        updates["doc_metadata"] = updates.pop("metadata")
+                    
+                    # Set all attributes
                     for key, value in updates.items():
+                        logger.debug(f"Setting document attribute {key} = {value}")
                         setattr(doc_model, key, value)
+                        
                     await session.commit()
+                    logger.info(f"Document {document_id} updated successfully")
                     return True
                 return False
 
@@ -1108,7 +1146,8 @@ class PostgresDatabase(BaseDatabase):
                     owner=folder_dict["owner"],
                     document_ids=folder_dict.get("document_ids", []),
                     system_metadata=folder_dict.get("system_metadata", {}),
-                    access_control=access_control
+                    access_control=access_control,
+                    rules=folder_dict.get("rules", [])
                 )
                 
                 session.add(folder_model)
@@ -1144,7 +1183,8 @@ class PostgresDatabase(BaseDatabase):
                     "owner": folder_model.owner,
                     "document_ids": folder_model.document_ids,
                     "system_metadata": folder_model.system_metadata,
-                    "access_control": folder_model.access_control
+                    "access_control": folder_model.access_control,
+                    "rules": folder_model.rules
                 }
                 
                 folder = Folder(**folder_dict)
@@ -1190,7 +1230,8 @@ class PostgresDatabase(BaseDatabase):
                             "owner": folder_row.owner,
                             "document_ids": folder_row.document_ids,
                             "system_metadata": folder_row.system_metadata,
-                            "access_control": folder_row.access_control
+                            "access_control": folder_row.access_control,
+                            "rules": folder_row.rules
                         }
                         
                         return Folder(**folder_dict)
@@ -1210,7 +1251,8 @@ class PostgresDatabase(BaseDatabase):
                         "owner": folder_model.owner,
                         "document_ids": folder_model.document_ids,
                         "system_metadata": folder_model.system_metadata,
-                        "access_control": folder_model.access_control
+                        "access_control": folder_model.access_control,
+                        "rules": folder_model.rules
                     }
                     
                     folder = Folder(**folder_dict)
@@ -1244,7 +1286,8 @@ class PostgresDatabase(BaseDatabase):
                         "owner": folder_model.owner,
                         "document_ids": folder_model.document_ids,
                         "system_metadata": folder_model.system_metadata,
-                        "access_control": folder_model.access_control
+                        "access_control": folder_model.access_control,
+                        "rules": folder_model.rules
                     }
                     
                     folder = Folder(**folder_dict)

core/models/folders.py

@@ -21,6 +21,7 @@ class Folder(BaseModel):
     access_control: Dict[str, List[str]] = Field(
         default_factory=lambda: {"readers": [], "writers": [], "admins": []}
     )
+    rules: List[Dict[str, Any]] = Field(default_factory=list)
 
     def __hash__(self):
         return hash(self.id)

core/models/request.py

@@ -109,3 +109,14 @@ class GenerateUriRequest(BaseModel):
     name: str = Field(..., description="Name of the application")
     user_id: str = Field(..., description="ID of the user who owns the app")
     expiry_days: int = Field(default=30, description="Number of days until the token expires")
+
+
+# Add these classes before the extract_folder_data endpoint
+class MetadataExtractionRuleRequest(BaseModel):
+    """Request model for metadata extraction rule"""
+    type: str = "metadata_extraction"  # Only metadata_extraction supported for now
+    schema: Dict[str, Any]
+
+class SetFolderRuleRequest(BaseModel):
+    """Request model for setting folder rules"""
+    rules: List[MetadataExtractionRuleRequest]

core/models/rules.py

@@ -74,14 +74,28 @@ class MetadataExtractionRule(BaseRule):
         # Create the dynamic model
         DynamicMetadataModel = create_model("DynamicMetadataModel", **field_definitions)
 
+        # Create a more explicit instruction that clearly shows expected output format
+        schema_descriptions = []
+        for field_name, field_config in self.schema.items():
+            field_type = field_config.get("type", "string") if isinstance(field_config, dict) else "string"
+            description = field_config.get("description", "No description") if isinstance(field_config, dict) else field_config
+            schema_descriptions.append(f"- {field_name}: {description} (type: {field_type})")
+        
+        schema_text = "\n".join(schema_descriptions)
+
         prompt = f"""
         Extract metadata from the following text according to this schema:
-        {self.schema}
+        
-
+        {schema_text}
+        
         Text to extract from:
         {content}
-
+        
-        Extract all relevant information that matches the schema.
+        Follow these guidelines:
+        1. Extract all requested information as simple strings, numbers, or booleans (not as objects or nested structures)
+        2. If information is not present, indicate this with null instead of making something up
+        3. Answer directly with the requested information - don't include explanations or reasoning
+        4. Be concise but accurate in your extractions
         """
 
         # Get the model configuration from registered_models
@@ -93,7 +107,7 @@ class MetadataExtractionRule(BaseRule):
 
         system_message = {
             "role": "system",
-            "content": "You are a metadata extraction assistant. Extract structured metadata from text precisely following the provided schema.",
+            "content": "You are a metadata extraction assistant. Extract structured metadata from text precisely following the provided schema. Always return the metadata as direct values (strings, numbers, booleans), not as objects with additional properties.",
         }
 
         user_message = {"role": "user", "content": prompt}

morphik.toml

@@ -15,6 +15,7 @@ dev_permissions = ["read", "write", "admin"]  # Default dev permissions
 # OpenAI models
 openai_gpt4o = { model_name = "gpt-4o", vision = true }
 openai_gpt4 = { model_name = "gpt-4" }
+openai_gpt4o_extraction = { model_name = "gpt-4o" }
 
 # Azure OpenAI models
 azure_gpt4 = { model_name = "gpt-4", api_base = "YOUR_AZURE_URL_HERE", api_version = "2023-05-15", deployment_id = "gpt-4-deployment" }

ui-component/app/page.tsx

@@ -1,8 +1,20 @@
 "use client";
 
-import React from 'react';
+import React, { Suspense } from 'react';
 import MorphikUI from '@/components/MorphikUI';
+import { useSearchParams } from 'next/navigation';
+
+function HomeContent() {
+  const searchParams = useSearchParams();
+  const folderParam = searchParams.get('folder');
+  
+  return <MorphikUI initialFolder={folderParam} />;
+}
 
 export default function Home() {
-  return <MorphikUI />;
+  return (
+    <Suspense fallback={<div>Loading...</div>}>
+      <HomeContent />
+    </Suspense>
+  );
 }

ui-component/components/MorphikUI.tsx

@@ -21,7 +21,8 @@ const MorphikUI: React.FC<MorphikUIProps> = ({
   apiBaseUrl = DEFAULT_API_BASE_URL,
   isReadOnlyUri = false, // Default to editable URI
   onUriChange,
-  onBackClick
+  onBackClick,
+  initialFolder = null
 }) => {
   // State to manage connectionUri internally if needed
   const [currentUri, setCurrentUri] = useState(connectionUri);
@@ -41,6 +42,7 @@ const MorphikUI: React.FC<MorphikUIProps> = ({
   };
   const [activeSection, setActiveSection] = useState('documents');
   const [selectedGraphName, setSelectedGraphName] = useState<string | undefined>(undefined);
+  const [isSidebarCollapsed, setIsSidebarCollapsed] = useState(false);
   
   // Extract auth token and API URL from connection URI if provided
   const authToken = currentUri ? extractTokenFromUri(currentUri) : null;
@@ -65,6 +67,8 @@ const MorphikUI: React.FC<MorphikUIProps> = ({
           connectionUri={currentUri}
           isReadOnlyUri={isReadOnlyUri}
           onUriChange={handleUriChange}
+          isCollapsed={isSidebarCollapsed}
+          setIsCollapsed={setIsSidebarCollapsed}
         />
         
         <div className="flex-1 flex flex-col h-screen overflow-hidden">
@@ -88,7 +92,9 @@ const MorphikUI: React.FC<MorphikUIProps> = ({
             {activeSection === 'documents' && (
               <DocumentsSection 
                 apiBaseUrl={effectiveApiBaseUrl} 
-                authToken={authToken} 
+                authToken={authToken}
+                initialFolder={initialFolder}
+                setSidebarCollapsed={setIsSidebarCollapsed}
               />
             )}
             

ui-component/components/documents/DocumentDetail.tsx

@@ -6,8 +6,9 @@ import { ScrollArea } from '@/components/ui/scroll-area';
 import { Button } from '@/components/ui/button';
 import { Accordion, AccordionContent, AccordionItem, AccordionTrigger } from '@/components/ui/accordion';
 import { Dialog, DialogContent, DialogDescription, DialogFooter, DialogHeader, DialogTitle, DialogTrigger } from '@/components/ui/dialog';
-import { Info, Folder as FolderIcon } from 'lucide-react';
+import { Info } from 'lucide-react';
 import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import Image from 'next/image';
 
 import { Document, Folder } from '@/components/types';
 
@@ -20,6 +21,7 @@ interface DocumentDetailProps {
   refreshDocuments: () => void;
   refreshFolders: () => void;
   loading: boolean;
+  onClose: () => void;
 }
 
 const DocumentDetail: React.FC<DocumentDetailProps> = ({
@@ -30,7 +32,8 @@ const DocumentDetail: React.FC<DocumentDetailProps> = ({
   authToken,
   refreshDocuments,
   refreshFolders,
-  loading
+  loading,
+  onClose
 }) => {
   const [isMovingToFolder, setIsMovingToFolder] = useState(false);
   
@@ -103,8 +106,20 @@ const DocumentDetail: React.FC<DocumentDetailProps> = ({
 
   return (
     <div className="border rounded-lg">
-      <div className="bg-muted px-4 py-3 border-b sticky top-0">
+      <div className="bg-muted px-4 py-3 border-b sticky top-0 flex justify-between items-center">
         <h3 className="text-lg font-semibold">Document Details</h3>
+        <Button
+          variant="ghost"
+          size="icon"
+          onClick={onClose}
+          className="rounded-full hover:bg-background/80"
+        >
+          <svg xmlns="http://www.w3.org/2000/svg" width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round">
+            <line x1="18" y1="6" x2="6" y2="18"></line>
+            <line x1="6" y1="6" x2="18" y2="18"></line>
+          </svg>
+          <span className="sr-only">Close panel</span>
+        </Button>
       </div>
       
       <ScrollArea className="h-[calc(100vh-200px)]">
@@ -122,7 +137,7 @@ const DocumentDetail: React.FC<DocumentDetailProps> = ({
           <div>
             <h3 className="font-medium mb-1">Folder</h3>
             <div className="flex items-center gap-2">
-              <FolderIcon className="h-4 w-4 text-muted-foreground" />
+              <Image src="/icons/folder-icon.png" alt="Folder" width={16} height={16} />
               <Select
                 value={currentFolder || "_none"}
                 onValueChange={(value) => handleMoveToFolder(value === "_none" ? null : value)}

ui-component/components/documents/DocumentList.tsx

@@ -1,11 +1,32 @@
 "use client";
 
-import React from 'react';
+import React, { useState } from 'react';
 import { Checkbox } from "@/components/ui/checkbox";
 import { Badge } from '@/components/ui/badge';
 import { ScrollArea } from '@/components/ui/scroll-area';
+import { Button } from '@/components/ui/button';
+import { Dialog, DialogContent, DialogDescription, DialogFooter, DialogHeader, DialogTitle } from '@/components/ui/dialog';
+import { Input } from '@/components/ui/input';
+import { Textarea } from '@/components/ui/textarea';
+import { Plus, Wand2, Upload } from 'lucide-react';
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import { showAlert } from '@/components/ui/alert-system';
 
-import { Document } from '@/components/types';
+import { Document, Folder } from '@/components/types';
+
+type ColumnType = 'string' | 'int' | 'float' | 'bool' | 'Date' | 'json';
+
+interface CustomColumn {
+  name: string;
+  description: string;
+  _type: ColumnType;
+  schema?: string;
+}
+
+interface MetadataExtractionRule {
+  type: "metadata_extraction";
+  schema: Record<string, unknown>;
+}
 
 interface DocumentListProps {
   documents: Document[];
@@ -15,9 +36,151 @@ interface DocumentListProps {
   handleCheckboxChange: (checked: boolean | "indeterminate", docId: string) => void;
   getSelectAllState: () => boolean | "indeterminate";
   setSelectedDocuments: (docIds: string[]) => void;
+  setDocuments: (docs: Document[]) => void;
   loading: boolean;
+  apiBaseUrl: string;
+  authToken: string | null;
+  selectedFolder?: string | null;
 }
 
+// Create a separate Column Dialog component to isolate its state
+const AddColumnDialog = ({ 
+  isOpen, 
+  onClose,
+  onAddColumn
+}: { 
+  isOpen: boolean; 
+  onClose: () => void;
+  onAddColumn: (column: CustomColumn) => void;
+}) => {
+  const [localColumnName, setLocalColumnName] = useState('');
+  const [localColumnDescription, setLocalColumnDescription] = useState('');
+  const [localColumnType, setLocalColumnType] = useState<ColumnType>('string');
+  const [localColumnSchema, setLocalColumnSchema] = useState<string>('');
+
+  const handleLocalSchemaFileChange = (e: React.ChangeEvent<HTMLInputElement>) => {
+    const file = e.target.files?.[0];
+    if (file) {
+      const reader = new FileReader();
+      reader.onload = (event) => {
+        setLocalColumnSchema(event.target?.result as string);
+      };
+      reader.readAsText(file);
+    }
+  };
+
+  const handleSubmit = (e: React.FormEvent) => {
+    e.preventDefault();
+    if (localColumnName.trim()) {
+      const column: CustomColumn = {
+        name: localColumnName.trim(),
+        description: localColumnDescription.trim(),
+        _type: localColumnType
+      };
+      
+      if (localColumnType === 'json' && localColumnSchema) {
+        column.schema = localColumnSchema;
+      }
+      
+      onAddColumn(column);
+      
+      // Reset form values
+      setLocalColumnName('');
+      setLocalColumnDescription('');
+      setLocalColumnType('string');
+      setLocalColumnSchema('');
+      
+      // Close the dialog
+      onClose();
+    }
+  };
+
+  return (
+    <Dialog open={isOpen} onOpenChange={(open) => !open && onClose()}>
+      <DialogContent onPointerDownOutside={(e) => e.preventDefault()}>
+        <form onSubmit={handleSubmit}>
+          <DialogHeader>
+            <DialogTitle>Add Custom Column</DialogTitle>
+            <DialogDescription>
+              Add a new column and specify its type and description.
+            </DialogDescription>
+          </DialogHeader>
+          <div className="space-y-4 py-4">
+            <div className="space-y-2">
+              <label htmlFor="column-name" className="text-sm font-medium">Column Name</label>
+              <Input
+                id="column-name"
+                placeholder="e.g. Author, Category, etc."
+                value={localColumnName}
+                onChange={(e) => setLocalColumnName(e.target.value)}
+                autoFocus
+              />
+            </div>
+            <div className="space-y-2">
+              <label htmlFor="column-type" className="text-sm font-medium">Type</label>
+              <Select 
+                value={localColumnType} 
+                onValueChange={(value) => setLocalColumnType(value as ColumnType)}
+              >
+                <SelectTrigger id="column-type">
+                  <SelectValue placeholder="Select data type" />
+                </SelectTrigger>
+                <SelectContent>
+                  <SelectItem value="string">String</SelectItem>
+                  <SelectItem value="int">Integer</SelectItem>
+                  <SelectItem value="float">Float</SelectItem>
+                  <SelectItem value="bool">Boolean</SelectItem>
+                  <SelectItem value="Date">Date</SelectItem>
+                  <SelectItem value="json">JSON</SelectItem>
+                </SelectContent>
+              </Select>
+            </div>
+            {localColumnType === 'json' && (
+              <div className="space-y-2">
+                <label htmlFor="column-schema" className="text-sm font-medium">JSON Schema</label>
+                <div className="flex items-center space-x-2">
+                  <Input
+                    id="column-schema-file"
+                    type="file"
+                    accept=".json"
+                    className="hidden"
+                    onChange={handleLocalSchemaFileChange}
+                  />
+                  <Button 
+                    type="button" 
+                    variant="outline" 
+                    onClick={() => document.getElementById('column-schema-file')?.click()}
+                    className="flex items-center gap-2"
+                  >
+                    <Upload className="h-4 w-4" />
+                    Upload Schema
+                  </Button>
+                  <span className="text-sm text-muted-foreground">
+                    {localColumnSchema ? 'Schema loaded' : 'No schema uploaded'}
+                  </span>
+                </div>
+              </div>
+            )}
+            <div className="space-y-2">
+              <label htmlFor="column-description" className="text-sm font-medium">Description</label>
+              <Textarea
+                id="column-description"
+                placeholder="Describe in natural language what information this column should contain..."
+                value={localColumnDescription}
+                onChange={(e) => setLocalColumnDescription(e.target.value)}
+              />
+            </div>
+          </div>
+          <DialogFooter>
+            <Button type="button" variant="outline" onClick={onClose}>Cancel</Button>
+            <Button type="submit">Add Column</Button>
+          </DialogFooter>
+        </form>
+      </DialogContent>
+    </Dialog>
+  );
+};
+
 const DocumentList: React.FC<DocumentListProps> = ({
   documents,
   selectedDocument,
@@ -26,38 +189,265 @@ const DocumentList: React.FC<DocumentListProps> = ({
   handleCheckboxChange,
   getSelectAllState,
   setSelectedDocuments,
-  loading
+  setDocuments,
+  loading,
+  apiBaseUrl,
+  authToken,
+  selectedFolder
 }) => {
-  if (loading && !documents.length) {
+  const [customColumns, setCustomColumns] = useState<CustomColumn[]>([]);
-    return <div className="text-center py-8 flex-1">Loading documents...</div>;
+  const [showAddColumnDialog, setShowAddColumnDialog] = useState(false);
-  }
+  const [isExtracting, setIsExtracting] = useState(false);
 
-  // Status badge helper component (used in the document list items)
+  // Get unique metadata fields from all documents
-  // Status rendering is handled inline in the component instead
+  const existingMetadataFields = React.useMemo(() => {
+    const fields = new Set<string>();
+    documents.forEach(doc => {
+      if (doc.metadata) {
+        Object.keys(doc.metadata).forEach(key => fields.add(key));
+      }
+    });
+    return Array.from(fields);
+  }, [documents]);
 
-  return (
+  // Combine existing metadata fields with custom columns
-    <div className="border rounded-md">
+  const allColumns = React.useMemo(() => {
-      <div className="bg-muted border-b p-3 font-medium sticky top-0">
+    const metadataColumns: CustomColumn[] = existingMetadataFields.map(field => ({
-        <div className="grid grid-cols-12">
+      name: field,
-          <div className="col-span-1 flex items-center justify-center">
+      description: `Extracted ${field}`,
-            <Checkbox
+      _type: 'string' // Default to string type for existing metadata
-              id="select-all-documents"
+    }));
-              checked={getSelectAllState()}
+    
-              onCheckedChange={(checked) => {
+    // Merge with custom columns, preferring custom column definitions if they exist
-                if (checked) {
+    const mergedColumns = [...metadataColumns];
-                  setSelectedDocuments(documents.map(doc => doc.external_id));
+    customColumns.forEach(customCol => {
-                } else {
+      const existingIndex = mergedColumns.findIndex(col => col.name === customCol.name);
-                  setSelectedDocuments([]);
+      if (existingIndex >= 0) {
+        mergedColumns[existingIndex] = customCol;
+      } else {
+        mergedColumns.push(customCol);
+      }
+    });
+    
+    return mergedColumns;
+  }, [existingMetadataFields, customColumns]);
+
+  const handleAddColumn = (column: CustomColumn) => {
+    setCustomColumns([...customColumns, column]);
+  };
+
+  // Handle data extraction
+  const handleExtract = async () => {
+    // First, find the folder object to get its ID
+    if (!selectedFolder || customColumns.length === 0) {
+      console.error("Cannot extract: No folder selected or no columns defined");
+      return;
+    }
+
+    // We need to get the folder ID for the API call
+    try {
+      setIsExtracting(true);
+
+      // First, get folders to find the current folder ID
+      const foldersResponse = await fetch(`${apiBaseUrl}/folders`, {
+        headers: authToken ? { 'Authorization': `Bearer ${authToken}` } : {}
+      });
+      
+      if (!foldersResponse.ok) {
+        throw new Error(`Failed to fetch folders: ${foldersResponse.statusText}`);
+      }
+      
+      const folders = await foldersResponse.json();
+      const currentFolder = folders.find((folder: Folder) => folder.name === selectedFolder);
+      
+      if (!currentFolder) {
+        throw new Error(`Folder "${selectedFolder}" not found`);
+      }
+      
+      // Convert columns to metadata extraction rule
+      const rule: MetadataExtractionRule = {
+        type: "metadata_extraction",
+        schema: Object.fromEntries(
+          customColumns.map(col => [
+            col.name,
+            {
+              type: col._type,
+              description: col.description,
+              ...(col.schema ? { schema: JSON.parse(col.schema) } : {})
+            }
+          ])
+        )
+      };
+      
+      // Set the rule
+      const setRuleResponse = await fetch(`${apiBaseUrl}/folders/${currentFolder.id}/set_rule`, {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+          ...(authToken ? { 'Authorization': `Bearer ${authToken}` } : {})
+        },
+        body: JSON.stringify({
+          rules: [rule]
+        })
+      });
+      
+      if (!setRuleResponse.ok) {
+        throw new Error(`Failed to set rule: ${setRuleResponse.statusText}`);
+      }
+      
+      const result = await setRuleResponse.json();
+      console.log("Rule set successfully:", result);
+      
+      // Show success message
+      showAlert("Extraction rule set successfully!", {
+        type: 'success',
+        duration: 3000
+      });
+
+      // Force a fresh refresh after setting the rule
+      // This is a special function to ensure we get truly fresh data
+      const refreshAfterRule = async () => {
+        try {
+          console.log("Performing fresh refresh after setting extraction rule");
+          // Clear folder data to force a clean refresh
+          const folderResponse = await fetch(`${apiBaseUrl}/folders`, {
+            method: 'GET',
+            headers: {
+              ...(authToken ? { 'Authorization': `Bearer ${authToken}` } : {})
+            }
+          });
+          
+          if (!folderResponse.ok) {
+            throw new Error(`Failed to fetch folders: ${folderResponse.statusText}`);
+          }
+          
+          const freshFolders = await folderResponse.json();
+          console.log(`Rule: Fetched ${freshFolders.length} folders with fresh data`);
+          
+          // Now fetch documents based on the current folder
+          if (selectedFolder && selectedFolder !== "all") {
+            // Find the folder by name
+            const targetFolder = freshFolders.find((folder: Folder) => folder.name === selectedFolder);
+            
+            if (targetFolder) {
+              console.log(`Rule: Found folder ${targetFolder.name} in fresh data`);
+              
+              // Get the document IDs from the folder
+              const documentIds = Array.isArray(targetFolder.document_ids) ? targetFolder.document_ids : [];
+              console.log(`Rule: Folder has ${documentIds.length} documents`);
+              
+              if (documentIds.length > 0) {
+                // Fetch document details for the IDs
+                const docResponse = await fetch(`${apiBaseUrl}/batch/documents`, {
+                  method: 'POST',
+                  headers: {
+                    'Content-Type': 'application/json',
+                    ...(authToken ? { 'Authorization': `Bearer ${authToken}` } : {})
+                  },
+                  body: JSON.stringify({
+                    document_ids: [...documentIds]
+                  })
+                });
+                
+                if (!docResponse.ok) {
+                  throw new Error(`Failed to fetch documents: ${docResponse.statusText}`);
                 }
-              }}
+                
-              aria-label="Select all documents"
+                const freshDocs = await docResponse.json();
-            />
+                console.log(`Rule: Fetched ${freshDocs.length} document details`);
+                
+                // Update documents state
+                setDocuments(freshDocs);
+              } else {
+                // Empty folder
+                setDocuments([]);
+              }
+            } else {
+              console.log(`Rule: Selected folder ${selectedFolder} not found in fresh data`);
+              setDocuments([]);
+            }
+          } else {
+            // For "all" documents view, fetch all documents
+            const allDocsResponse = await fetch(`${apiBaseUrl}/documents`, {
+              method: 'POST',
+              headers: {
+                'Content-Type': 'application/json',
+                ...(authToken ? { 'Authorization': `Bearer ${authToken}` } : {})
+              },
+              body: JSON.stringify({})
+            });
+            
+            if (!allDocsResponse.ok) {
+              throw new Error(`Failed to fetch all documents: ${allDocsResponse.statusText}`);
+            }
+            
+            const allDocs = await allDocsResponse.json();
+            console.log(`Rule: Fetched ${allDocs.length} documents for "all" view`);
+            setDocuments(allDocs);
+          }
+        } catch (err) {
+          console.error('Error refreshing after setting rule:', err);
+          showAlert('Error refreshing data after setting rule', {
+            type: 'error',
+            duration: 3000
+          });
+        }
+      };
+      
+      // Execute the refresh
+      await refreshAfterRule();
+    } catch (error) {
+      console.error("Error setting extraction rule:", error);
+      showAlert(`Failed to set extraction rule: ${error instanceof Error ? error.message : String(error)}`, {
+        type: 'error',
+        title: 'Error',
+        duration: 5000
+      });
+    } finally {
+      setIsExtracting(false);
+    }
+  };
+
+  const DocumentListHeader = () => (
+    <div className="bg-muted border-b font-medium sticky top-0 z-10 relative">
+      <div className="grid items-center w-full" style={{ 
+        gridTemplateColumns: `48px minmax(200px, 350px) 100px 120px ${allColumns.map(() => '140px').join(' ')}` 
+      }}>
+        <div className="flex items-center justify-center p-3">
+          <Checkbox
+            id="select-all-documents"
+            checked={getSelectAllState()}
+            onCheckedChange={(checked) => {
+              if (checked) {
+                setSelectedDocuments(documents.map(doc => doc.external_id));
+              } else {
+                setSelectedDocuments([]);
+              }
+            }}
+            aria-label="Select all documents"
+          />
+        </div>
+        <div className="text-sm font-semibold p-3">Filename</div>
+        <div className="text-sm font-semibold p-3">Type</div>
+        <div className="text-sm font-semibold p-3">
+          <div className="group relative inline-flex items-center">
+            Status
+            <span className="ml-1 text-muted-foreground cursor-help">
+              <svg xmlns="http://www.w3.org/2000/svg" width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round">
+                <circle cx="12" cy="12" r="10"></circle>
+                <line x1="12" y1="16" x2="12" y2="12"></line>
+                <line x1="12" y1="8" x2="12.01" y2="8"></line>
+              </svg>
+            </span>
+            <div className="absolute left-0 top-6 hidden group-hover:block bg-background border text-foreground text-xs p-3 rounded-md w-64 z-[100] shadow-lg">
+              Documents with &quot;Processing&quot; status are queryable, but visual features like direct visual context will only be available after processing completes.
+            </div>
           </div>
-          <div className="col-span-4">Filename</div>
+        </div>
-          <div className="col-span-2">Type</div>
+        {allColumns.map((column) => (
-          <div className="col-span-2">
+          <div key={column.name} className="text-sm font-semibold p-3">
             <div className="group relative inline-flex items-center">
-              Status
+              {column.name}
               <span className="ml-1 text-muted-foreground cursor-help">
                 <svg xmlns="http://www.w3.org/2000/svg" width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round">
                   <circle cx="12" cy="12" r="10"></circle>
@@ -65,25 +455,72 @@ const DocumentList: React.FC<DocumentListProps> = ({
                   <line x1="12" y1="8" x2="12.01" y2="8"></line>
                 </svg>
               </span>
-              <div className="absolute left-0 -top-24 hidden group-hover:block bg-gray-800 text-white text-xs p-2 rounded w-60 z-50 shadow-lg">
+              <div className="absolute left-0 top-6 hidden group-hover:block bg-background border text-foreground text-xs p-3 rounded-md w-64 z-[100] shadow-lg">
-                Documents with &quot;Processing&quot; status are queryable, but visual features like direct visual context will only be available after processing completes.
+                <p>{column.description}</p>
+                <p className="mt-1 font-medium">Type: {column._type}</p>
+                {column.schema && (
+                  <p className="mt-1 text-xs">Schema provided</p>
+                )}
               </div>
             </div>
           </div>
-          <div className="col-span-3">ID</div>
+        ))}
+      </div>
+      
+      <div className="absolute right-3 top-1/2 transform -translate-y-1/2">
+        <Button 
+          variant="ghost" 
+          size="icon" 
+          className="h-6 w-6 rounded-full" 
+          title="Add column"
+          onClick={() => setShowAddColumnDialog(true)}
+        >
+          <Plus className="h-4 w-4" />
+          <span className="sr-only">Add column</span>
+        </Button>
+        
+        {/* Render the dialog separately */}
+        <AddColumnDialog 
+          isOpen={showAddColumnDialog}
+          onClose={() => setShowAddColumnDialog(false)}
+          onAddColumn={handleAddColumn}
+        />
+      </div>
+    </div>
+  );
+
+  if (loading && !documents.length) {
+    return (
+      <div className="border rounded-md overflow-hidden shadow-sm w-full">
+        <DocumentListHeader />
+        <div className="p-8">
+          <div className="flex flex-col items-center justify-center">
+            <div className="animate-spin rounded-full h-8 w-8 border-b-2 border-primary mb-4"></div>
+            <p className="text-muted-foreground">Loading documents...</p>
+          </div>
         </div>
       </div>
+    );
+  }
 
-      <ScrollArea className="h-[calc(100vh-200px)]">
+  return (
+    <div className="border rounded-md overflow-hidden shadow-sm w-full">
+      <DocumentListHeader />
+      <ScrollArea className="h-[calc(100vh-220px)]">
         {documents.map((doc) => (
           <div 
             key={doc.external_id}
             onClick={() => handleDocumentClick(doc)}
-            className={`grid grid-cols-12 p-3 cursor-pointer hover:bg-muted/50 border-b ${
+            className={`grid items-center w-full border-b ${
-              doc.external_id === selectedDocument?.external_id ? 'bg-muted' : ''
+              doc.external_id === selectedDocument?.external_id 
+                ? 'bg-primary/10 hover:bg-primary/15' 
+                : 'hover:bg-muted/70'
             }`}
+            style={{ 
+              gridTemplateColumns: `48px minmax(200px, 350px) 100px 120px ${allColumns.map(() => '140px').join(' ')}` 
+            }}
           >
-            <div className="col-span-1 flex items-center justify-center">
+            <div className="flex items-center justify-center p-3">
               <Checkbox 
                 id={`doc-${doc.external_id}`}
                 checked={selectedDocuments.includes(doc.external_id)}
@@ -92,46 +529,77 @@ const DocumentList: React.FC<DocumentListProps> = ({
                 aria-label={`Select ${doc.filename || 'document'}`}
               />
             </div>
-            <div className="col-span-4 flex items-center">
+            <div className="flex items-center p-3">
-              <span className="truncate">{doc.filename || 'N/A'}</span>
+              <span className="truncate font-medium">{doc.filename || 'N/A'}</span>
             </div>
-            <div className="col-span-2">
+            <div className="p-3">
-              <Badge variant="secondary">
+              <Badge variant="secondary" className="capitalize text-xs">
                 {doc.content_type.split('/')[0]}
               </Badge>
             </div>
-            <div className="col-span-2">
+            <div className="p-3">
               {doc.system_metadata?.status === "completed" ? (
-                <Badge variant="outline" className="bg-green-50 text-green-800 border-green-200">
+                <Badge variant="outline" className="bg-green-50 text-green-700 border-green-200 flex items-center gap-1 font-normal text-xs">
+                  <span className="h-1.5 w-1.5 rounded-full bg-green-500"></span>
                   Completed
                 </Badge>
               ) : doc.system_metadata?.status === "failed" ? (
-                <Badge variant="outline" className="bg-red-50 text-red-700 border-red-200">
+                <Badge variant="outline" className="bg-red-50 text-red-700 border-red-200 flex items-center gap-1 font-normal text-xs">
+                  <span className="h-1.5 w-1.5 rounded-full bg-red-500"></span>
                   Failed
                 </Badge>
               ) : (
                 <div className="group relative flex items-center">
-                  <Badge variant="outline" className="bg-amber-50 text-amber-700 border-amber-200 px-2 py-1">
+                  <Badge variant="outline" className="bg-amber-50 text-amber-700 border-amber-200 flex items-center gap-1 font-normal text-xs">
+                    <div className="h-1.5 w-1.5 rounded-full bg-amber-500 animate-pulse"></div>
                     Processing
                   </Badge>
-                  <div className="absolute left-0 -bottom-10 hidden group-hover:block bg-gray-800 text-white text-xs p-2 rounded whitespace-nowrap z-10">
+                  <div className="absolute left-0 -bottom-14 hidden group-hover:block bg-popover border text-foreground text-xs p-2 rounded-md whitespace-nowrap z-10 shadow-md">
                     Document is being processed. Partial search available.
                   </div>
                 </div>
               )}
             </div>
-            <div className="col-span-3 font-mono text-xs">
+            {/* Render metadata values for each column */}
-              {doc.external_id.substring(0, 8)}...
+            {allColumns.map((column) => (
-            </div>
+              <div key={column.name} className="p-3 truncate" title={String(doc.metadata?.[column.name] ?? '')}>
+                {String(doc.metadata?.[column.name] ?? '-')}
+              </div>
+            ))}
           </div>
         ))}
         
         {documents.length === 0 && (
-          <div className="p-8 text-center text-muted-foreground">
+          <div className="p-12 text-center flex flex-col items-center justify-center">
-            No documents found in this view. Try uploading a document or selecting a different folder.
+            <div className="h-12 w-12 rounded-full bg-muted flex items-center justify-center mb-4">
+              <svg xmlns="http://www.w3.org/2000/svg" width="24" height="24" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round" className="text-muted-foreground">
+                <path d="M14 2H6a2 2 0 0 0-2 2v16a2 2 0 0 0 2 2h12a2 2 0 0 0 2-2V8z"></path>
+                <polyline points="14 2 14 8 20 8"></polyline>
+                <line x1="9" y1="15" x2="15" y2="15"></line>
+              </svg>
+            </div>
+            <p className="text-muted-foreground">
+              No documents found in this view.
+            </p>
+            <p className="text-xs text-muted-foreground mt-1">
+              Try uploading a document or selecting a different folder.
+            </p>
           </div>
         )}
       </ScrollArea>
+      
+      {customColumns.length > 0 && (
+        <div className="border-t p-3 flex justify-end">
+          <Button 
+            className="gap-2" 
+            onClick={handleExtract}
+            disabled={isExtracting || !selectedFolder}
+          >
+            <Wand2 className="h-4 w-4" />
+            {isExtracting ? 'Processing...' : 'Extract'}
+          </Button>
+        </div>
+      )}
     </div>
   );
 };

ui-component/components/documents/DocumentsSection.tsx

@@ -8,18 +8,80 @@ import DocumentList from './DocumentList';
 import DocumentDetail from './DocumentDetail';
 import FolderList from './FolderList';
 import { UploadDialog, useUploadDialog } from './UploadDialog';
+import { cn } from '@/lib/utils';
 
 import { Document, Folder } from '@/components/types';
 
+// Custom hook for drag and drop functionality
+function useDragAndDrop({
+  onDrop,
+  disabled = false
+}: {
+  onDrop: (files: File[]) => void;
+  disabled?: boolean;
+}) {
+  const [isDragging, setIsDragging] = useState(false);
+  
+  const handleDragOver = useCallback((e: React.DragEvent<HTMLDivElement>) => {
+    if (disabled) return;
+    e.preventDefault();
+    e.stopPropagation();
+    setIsDragging(true);
+  }, [disabled]);
+  
+  const handleDragEnter = useCallback((e: React.DragEvent<HTMLDivElement>) => {
+    if (disabled) return;
+    e.preventDefault();
+    e.stopPropagation();
+    setIsDragging(true);
+  }, [disabled]);
+  
+  const handleDragLeave = useCallback((e: React.DragEvent<HTMLDivElement>) => {
+    if (disabled) return;
+    e.preventDefault();
+    e.stopPropagation();
+    setIsDragging(false);
+  }, [disabled]);
+  
+  const handleDrop = useCallback((e: React.DragEvent<HTMLDivElement>) => {
+    if (disabled) return;
+    e.preventDefault();
+    e.stopPropagation();
+    setIsDragging(false);
+    
+    const files = Array.from(e.dataTransfer.files);
+    if (files.length > 0) {
+      onDrop(files);
+    }
+  }, [disabled, onDrop]);
+  
+  return {
+    isDragging,
+    dragHandlers: {
+      onDragOver: handleDragOver,
+      onDragEnter: handleDragEnter,
+      onDragLeave: handleDragLeave,
+      onDrop: handleDrop
+    }
+  };
+}
+
 interface DocumentsSectionProps {
   apiBaseUrl: string;
   authToken: string | null;
+  initialFolder?: string | null;
+  setSidebarCollapsed?: (collapsed: boolean) => void;
 }
 
 // Debug render counter
 let renderCount = 0;
 
-const DocumentsSection: React.FC<DocumentsSectionProps> = ({ apiBaseUrl, authToken }) => {
+const DocumentsSection: React.FC<DocumentsSectionProps> = ({ 
+  apiBaseUrl, 
+  authToken, 
+  initialFolder = null,
+  setSidebarCollapsed
+}) => {
   // Increment render counter for debugging
   renderCount++;
   console.log(`DocumentsSection rendered: #${renderCount}`);
@@ -38,7 +100,7 @@ const DocumentsSection: React.FC<DocumentsSectionProps> = ({ apiBaseUrl, authTok
   // State for documents and folders
   const [documents, setDocuments] = useState<Document[]>([]);
   const [folders, setFolders] = useState<Folder[]>([]);
-  const [selectedFolder, setSelectedFolder] = useState<string | null>(null);
+  const [selectedFolder, setSelectedFolder] = useState<string | null>(initialFolder);
   const [selectedDocument, setSelectedDocument] = useState<Document | null>(null);
   const [selectedDocuments, setSelectedDocuments] = useState<string[]>([]);
   const [loading, setLoading] = useState(false);
@@ -57,6 +119,17 @@ const DocumentsSection: React.FC<DocumentsSectionProps> = ({ apiBaseUrl, authTok
     useColpali,
     resetUploadDialog
   } = uploadDialogState;
+  
+  // Initialize drag and drop
+  const { isDragging, dragHandlers } = useDragAndDrop({
+    onDrop: (files) => {
+      // Only allow drag and drop when inside a folder
+      if (selectedFolder && selectedFolder !== null) {
+        handleBatchFileUpload(files, true);
+      }
+    },
+    disabled: !selectedFolder || selectedFolder === null
+  });
 
   // No need for a separate header function, use authToken directly
 
@@ -357,6 +430,15 @@ const DocumentsSection: React.FC<DocumentsSectionProps> = ({ apiBaseUrl, authTok
     }
   };
 
+  // Collapse sidebar when a folder is selected
+  useEffect(() => {
+    if (selectedFolder !== null && setSidebarCollapsed) {
+      setSidebarCollapsed(true);
+    } else if (setSidebarCollapsed) {
+      setSidebarCollapsed(false);
+    }
+  }, [selectedFolder, setSidebarCollapsed]);
+
   // Fetch documents when selected folder changes
   useEffect(() => {
     // Skip initial render to prevent double fetching with the auth useEffect
@@ -790,7 +872,7 @@ const DocumentsSection: React.FC<DocumentsSectionProps> = ({ apiBaseUrl, authTok
   };
 
   // Handle batch file upload
-  const handleBatchFileUpload = async (files: File[]) => {
+  const handleBatchFileUpload = async (files: File[], fromDragAndDrop: boolean = false) => {
     if (files.length === 0) {
       showAlert('Please select files to upload', {
         type: 'error',
@@ -799,8 +881,11 @@ const DocumentsSection: React.FC<DocumentsSectionProps> = ({ apiBaseUrl, authTok
       return;
     }
 
-    // Close dialog and update upload count using alert system
+    // Close dialog if it's open (but not if drag and drop)
-    setShowUploadDialog(false);
+    if (!fromDragAndDrop) {
+      setShowUploadDialog(false);
+    }
+    
     const fileCount = files.length;
     const uploadId = 'batch-upload-progress';
     showAlert(`Uploading ${fileCount} files...`, {
@@ -809,14 +894,16 @@ const DocumentsSection: React.FC<DocumentsSectionProps> = ({ apiBaseUrl, authTok
       id: uploadId
     });
     
-    // Save form data locally before resetting
+    // Save form data locally
     const batchFilesRef = [...files];
     const metadataRef = metadata;
     const rulesRef = rules;
     const useColpaliRef = useColpali;
     
-    // Reset form immediately
+    // Only reset form if not from drag and drop
-    resetUploadDialog();
+    if (!fromDragAndDrop) {
+      resetUploadDialog();
+    }
     
     try {      
       const formData = new FormData();
@@ -1100,115 +1187,121 @@ const DocumentsSection: React.FC<DocumentsSectionProps> = ({ apiBaseUrl, authTok
     }
   };
 
-  // Title based on selected folder
+  // Function to trigger refresh
-  const sectionTitle = selectedFolder === null 
+  const handleRefresh = () => {
-    ? "Folders" 
+    console.log("Manual refresh triggered");
-    : selectedFolder === "all" 
+    // Show a loading indicator
-      ? "All Documents" 
+    showAlert("Refreshing documents and folders...", {
-      : `Folder: ${selectedFolder}`;
+      type: 'info',
+      duration: 1500
+    });
+    
+    // First clear folder data to force a clean refresh
+    setLoading(true);
+    setFolders([]);
+    
+    // Create a new function to perform a truly fresh fetch
+    const performFreshFetch = async () => {
+      try {
+        // First get fresh folder data from the server
+        const folderResponse = await fetch(`${effectiveApiUrl}/folders`, {
+          method: 'GET',
+          headers: {
+            ...(authToken ? { 'Authorization': `Bearer ${authToken}` } : {})
+          }
+        });
+        
+        if (!folderResponse.ok) {
+          throw new Error(`Failed to fetch folders: ${folderResponse.statusText}`);
+        }
+        
+        // Get the fresh folder data
+        const freshFolders = await folderResponse.json();
+        console.log(`Refresh: Fetched ${freshFolders.length} folders with fresh data`);
+        
+        // Update folders state with fresh data
+        setFolders(freshFolders);
+        
+        // Use our helper function to refresh documents with fresh folder data
+        await refreshDocuments(freshFolders);
+        
+        // Show success message
+        showAlert("Refresh completed successfully", {
+          type: 'success',
+          duration: 1500
+        });
+      } catch (error) {
+        console.error("Error during refresh:", error);
+        showAlert(`Error refreshing: ${error instanceof Error ? error.message : 'Unknown error'}`, {
+          type: 'error',
+          duration: 3000
+        });
+      } finally {
+        setLoading(false);
+      }
+    };
+    
+    // Execute the fresh fetch
+    performFreshFetch();
+  };
 
   return (
-    <div className="flex-1 flex flex-col h-full">
+    <div 
-      <div className="flex justify-between items-center py-3 mb-4">
+      className={cn(
-        <div className="flex items-center gap-4">
+        "flex-1 flex flex-col h-full relative",
+        selectedFolder && isDragging ? "drag-active" : ""
+      )}
+      {...(selectedFolder ? dragHandlers : {})}
+    >
+      {/* Drag overlay - only visible when dragging files over the folder */}
+      {isDragging && selectedFolder && (
+        <div className="absolute inset-0 bg-primary/10 backdrop-blur-sm z-50 flex items-center justify-center rounded-lg border-2 border-dashed border-primary animate-pulse">
+          <div className="bg-background p-8 rounded-lg shadow-lg text-center">
+            <Upload className="h-12 w-12 mx-auto mb-4 text-primary" />
+            <h3 className="text-xl font-medium mb-2">Drop to Upload</h3>
+            <p className="text-muted-foreground">
+              Files will be added to {selectedFolder === "all" ? "your documents" : `folder "${selectedFolder}"`}
+            </p>
+          </div>
+        </div>
+      )}
+      {/* Hide the main header when viewing a specific folder - it will be merged with the FolderList header */}
+      {selectedFolder === null && (
+        <div className="flex justify-between items-center py-3 mb-4">
           <div>
-            <h2 className="text-2xl font-bold leading-tight">{sectionTitle}</h2>
+            <h2 className="text-2xl font-bold leading-tight">Folders</h2>
             <p className="text-muted-foreground">Manage your uploaded documents and view their metadata.</p>
           </div>
-          {selectedDocuments.length > 0 && (
+          <div className="flex items-center gap-2">
-            <Button 
+            <Button
-              variant="outline" 
+              variant="outline"
-              onClick={handleDeleteMultipleDocuments} 
+              onClick={handleRefresh}
               disabled={loading}
-              className="border-red-500 text-red-500 hover:bg-red-50 ml-4"
+              className="flex items-center"
+              title="Refresh folders"
             >
-              Delete {selectedDocuments.length} selected
+              <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round" className="mr-1">
+                <path d="M3 12a9 9 0 0 1 9-9 9.75 9.75 0 0 1 6.74 2.74L21 8"></path>
+                <path d="M21 3v5h-5"></path>
+                <path d="M21 12a9 9 0 0 1-9 9 9.75 9.75 0 0 1-6.74-2.74L3 16"></path>
+                <path d="M8 16H3v5"></path>
+              </svg>
+              Refresh
             </Button>
-          )}
+            <UploadDialog
+              showUploadDialog={showUploadDialog}
+              setShowUploadDialog={setShowUploadDialog}
+              loading={loading}
+              onFileUpload={handleFileUpload}
+              onBatchFileUpload={handleBatchFileUpload}
+              onTextUpload={handleTextUpload}
+            />
+          </div>
         </div>
-        <div className="flex items-center gap-2">
+      )}
-          <Button
-            variant="outline"
-            onClick={() => {
-              console.log("Manual refresh triggered");
-              // Show a loading indicator
-              showAlert("Refreshing documents and folders...", {
-                type: 'info',
-                duration: 1500
-              });
-              
-              // First clear folder data to force a clean refresh
-              setLoading(true);
-              setFolders([]);
-              
-              // Create a new function to perform a truly fresh fetch
-              const performFreshFetch = async () => {
-                try {
-                  // First get fresh folder data from the server
-                  const folderResponse = await fetch(`${effectiveApiUrl}/folders`, {
-                    method: 'GET',
-                    headers: {
-                      ...(authToken ? { 'Authorization': `Bearer ${authToken}` } : {})
-                    }
-                  });
-                  
-                  if (!folderResponse.ok) {
-                    throw new Error(`Failed to fetch folders: ${folderResponse.statusText}`);
-                  }
-                  
-                  // Get the fresh folder data
-                  const freshFolders = await folderResponse.json();
-                  console.log(`Refresh: Fetched ${freshFolders.length} folders with fresh data`);
-                  
-                  // Update folders state with fresh data
-                  setFolders(freshFolders);
-                  
-                  // Use our helper function to refresh documents with fresh folder data
-                  await refreshDocuments(freshFolders);
-                  
-                  // Show success message
-                  showAlert("Refresh completed successfully", {
-                    type: 'success',
-                    duration: 1500
-                  });
-                } catch (error) {
-                  console.error("Error during refresh:", error);
-                  showAlert(`Error refreshing: ${error instanceof Error ? error.message : 'Unknown error'}`, {
-                    type: 'error',
-                    duration: 3000
-                  });
-                } finally {
-                  setLoading(false);
-                }
-              };
-              
-              // Execute the fresh fetch
-              performFreshFetch();
-            }}
-            disabled={loading}
-            className="flex items-center"
-            title="Refresh documents"
-          >
-            <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round" className="mr-1">
-              <path d="M3 12a9 9 0 0 1 9-9 9.75 9.75 0 0 1 6.74 2.74L21 8"></path>
-              <path d="M21 3v5h-5"></path>
-              <path d="M21 12a9 9 0 0 1-9 9 9.75 9.75 0 0 1-6.74-2.74L3 16"></path>
-              <path d="M8 16H3v5"></path>
-            </svg>
-            Refresh
-          </Button>
-          <UploadDialog
-            showUploadDialog={showUploadDialog}
-            setShowUploadDialog={setShowUploadDialog}
-            loading={loading}
-            onFileUpload={handleFileUpload}
-            onBatchFileUpload={handleBatchFileUpload}
-            onTextUpload={handleTextUpload}
-          />
-        </div>
-      </div>
 
-      <div className="flex flex-col gap-4 flex-1">
+      {/* Render the FolderList with header at all times when selectedFolder is not null */}
+      {selectedFolder !== null && (
         <FolderList
           folders={folders}
           selectedFolder={selectedFolder}
@@ -1217,47 +1310,108 @@ const DocumentsSection: React.FC<DocumentsSectionProps> = ({ apiBaseUrl, authTok
           authToken={authToken}
           refreshFolders={fetchFolders}
           loading={foldersLoading}
+          refreshAction={handleRefresh}
+          selectedDocuments={selectedDocuments}
+          handleDeleteMultipleDocuments={handleDeleteMultipleDocuments}
+          uploadDialogComponent={
+            <UploadDialog
+              showUploadDialog={showUploadDialog}
+              setShowUploadDialog={setShowUploadDialog}
+              loading={loading}
+              onFileUpload={handleFileUpload}
+              onBatchFileUpload={handleBatchFileUpload}
+              onTextUpload={handleTextUpload}
+            />
+          }
         />
-        
+      )}
-        {selectedFolder !== null ? (
+      
-          documents.length === 0 && !loading ? (
+      {documents.length === 0 && !loading && folders.length === 0 && !foldersLoading ? (
-            <div className="text-center py-8 border border-dashed rounded-lg flex-1 flex items-center justify-center">
+        <div className="text-center py-8 border border-dashed rounded-lg flex-1 flex items-center justify-center">
-              <div>
+          <div>
-                <Upload className="mx-auto h-12 w-12 mb-2 text-muted-foreground" />
+            <Upload className="mx-auto h-12 w-12 mb-2 text-muted-foreground" />
-                <p className="text-muted-foreground">No documents found in this folder. Upload a document.</p>
+            <p className="text-muted-foreground">No documents found. Upload your first document.</p>
-              </div>
+          </div>
+        </div>
+      ) : selectedFolder && documents.length === 0 && !loading ? (
+        <div className="text-center py-8 border border-dashed rounded-lg flex-1 flex items-center justify-center">
+          <div>
+            <Upload className="mx-auto h-12 w-12 mb-2 text-muted-foreground" />
+            <p className="text-muted-foreground">Drag and drop files here to upload to this folder.</p>
+            <p className="text-xs text-muted-foreground mt-2">Or use the upload button in the top right.</p>
+          </div>
+        </div>
+      ) : selectedFolder && loading ? (
+        <div className="text-center py-8 flex-1 flex items-center justify-center">
+          <div className="flex flex-col items-center">
+            <div className="animate-spin rounded-full h-8 w-8 border-b-2 border-primary mb-4"></div>
+            <p className="text-muted-foreground">Loading documents...</p>
+          </div>
+        </div>
+      ) : selectedFolder === null ? (
+        <div className="flex flex-col gap-4 flex-1">
+          <FolderList
+            folders={folders}
+            selectedFolder={selectedFolder}
+            setSelectedFolder={setSelectedFolder}
+            apiBaseUrl={effectiveApiUrl}
+            authToken={authToken}
+            refreshFolders={fetchFolders}
+            loading={foldersLoading}
+            refreshAction={handleRefresh}
+            selectedDocuments={selectedDocuments}
+            handleDeleteMultipleDocuments={handleDeleteMultipleDocuments}
+            uploadDialogComponent={
+              <UploadDialog
+                showUploadDialog={showUploadDialog}
+                setShowUploadDialog={setShowUploadDialog}
+                loading={loading}
+                onFileUpload={handleFileUpload}
+                onBatchFileUpload={handleBatchFileUpload}
+                onTextUpload={handleTextUpload}
+              />
+            }
+          />
+        </div>
+      ) : (
+        <div className="flex flex-col md:flex-row gap-4 flex-1">
+          <div className={cn(
+            "w-full transition-all duration-300",
+            selectedDocument ? "md:w-2/3" : "md:w-full"
+          )}>
+            <DocumentList
+              documents={documents}
+              selectedDocument={selectedDocument}
+              selectedDocuments={selectedDocuments}
+              handleDocumentClick={handleDocumentClick}
+              handleCheckboxChange={handleCheckboxChange}
+              getSelectAllState={getSelectAllState}
+              setSelectedDocuments={setSelectedDocuments}
+              setDocuments={setDocuments}
+              loading={loading}
+              apiBaseUrl={effectiveApiUrl}
+              authToken={authToken}
+              selectedFolder={selectedFolder}
+            />
+          </div>
+          
+          {selectedDocument && (
+            <div className="w-full md:w-1/3 animate-in slide-in-from-right duration-300">
+              <DocumentDetail
+                selectedDocument={selectedDocument}
+                handleDeleteDocument={handleDeleteDocument}
+                folders={folders}
+                apiBaseUrl={effectiveApiUrl}
+                authToken={authToken}
+                refreshDocuments={fetchDocuments}
+                refreshFolders={fetchFolders}
+                loading={loading}
+                onClose={() => setSelectedDocument(null)}
+              />
             </div>
-          ) : (
+          )}
-            <div className="flex flex-col md:flex-row gap-4 flex-1">
+        </div>
-              <div className="w-full md:w-2/3">
+      )}
-                <DocumentList
-                  documents={documents}
-                  selectedDocument={selectedDocument}
-                  selectedDocuments={selectedDocuments}
-                  handleDocumentClick={handleDocumentClick}
-                  handleCheckboxChange={handleCheckboxChange}
-                  getSelectAllState={getSelectAllState}
-                  setSelectedDocuments={setSelectedDocuments}
-                  loading={loading}
-                />
-              </div>
-              
-              <div className="w-full md:w-1/3">
-                <DocumentDetail
-                  selectedDocument={selectedDocument}
-                  handleDeleteDocument={handleDeleteDocument}
-                  folders={folders}
-                  apiBaseUrl={effectiveApiUrl}
-                  authToken={authToken}
-                  refreshDocuments={fetchDocuments}
-                  refreshFolders={fetchFolders}
-                  loading={loading}
-                />
-              </div>
-            </div>
-          )
-        ) : null}
-      </div>
     </div>
   );
 };

ui-component/components/documents/FolderList.tsx

@@ -2,14 +2,14 @@
 
 import React from 'react';
 import { Button } from '@/components/ui/button';
-import { PlusCircle, Folder as FolderIcon, File, ArrowLeft } from 'lucide-react';
+import { PlusCircle, ArrowLeft } from 'lucide-react';
 import { Dialog, DialogContent, DialogDescription, DialogFooter, DialogHeader, DialogTitle, DialogTrigger } from '@/components/ui/dialog';
 import { Label } from '@/components/ui/label';
 import { Input } from '@/components/ui/input';
 import { Textarea } from '@/components/ui/textarea';
-import { cn } from '@/lib/utils';
-import { Card, CardContent } from '@/components/ui/card';
 import { Folder } from '@/components/types';
+import { useRouter, usePathname } from 'next/navigation';
+import Image from 'next/image';
 
 interface FolderListProps {
   folders: Folder[];
@@ -19,6 +19,12 @@ interface FolderListProps {
   authToken: string | null;
   refreshFolders: () => void;
   loading: boolean;
+  refreshAction?: () => void;
+  selectedDocuments?: string[];
+  handleDeleteMultipleDocuments?: () => void;
+  showUploadDialog?: boolean;
+  setShowUploadDialog?: (show: boolean) => void;
+  uploadDialogComponent?: React.ReactNode;
 }
 
 const FolderList: React.FC<FolderListProps> = ({
@@ -28,12 +34,30 @@ const FolderList: React.FC<FolderListProps> = ({
   apiBaseUrl,
   authToken,
   refreshFolders,
-  loading
+  loading,
+  refreshAction,
+  selectedDocuments = [],
+  handleDeleteMultipleDocuments,
+  uploadDialogComponent
 }) => {
+  const router = useRouter();
+  const pathname = usePathname();
   const [showNewFolderDialog, setShowNewFolderDialog] = React.useState(false);
   const [newFolderName, setNewFolderName] = React.useState('');
   const [newFolderDescription, setNewFolderDescription] = React.useState('');
   const [isCreatingFolder, setIsCreatingFolder] = React.useState(false);
+  
+  // Function to update both state and URL
+  const updateSelectedFolder = (folderName: string | null) => {
+    setSelectedFolder(folderName);
+    
+    // Update URL to reflect the selected folder
+    if (folderName) {
+      router.push(`${pathname}?folder=${encodeURIComponent(folderName)}`);
+    } else {
+      router.push(pathname);
+    }
+  };
 
   const handleCreateFolder = async () => {
     if (!newFolderName.trim()) return;
@@ -73,7 +97,7 @@ const FolderList: React.FC<FolderListProps> = ({
       
       // Auto-select this newly created folder so user can immediately add files to it
       // This ensures we start with a clean empty folder view
-      setSelectedFolder(folderData.name);
+      updateSelectedFolder(folderData.name);
       
     } catch (error) {
       console.error('Error creating folder:', error);
@@ -86,28 +110,61 @@ const FolderList: React.FC<FolderListProps> = ({
   if (selectedFolder !== null) {
     return (
       <div className="mb-4">
-        <div className="flex items-center gap-2 mb-4">
+        <div className="flex justify-between items-center py-2">
-          <Button 
+          <div className="flex items-center gap-4">
-            variant="ghost" 
+            <Button 
-            size="sm" 
+              variant="ghost" 
-            className="p-1 h-8 w-8" 
+              size="icon"
-            onClick={() => setSelectedFolder(null)}
+              className="rounded-full hover:bg-muted/50" 
-          >
+              onClick={() => updateSelectedFolder(null)}
-            <ArrowLeft size={16} />
+            >
-          </Button>
+              <ArrowLeft size={18} />
-          <h2 className="font-medium text-lg flex items-center">
+            </Button>
-            {selectedFolder === "all" ? (
+            <div className="flex items-center">
-              <>
+              {selectedFolder === "all" ? (
-                <File className="h-5 w-5 mr-2" />
+                <span className="text-3xl mr-3" aria-hidden="true">📄</span>
-                All Documents
+              ) : (
-              </>
+                <Image src="/icons/folder-icon.png" alt="Folder" width={32} height={32} className="mr-3" />
-            ) : (
+              )}
-              <>
+              <h2 className="font-medium text-xl">
-                <FolderIcon className="h-5 w-5 mr-2" />
+                {selectedFolder === "all" ? "All Documents" : selectedFolder}
-                {selectedFolder}
+              </h2>
-              </>
+            </div>
+            
+            {/* Show delete button if documents are selected */}
+            {selectedDocuments.length > 0 && handleDeleteMultipleDocuments && (
+              <Button 
+                variant="outline" 
+                onClick={handleDeleteMultipleDocuments}
+                className="border-red-500 text-red-500 hover:bg-red-50 ml-4"
+              >
+                Delete {selectedDocuments.length} selected
+              </Button>
             )}
-          </h2>
+          </div>
+          
+          {/* Action buttons */}
+          <div className="flex items-center gap-2">
+            {refreshAction && (
+              <Button
+                variant="outline"
+                onClick={refreshAction}
+                className="flex items-center"
+                title="Refresh documents"
+              >
+                <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round" className="mr-1">
+                  <path d="M3 12a9 9 0 0 1 9-9 9.75 9.75 0 0 1 6.74 2.74L21 8"></path>
+                  <path d="M21 3v5h-5"></path>
+                  <path d="M21 12a9 9 0 0 1-9 9 9.75 9.75 0 0 1-6.74-2.74L3 16"></path>
+                  <path d="M8 16H3v5"></path>
+                </svg>
+                Refresh
+              </Button>
+            )}
+            
+            {/* Upload dialog component */}
+            {uploadDialogComponent}
+          </div>
         </div>
       </div>
     );
@@ -165,46 +222,44 @@ const FolderList: React.FC<FolderListProps> = ({
         </Dialog>
       </div>
       
-      <div className="grid grid-cols-2 sm:grid-cols-3 md:grid-cols-4 lg:grid-cols-5 gap-4">
+      <div className="grid grid-cols-2 sm:grid-cols-3 md:grid-cols-4 lg:grid-cols-6 gap-6 py-2">
-        <Card
+        <div 
-          className={cn(
+          className="cursor-pointer group flex flex-col items-center"
-            "cursor-pointer hover:border-primary transition-colors",
+          onClick={() => updateSelectedFolder("all")}
-            "flex flex-col items-center justify-center h-24"
-          )}
-          onClick={() => setSelectedFolder("all")}
         >
-          <CardContent className="p-4 flex flex-col items-center justify-center">
+          <div className="mb-2 group-hover:scale-110 transition-transform">
-            <File className="h-10 w-10 mb-1" />
+            <span className="text-4xl" aria-hidden="true">📄</span>
-            <span className="text-sm font-medium text-center">All Documents</span>
+          </div>
-          </CardContent>
+          <span className="text-sm font-medium text-center group-hover:text-primary transition-colors">All Documents</span>
-        </Card>
+        </div>
         
         {folders.map((folder) => (
-          <Card
+          <div
             key={folder.name}
-            className={cn(
+            className="cursor-pointer group flex flex-col items-center"
-              "cursor-pointer hover:border-primary transition-colors",
+            onClick={() => updateSelectedFolder(folder.name)}
-              "flex flex-col items-center justify-center h-24"
-            )}
-            onClick={() => setSelectedFolder(folder.name)}
           >
-            <CardContent className="p-4 flex flex-col items-center justify-center">
+            <div className="mb-2 group-hover:scale-110 transition-transform">
-              <FolderIcon className="h-10 w-10 mb-1" />
+              <Image src="/icons/folder-icon.png" alt="Folder" width={64} height={64} />
-              <span className="text-sm font-medium truncate text-center w-full">{folder.name}</span>
+            </div>
-            </CardContent>
+            <span className="text-sm font-medium truncate text-center w-full max-w-[100px] group-hover:text-primary transition-colors">{folder.name}</span>
-          </Card>
+          </div>
         ))}
       </div>
       
       {folders.length === 0 && !loading && (
-        <div className="text-center p-8 text-sm text-muted-foreground">
+        <div className="flex flex-col items-center justify-center p-8 mt-4">
-          No folders yet. Create one to organize your documents.
+          <Image src="/icons/folder-icon.png" alt="Folder" width={80} height={80} className="opacity-50 mb-3" />
+          <p className="text-sm text-muted-foreground">No folders yet. Create one to organize your documents.</p>
         </div>
       )}
       
       {loading && folders.length === 0 && (
-        <div className="text-center p-8 text-sm text-muted-foreground">
+        <div className="flex items-center justify-center p-8 mt-4">
-          Loading folders...
+          <div className="flex items-center space-x-2">
+            <div className="animate-spin rounded-full h-5 w-5 border-b-2 border-primary"></div>
+            <p className="text-sm text-muted-foreground">Loading folders...</p>
+          </div>
         </div>
       )}
     </div>

ui-component/components/types.ts

@@ -7,6 +7,7 @@ export interface MorphikUIProps {
   onUriChange?: (uri: string) => void; // Callback when URI is changed
   onBackClick?: () => void; // Callback when back button is clicked
   appName?: string; // Name of the app to display in UI
+  initialFolder?: string | null; // Initial folder to show
 }
 
 export interface Document {

ui-component/components/ui/checkbox.tsx

@@ -13,7 +13,7 @@ const Checkbox = React.forwardRef<
   <CheckboxPrimitive.Root
     ref={ref}
     className={cn(
-      "peer h-4 w-4 shrink-0 rounded-sm border border-primary shadow focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring disabled:cursor-not-allowed disabled:opacity-50 data-[state=checked]:bg-primary data-[state=checked]:text-primary-foreground",
+      "peer h-5 w-5 shrink-0 rounded-sm border-2 border-primary/50 shadow focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring hover:border-primary transition-colors disabled:cursor-not-allowed disabled:opacity-50 data-[state=checked]:bg-primary data-[state=checked]:text-primary-foreground data-[state=checked]:border-primary",
       className
     )}
     {...props}
@@ -21,7 +21,7 @@ const Checkbox = React.forwardRef<
     <CheckboxPrimitive.Indicator
       className={cn("flex items-center justify-center text-current")}
     >
-      <Check className="h-4 w-4" />
+      <Check className="h-3.5 w-3.5 stroke-[3]" />
     </CheckboxPrimitive.Indicator>
   </CheckboxPrimitive.Root>
 ))

ui-component/components/ui/sidebar.tsx

@@ -14,6 +14,8 @@ interface SidebarProps extends React.HTMLAttributes<HTMLDivElement> {
   connectionUri?: string
   isReadOnlyUri?: boolean
   onUriChange?: (uri: string) => void
+  isCollapsed?: boolean
+  setIsCollapsed?: (collapsed: boolean) => void
 }
 
 export function Sidebar({ 
@@ -23,12 +25,26 @@ export function Sidebar({
   connectionUri, 
   isReadOnlyUri = false,
   onUriChange,
+  isCollapsed: externalIsCollapsed,
+  setIsCollapsed: externalSetIsCollapsed,
   ...props 
 }: SidebarProps) {
-  const [isCollapsed, setIsCollapsed] = React.useState(false)
+  // Use internal state that syncs with external state if provided
+  const [internalIsCollapsed, setInternalIsCollapsed] = React.useState(false)
   const [editableUri, setEditableUri] = React.useState('')
   const [isEditingUri, setIsEditingUri] = React.useState(false)
   
+  // Determine if sidebar is collapsed based on props or internal state
+  const isCollapsed = externalIsCollapsed !== undefined ? externalIsCollapsed : internalIsCollapsed
+  
+  // Toggle function that updates both internal and external state if provided
+  const toggleCollapsed = () => {
+    if (externalSetIsCollapsed) {
+      externalSetIsCollapsed(!isCollapsed)
+    }
+    setInternalIsCollapsed(!isCollapsed)
+  }
+  
   // Initialize from localStorage or props
   React.useEffect(() => {
     // For development/testing - check if we have a stored URI
@@ -110,7 +126,7 @@ export function Sidebar({
             variant="ghost"
             size="icon"
             className="ml-auto"
-            onClick={() => setIsCollapsed(!isCollapsed)}
+            onClick={toggleCollapsed}
           >
             {isCollapsed ? <ChevronRight className="h-4 w-4" /> : <ChevronLeft className="h-4 w-4" />}
           </Button>