Merge branch 'main' of https://github.com/GeneralUserModels/gum

Author: vanihba14

gum/batcher.py

@@ -0,0 +1,99 @@
+import asyncio
+import logging
+from datetime import datetime, timezone
+from typing import List, Optional, Dict, Any
+from pathlib import Path
+import uuid
+from persistqueue import Queue
+
+class ObservationBatcher:
+    """A persistent queue for batching observations to reduce API calls."""
+    
+    def __init__(self, data_directory: str, min_batch_size: int = 5, max_batch_size: int = 50):
+        self.data_directory = Path(data_directory)
+        self.min_batch_size = min_batch_size
+        self.max_batch_size = max_batch_size
+        
+        # Create persistent queue backed by SQLite
+        queue_dir = self.data_directory / "batches"
+        queue_dir.mkdir(parents=True, exist_ok=True)
+        self._queue = Queue(path=str(queue_dir / "queue"))
+        
+        self._batch_ready_event = asyncio.Event()        
+        self.logger = logging.getLogger("gum.batcher")
+        
+    async def start(self):
+        """Start the batching system."""
+        self.logger.info(f"Started batcher with {self._queue.qsize()} items in queue")
+        
+        if self.should_process_batch():
+            self._batch_ready_event.set()
+        
+    async def stop(self):
+        """Stop the batching system."""
+        self.logger.info("Stopped batcher")
+        
+    def push(self, observer_name: str, content: str, content_type: str) -> str:
+        """Push an observation onto the queue.
+        
+        Args:
+            observer_name: Name of the observer
+            content: Observation content
+            content_type: Type of content
+            
+        Returns:
+            str: Observation ID
+        """
+        observation_id = str(uuid.uuid4())
+        observation_dict = {
+            'id': observation_id,
+            'observer_name': observer_name,
+            'content': content,
+            'content_type': content_type,
+            'timestamp': datetime.now(timezone.utc).isoformat()
+        }
+        
+        # Add to queue - automatically persisted by persist-queue
+        self._queue.put(observation_dict)
+        self.logger.debug(f"Pushed observation {observation_id} to queue (size: {self._queue.qsize()})")
+        
+        # Signal that a batch is ready if we've reached minimum size
+        if self.should_process_batch():
+            self._batch_ready_event.set()
+        
+        return observation_id
+        
+    def size(self) -> int:
+        """Get the current size of the queue."""
+        return self._queue.qsize()
+        
+    def should_process_batch(self) -> bool:
+        """Check if the batch should be processed based on minimum batch size."""
+        return self._queue.qsize() >= self.min_batch_size
+        
+    def pop_batch(self, batch_size: Optional[int] = None) -> List[Dict[str, Any]]:
+        """Pop a batch of observations from the front of the queue (FIFO).
+        
+        Args:
+            batch_size: Number of items to pop. Defaults to max_batch_size
+            
+        Returns:
+            List of observation dictionaries popped from queue
+        """
+        batch_size = batch_size or self.max_batch_size
+        
+        batch = []
+        for _ in range(min(batch_size, self._queue.qsize())):
+            batch.append(self._queue.get_nowait())
+        
+        if batch:
+            self.logger.debug(f"Popped batch of {len(batch)} observations (queue size: {self._queue.qsize()})")
+        
+        if not self.should_process_batch():
+            self._batch_ready_event.clear()
+        
+        return batch
+    
+    async def wait_for_batch_ready(self):
+        """Wait for a batch to be ready for processing."""
+        await self._batch_ready_event.wait()

gum/cli.py

@@ -4,7 +4,7 @@
 import os
 import argparse
 import asyncio
-import shutil  # Add this import for deleting directories
+import shutil  
 from gum import gum
 from gum.observers import Screen
 
@@ -29,6 +29,10 @@ def parse_args():
     parser.add_argument('--limit', '-l', type=int, help='Limit the number of results', default=10)
     parser.add_argument('--model', '-m', type=str, help='Model to use')
     parser.add_argument('--reset-cache', action='store_true', help='Reset the GUM cache and exit')  # Add this line
+    
+    # Batching configuration arguments
+    parser.add_argument('--min-batch-size', type=int, help='Minimum number of observations to trigger batch processing')
+    parser.add_argument('--max-batch-size', type=int, help='Maximum number of observations per batch')
 
     args = parser.parse_args()
 
@@ -53,7 +57,11 @@ async def main():
     model = args.model or os.getenv('MODEL_NAME') or 'gpt-4o-mini'
     user_name = args.user_name or os.getenv('USER_NAME')
 
-    # you need one or the other-
+    # Batching configuration - follow same pattern as other args    
+    min_batch_size = args.min_batch_size or int(os.getenv('MIN_BATCH_SIZE', '5'))
+    max_batch_size = args.max_batch_size or int(os.getenv('MAX_BATCH_SIZE', '15'))
+
+    # you need one or the other
     if user_name is None and args.query is None:
         print("Please provide a user name (as an argument, -u, or as an env variable) or a query (as an argument, -q)")
         return
@@ -63,7 +71,7 @@ async def main():
         await gum_instance.connect_db()
         result = await gum_instance.query(args.query, limit=args.limit)
 
-        # pretty print confidences / propositions / number of items returned
+        # confidences / propositions / number of items returned
         print(f"\nFound {len(result)} results:")
         for prop, score in result:
             print(f"\nProposition: {prop.text}")
@@ -75,7 +83,14 @@ async def main():
             print("-" * 80)
     else:
         print(f"Listening to {user_name} with model {model}")
-        async with gum(user_name, model, Screen(model)) as gum_instance:
+            
+        async with gum(
+            user_name, 
+            model, 
+            Screen(model),
+            min_batch_size=min_batch_size,
+            max_batch_size=max_batch_size
+        ) as gum_instance:
             await asyncio.Future()  # run forever (Ctrl-C to stop)
 
 def cli():

gum/db_utils.py

@@ -26,7 +26,6 @@
 from .models import (
     Observation,
     Proposition,
-    proposition_parent,
     observation_proposition,
 )
 
@@ -45,13 +44,7 @@ def build_fts_query(raw: str, mode: str = "OR") -> str:
     else:  # implicit AND
         return " ".join(tokens)
 
-def _has_child_subquery() -> select:
-    return (
-        select(literal_column("1"))
-        .select_from(proposition_parent)
-        .where(proposition_parent.c.parent_id == Proposition.id)
-        .exists()
-    )
+
 
 
 async def search_propositions_bm25(
@@ -74,7 +67,6 @@ async def search_propositions_bm25(
     # 1  Build candidate list
     # --------------------------------------------------------
     candidate_pool = limit * 10 if enable_mmr else limit
-    has_child      = _has_child_subquery()
 
     if has_query:
         fts_prop = Table("propositions_fts", MetaData())
@@ -143,14 +135,12 @@ async def search_propositions_bm25(
         stmt = (
             select(Proposition, best_scores.c.bm25)
             .join(best_scores, best_scores.c.pid == Proposition.id)
-            .where(~has_child)
             .order_by(best_scores.c.bm25.asc())          # smallest→best
         )
     else:
         # --- 1-b  No user query ------------------------------
         stmt = (
             select(Proposition, literal_column("0.0").label("bm25"))
-            .where(~has_child)
             .order_by(Proposition.created_at.desc())
         )