Upload 2 files

data/partition_loader.py

@@ -119,7 +119,7 @@ class SinglePartitionStore:
         if not HAS_NUMPY or id_to_tag is None:
             return Counter()
 
-        if event_indices is None or len(event_indices) == 0:
+        if event_indices is None:
             # Total tag counts
             return Counter({
                 id_to_tag[tag_id]: len(events)

data/tag_store.py

@@ -35,12 +35,13 @@ class QuickSearchResult:
             self.tags = []
 
 
+
 class TagStore:
     """
     High-level interface for Quick Search functionality.
 
     Provides:
-    - Partition loading and management
+    - Partition loading and management (supports multiple active partitions)
     - Tag filtering by rating and person category
     - Include/Exclude tag filtering
     - Random event sampling for prompt generation
@@ -66,38 +67,38 @@ class TagStore:
     def load_partition(self, rating: str, person: str) -> bool:
         """
         Load partition for given rating and person category.
-
+        
         Args:
             rating: 'g', 's', 'q', or 'e'
             person: Person category like '1girl_solo'
-
+            
         Returns:
             True if loaded successfully
         """
         partition_name = self._manager.get_partition_filename(rating, person)
+        
+        # If already loaded, do nothing
+        if self._current_partition_name == partition_name and self._current_partition is not None:
+            return True
 
-        if partition_name == self._current_partition_name and self._current_partition is not None:
-            return True  # Already loaded
-
-        # Unload previous partition to save memory
-        if self._current_partition_name:
-            self._manager.unload_partition(self._current_partition_name)
-
+        # Load new partition
         partition = self._manager.load_partition(partition_name)
-        if partition is None:
-            self._current_partition = None
-            self._current_partition_name = ""
-            return False
-
-        self._current_partition = partition
-        self._current_partition_name = partition_name
-        return True
+        if partition:
+            # Unload previous if different
+            if self._current_partition_name and self._current_partition_name != partition_name:
+                self._manager.unload_partition(self._current_partition_name)
+            
+            self._current_partition = partition
+            self._current_partition_name = partition_name
+            return True
+            
+        return False
 
     def get_event_count(self) -> int:
         """Get number of events in current partition"""
-        if self._current_partition is None:
-            return 0
-        return self._current_partition.num_events
+        if self._current_partition:
+            return self._current_partition.num_events
+        return 0
 
     def get_filtered_event_count(
         self,
@@ -105,20 +106,19 @@ class TagStore:
         exclude_tags: Optional[List[str]] = None
     ) -> int:
         """Get number of events matching filter criteria"""
-        if self._current_partition is None:
+        if not self._current_partition:
             return 0
 
         metadata = self._manager.get_metadata()
         if metadata is None:
             return 0
 
-        filtered = self._current_partition.filter_events(
+        filtered_indices = self._current_partition.filter_events(
             required_tags=include_tags,
             excluded_tags=exclude_tags,
             tag_to_id=metadata.tag_to_id
         )
-
-        return len(filtered)
+        return len(filtered_indices)
 
     def get_top_tags(
         self,
@@ -129,17 +129,14 @@ class TagStore:
     ) -> List[TagInfo]:
         """
         Get most frequent tags in current partition with filters applied.
-
-        Returns tags sorted by frequency (descending).
         """
-        if self._current_partition is None:
+        if not self._current_partition:
             return []
 
         metadata = self._manager.get_metadata()
         if metadata is None:
             return []
 
-        # Filter events
         filtered_indices = self._current_partition.filter_events(
             required_tags=include_tags,
             excluded_tags=exclude_tags,
@@ -149,23 +146,25 @@ class TagStore:
         if len(filtered_indices) == 0:
             return []
 
-        # Count tags in filtered events
+        # Get counts for these events
         tag_counts = self._current_partition.get_tag_counts(
             event_indices=filtered_indices,
             id_to_tag=metadata.id_to_tag
         )
 
-        # Exclude already included/excluded tags from results
+        # Remove excluded tags and sort
         excluded_set = set(include_tags or []) | set(exclude_tags or [])
-
-        # Sort by count and return with pagination
+        
         sorted_tags = sorted(
             [(tag, count) for tag, count in tag_counts.items() if tag not in excluded_set],
             key=lambda x: x[1],
             reverse=True
-        )[offset:offset + limit]
+        )
 
-        return [TagInfo(tag=tag, count=count) for tag, count in sorted_tags]
+        # Pagination
+        paged_tags = sorted_tags[offset:offset + limit]
+        
+        return [TagInfo(tag=tag, count=count) for tag, count in paged_tags]
 
     def get_total_tag_count(
         self,
@@ -173,14 +172,13 @@ class TagStore:
         exclude_tags: Optional[List[str]] = None
     ) -> int:
         """Get total number of unique tags matching filter criteria"""
-        if self._current_partition is None:
+        if not self._current_partition:
             return 0
 
         metadata = self._manager.get_metadata()
         if metadata is None:
             return 0
 
-        # Filter events
         filtered_indices = self._current_partition.filter_events(
             required_tags=include_tags,
             excluded_tags=exclude_tags,
@@ -190,15 +188,13 @@ class TagStore:
         if len(filtered_indices) == 0:
             return 0
 
-        # Count tags in filtered events
         tag_counts = self._current_partition.get_tag_counts(
             event_indices=filtered_indices,
             id_to_tag=metadata.id_to_tag
         )
-
-        # Exclude already included/excluded tags
+        
         excluded_set = set(include_tags or []) | set(exclude_tags or [])
-        return len([tag for tag in tag_counts.keys() if tag not in excluded_set])
+        return len([t for t in tag_counts if t not in excluded_set])
 
     def generate_random_prompt(
         self,
@@ -208,67 +204,43 @@ class TagStore:
         exclude_tags: Optional[List[str]] = None
     ) -> QuickSearchResult:
         """
-        Generate a random prompt from the partition.
-
-        Args:
-            rating: Rating code ('g', 's', 'q', 'e')
-            person: Person category
-            include_tags: Tags that must be present
-            exclude_tags: Tags that must not be present
-
-        Returns:
-            QuickSearchResult with generated prompt
+        Generate a random prompt from the dataset based on criteria.
         """
-        # Load partition
+        # Ensure correct partition is loaded
         if not self.load_partition(rating, person):
-            return QuickSearchResult(
-                success=False,
-                error_message=f"Failed to load partition: {rating}_{person}"
-            )
+            return QuickSearchResult(success=False, error_message=f"Failed to load partition for {rating}/{person}")
 
         metadata = self._manager.get_metadata()
         if metadata is None:
-            return QuickSearchResult(
-                success=False,
-                error_message="Metadata not available"
-            )
-
-        # Get auto-tags for person category
-        auto_tags = PERSON_AUTO_TAGS.get(person, [])
-
-        # Combine include tags with auto tags
-        all_include = list(set((include_tags or []) + auto_tags))
+            return QuickSearchResult(success=False, error_message="Metadata not available")
 
         # Filter events
         filtered_indices = self._current_partition.filter_events(
-            required_tags=all_include if all_include else None,
+            required_tags=include_tags,
             excluded_tags=exclude_tags,
             tag_to_id=metadata.tag_to_id
         )
 
         if len(filtered_indices) == 0:
             return QuickSearchResult(
-                success=False,
+                success=False, 
                 error_message="No events match the filter criteria",
                 event_count=0
             )
 
-        # Select random event
-        random_idx = random.choice(filtered_indices)
-
-        # Get tags for this event
+        # Pick random event
+        choice_idx = random.choice(filtered_indices)
+        
+        # Get tags for event
         event_tags = self._current_partition.get_event_tags(
-            event_idx=int(random_idx),
+            event_idx=int(choice_idx),
             id_to_tag=metadata.id_to_tag
         )
 
         if not event_tags:
-            return QuickSearchResult(
-                success=False,
-                error_message="Failed to get tags for selected event"
-            )
+            return QuickSearchResult(success=False, error_message="Failed to retrieve tags for selected event")
 
-        # Convert to list and create prompt
+        # Format prompt
         tags_list = sorted(list(event_tags))
         prompt = ", ".join(tags_list)
 
@@ -286,13 +258,6 @@ class TagStore:
     ) -> List[TagInfo]:
         """
         Search for tags matching query string.
-
-        Args:
-            query: Search query (partial match)
-            limit: Maximum results
-
-        Returns:
-            List of matching tags with frequencies
         """
         metadata = self._manager.get_metadata()
         if metadata is None or not query: