Merge branch 'support/anybase' of https://github.com/apecloud/ApeRAG into support/anybase

Author: sailwebs

aperag/db/repositories/collection.py

@@ -201,14 +201,13 @@ async def _query(session):
             return result.scalars().first()
 
         return await self._execute_query(_query)
-
-    async def query_collections_by_ids(self, user: str, collection_ids: List[str]):
+    
+    async def query_collections_by_ids(self, collection_ids: List[str]):
         """Query multiple collections by their IDs in a single database call"""
 
         async def _query(session):
             stmt = select(Collection).where(
                 Collection.id.in_(collection_ids),
-                Collection.user == user,
                 Collection.status != CollectionStatus.DELETED,
             )
             result = await session.execute(stmt)

aperag/db/repositories/marketplace.py

@@ -237,12 +237,23 @@ async def _operation(session):
 
     # Marketplace listing operations
     async def list_published_collections_with_subscription_status(
-        self, user_id: str, page: int = 1, page_size: int = 12, user_group_id: Optional[str] = None
+        self, user_id: str, page: int = 1, page_size: int = 12
     ) -> Tuple[List[dict], int]:
         """List all published collections accessible to user with subscription status"""
 
         async def _query(session):
+            # Get user's department information
+            user_result = await session.execute(
+                select(User).where(User.id == user_id)
+            )
+            user_obj = user_result.scalars().first()
+
+            if not user_obj:
+                return []
+
+            user_group_id = user_obj.department_id
             # Get user's accessible department IDs
+
             accessible_group_ids = ["*"]  # Global collections are always accessible
 
             if user_group_id:
@@ -338,6 +349,137 @@ async def _query(session):
 
         return await self._execute_query(_query)
 
+    async def list_all_accessible_collections_for_user(
+        self, user_id: str
+    ) -> List[dict]:
+        """List all collections accessible to user (owned + published accessible collections)"""
+
+        async def _query(session):
+            # Get user's department information
+            user_result = await session.execute(
+                select(User).where(User.id == user_id)
+            )
+            user_obj = user_result.scalars().first()
+
+            if not user_obj:
+                return []
+
+            user_group_id = user_obj.department_id
+
+            # Get user's accessible department IDs
+            accessible_group_ids = ["*"]  # Global collections are always accessible
+
+            if user_group_id:
+                # Get user's department hierarchy
+                dept_result = await session.execute(
+                    select(Department).where(Department.id == user_group_id)
+                )
+                user_dept = dept_result.scalars().first()
+
+                if user_dept and user_dept.group_path:
+                    # Extract all parent department IDs from group_path
+                    path_parts = user_dept.group_path.strip("/").split("/")
+                    accessible_group_ids.extend([part for part in path_parts if part])
+
+                # Also include the user's direct department
+                accessible_group_ids.append(user_group_id)
+
+            # Query for user's own collections
+            owned_collections_stmt = (
+                select(
+                    Collection.id.label("id"),
+                    Collection.title,
+                    Collection.description,
+                    Collection.config,
+                    Collection.type,
+                    Collection.status,
+                    Collection.gmt_created,
+                    Collection.gmt_updated,
+                    Collection.user.label("owner_user_id"),
+                    User.username.label("owner_username"),
+                    CollectionMarketplace.id.label("marketplace_id"),
+                    CollectionMarketplace.status.label("marketplace_status"),
+                    CollectionMarketplace.gmt_created.label("published_at"),
+                    CollectionMarketplace.group_id,
+                )
+                .select_from(Collection)
+                .join(User, Collection.user == User.id)
+                .outerjoin(
+                    CollectionMarketplace,
+                    and_(
+                        CollectionMarketplace.collection_id == Collection.id,
+                        CollectionMarketplace.gmt_deleted.is_(None)
+                    )
+                )
+                .where(
+                    Collection.user == user_id,
+                    Collection.status != CollectionStatus.DELETED,
+                    Collection.gmt_deleted.is_(None),
+                )
+            )
+
+            # Query for published collections accessible to user (not owned by user)
+            published_collections_stmt = (
+                select(
+                    Collection.id.label("id"),
+                    Collection.title,
+                    Collection.description,
+                    Collection.config,
+                    Collection.type,
+                    Collection.status,
+                    Collection.gmt_created,
+                    Collection.gmt_updated,
+                    Collection.user.label("owner_user_id"),
+                    User.username.label("owner_username"),
+                    CollectionMarketplace.id.label("marketplace_id"),
+                    CollectionMarketplace.status.label("marketplace_status"),
+                    CollectionMarketplace.gmt_created.label("published_at"),
+                    CollectionMarketplace.group_id,
+                )
+                .select_from(CollectionMarketplace)
+                .join(Collection, CollectionMarketplace.collection_id == Collection.id)
+                .join(User, Collection.user == User.id)
+                .where(
+                    CollectionMarketplace.status == CollectionMarketplaceStatusEnum.PUBLISHED.value,
+                    CollectionMarketplace.gmt_deleted.is_(None),
+                    Collection.status != CollectionStatus.DELETED,
+                    Collection.gmt_deleted.is_(None),
+                    CollectionMarketplace.group_id.in_(accessible_group_ids),
+                    Collection.user != user_id,  # Exclude user's own collections to avoid duplicates
+                )
+            )
+
+            # Combine both queries using UNION ALL
+            combined_stmt = owned_collections_stmt.union_all(published_collections_stmt)
+
+            # Execute the combined query
+            result = await session.execute(combined_stmt)
+            rows = result.all()
+            
+            # Convert rows to dictionaries for easier access
+            collections = []
+            for row in rows:
+                collections.append({
+                    'id': row.id,
+                    'title': row.title,
+                    'description': row.description,
+                    'config': row.config,
+                    'type': row.type,
+                    'status': row.status,
+                    'gmt_created': row.gmt_created,
+                    'gmt_updated': row.gmt_updated,
+                    'owner_user_id': row.owner_user_id,
+                    'owner_username': row.owner_username,
+                    'marketplace_id': row.marketplace_id,
+                    'marketplace_status': row.marketplace_status,
+                    'published_at': row.published_at,
+                    'group_id': row.group_id,
+                })
+            
+            return collections
+
+        return await self._execute_query(_query)
+
     # Subscription operations
     async def create_subscription(self, user_id: str, collection_marketplace_id: str) -> UserCollectionSubscription:
         """Create a new subscription"""

aperag/service/agent_chat_service.py

@@ -53,7 +53,9 @@
 from aperag.agent.response_types import AgentErrorResponse, AgentToolCallResultResponse
 from aperag.chat.history.message import StoredChatMessage, create_assistant_message
 from aperag.db.ops import AsyncDatabaseOps, async_db_ops
+from aperag.exceptions import ResourceNotFoundException
 from aperag.schema import view_models
+from aperag.schema.utils import parseCollectionConfig
 from aperag.service.prompt_template_service import build_agent_query_prompt, get_agent_system_prompt
 from aperag.trace import trace_async_function
 
@@ -100,25 +102,6 @@ def __init__(self, session: AsyncSession = None):
         self.memory_manager = AgentMemoryManager()
         self.history_manager = AgentHistoryManager()
 
-    async def _convert_db_collections_to_pydantic(self, db_collections) -> List[view_models.Collection]:
-        """Convert SQLAlchemy Collection models to Pydantic Collection models"""
-        from aperag.schema.utils import parseCollectionConfig
-
-        pydantic_collections = []
-        for db_collection in db_collections:
-            pydantic_collection = view_models.Collection(
-                id=db_collection.id,
-                title=db_collection.title,
-                description=db_collection.description,
-                type=db_collection.type,
-                status=getattr(db_collection, "status", None),
-                config=parseCollectionConfig(db_collection.config),
-                created=db_collection.gmt_created.isoformat(),
-                updated=db_collection.gmt_updated.isoformat(),
-            )
-            pydantic_collections.append(pydantic_collection)
-        return pydantic_collections
-
     def _parse_websocket_message(
         self, raw_data: str
     ) -> Tuple[Optional[view_models.AgentMessage], Optional[AgentErrorResponse]]:
@@ -191,10 +174,19 @@ async def handle_websocket_agent_chat(self, websocket: WebSocket, user: str, bot
 
             # Get default collections once for performance
             if bot_config.agent.collections:
-                collection_ids = [collection.id for collection in bot_config.agent.collections]
-                db_collections = await self.db_ops.query_collections_by_ids(user, collection_ids)
-                # Convert SQLAlchemy models to Pydantic models
-                default_collections = await self._convert_db_collections_to_pydantic(db_collections)
+                agent_collection_ids = [collection.id for collection in bot_config.agent.collections]
+                agent_collections = await self.db_ops.query_collections_by_ids(agent_collection_ids)
+                for agent_collection in agent_collections:
+                    default_collections.append(view_models.Collection(
+                        id=agent_collection.id,
+                        title=agent_collection.title,
+                        description=agent_collection.description,
+                        type=agent_collection.type,
+                        status=agent_collection.status,
+                        config=parseCollectionConfig(agent_collection.config),
+                        created=agent_collection.gmt_created.isoformat(),
+                        updated=agent_collection.gmt_updated.isoformat(),
+                    ))
 
         while True:
             # Receive message from WebSocket

aperag/service/bot_service.py

@@ -24,6 +24,7 @@
 )
 from aperag.schema import view_models
 from aperag.schema.view_models import Bot, BotList
+from aperag.service import collection_service, marketplace_service
 from aperag.service.quota_service import quota_service
 
 
@@ -49,18 +50,20 @@ async def build_bot_response(self, bot: db_models.Bot) -> view_models.Bot:
             id=bot.id,
             title=bot.title,
             description=bot.description,
-            type=bot.type,
+        type=bot.type,
             config=bot_config,
             created=bot.gmt_created.isoformat(),
             updated=bot.gmt_updated.isoformat(),
         )
 
     async def validate_collections(self, user: str, bot_config: view_models.BotConfig):
+        user_collections = await self.db_ops.list_all_accessible_collections_for_user(user)
+        user_collections_ids = [collection['id'] for collection in user_collections]
         if bot_config and bot_config.agent and bot_config.agent.collections:
-            collection_ids = [collection.id for collection in bot_config.agent.collections]
-            collections = await self.db_ops.query_collections_by_ids(user, collection_ids)
-            if not collections or len(collections) != len(collection_ids):
-                raise ResourceNotFoundException("Collection", collection_ids)
+            agent_collection_ids = [collection.id for collection in bot_config.agent.collections]
+            for collection_id in agent_collection_ids:
+                if collection_id not in user_collections_ids:
+                    raise ResourceNotFoundException("Collection", collection_id)
 
     async def create_bot(
         self, user: str, bot_in: view_models.BotCreate, skip_quota_check: bool = False

aperag/service/collection_service.py

@@ -616,49 +616,6 @@ async def delete_search(self, user: str, collection_id: str, search_id: str) ->
 
         return await self.db_ops.delete_search(user, collection_id, search_id)
 
-    async def validate_collections_batch(
-        self, user: str, collections: list[view_models.Collection]
-    ) -> tuple[bool, str]:
-        """
-        Validate multiple collections in a single database call.
-
-        Args:
-            user: User identifier
-            collections: List of collection objects to validate
-
-        Returns:
-            Tuple of (is_valid, error_message). If valid, error_message is empty.
-        """
-        if not collections:
-            return True, ""
-
-        # Extract collection IDs and validate they exist
-        collection_ids = []
-        for collection in collections:
-            if not collection.id:
-                return False, "Collection object missing 'id' field"
-            collection_ids.append(collection.id)
-
-        # Remove duplicates while preserving order
-        unique_collection_ids = list(dict.fromkeys(collection_ids))
-
-        try:
-            # Single database call to get all collections
-            db_collections = await self.db_ops.query_collections_by_ids(user, unique_collection_ids)
-
-            # Create a set of found collection IDs for fast lookup
-            found_collection_ids = {str(col.id) for col in db_collections}
-
-            # Check if all requested collections were found
-            for collection_id in unique_collection_ids:
-                if collection_id not in found_collection_ids:
-                    return False, f"Collection {collection_id} not found"
-
-            return True, ""
-
-        except Exception as e:
-            return False, f"Failed to validate collections: {str(e)}"
-
     async def test_mineru_token(self, token: str) -> dict:
         """Test the MinerU API token."""
         async with httpx.AsyncClient() as client:

envs/env.template

@@ -7,11 +7,16 @@ POSTGRES_PASSWORD=postgres
 
 # Database Connection Pool Settings
 # Adjust these values based on your server resources and expected load
-DB_POOL_SIZE=20          # Base connection pool size
-DB_MAX_OVERFLOW=40       # Maximum overflow connections (total = pool_size + max_overflow)
-DB_POOL_TIMEOUT=60       # Connection timeout in seconds
-DB_POOL_RECYCLE=3600     # Recycle connections after 1 hour (in seconds)
-DB_POOL_PRE_PING=True    # Validate connections before use
+# Base connection pool size
+DB_POOL_SIZE=20
+# Maximum overflow connections (total = pool_size + max_overflow)
+DB_MAX_OVERFLOW=40
+# Connection timeout in seconds
+DB_POOL_TIMEOUT=60
+# Recycle connections after 1 hour (in seconds)
+DB_POOL_RECYCLE=3600
+# Validate connections before use
+DB_POOL_PRE_PING=True
 
 # Redis
 REDIS_HOST=127.0.0.1

web/src/app/workspace/collections/[collectionId]/documents/upload/document-upload.tsx

@@ -378,7 +378,7 @@ export const DocumentUpload = () => {
     <>
       <FileUpload
         maxFiles={1000}
-        maxSize={10 * 1024 * 1024}
+        maxSize={100 * 1024 * 1024}
         className="w-full gap-4"
         accept=".pdf,.doc,.docx,.txt,.md,.ppt,.pptx,.xls,.xlsx"
         value={documents.map((f) => f.file)}