use context vars for human input repo

ra_aid/__main__.py

@@ -46,6 +46,9 @@ from ra_aid.database.repositories.key_fact_repository import KeyFactRepositoryMa
 from ra_aid.database.repositories.key_snippet_repository import (
     KeySnippetRepositoryManager, get_key_snippet_repository
 )
+from ra_aid.database.repositories.human_input_repository import (
+    HumanInputRepositoryManager, get_human_input_repository
+)
 from ra_aid.model_formatters import format_key_facts_dict
 from ra_aid.model_formatters.key_snippets_formatter import format_key_snippets_dict
 from ra_aid.console.output import cpm
@@ -396,10 +399,13 @@ def main():
                 logger.error(f"Database migration error: {str(e)}")
 
             # Initialize repositories with database connection
-            with KeyFactRepositoryManager(db) as key_fact_repo, KeySnippetRepositoryManager(db) as key_snippet_repo:
+            with KeyFactRepositoryManager(db) as key_fact_repo, \
-                # This initializes both repositories and makes them available via their respective get methods
+                 KeySnippetRepositoryManager(db) as key_snippet_repo, \
+                 HumanInputRepositoryManager(db) as human_input_repo:
+                # This initializes all repositories and makes them available via their respective get methods
                 logger.debug("Initialized KeyFactRepository")
                 logger.debug("Initialized KeySnippetRepository")
+                logger.debug("Initialized HumanInputRepository")
 
                 # Check dependencies before proceeding
                 check_dependencies()
@@ -479,10 +485,10 @@ def main():
                     # Record chat input in database (redundant as ask_human already records it,
                     # but needed in case the ask_human implementation changes)
                     try:
-                        from ra_aid.database.repositories.human_input_repository import HumanInputRepository
+                        # Using get_human_input_repository() to access the repository from context
-                        human_input_repo = HumanInputRepository(db)
+                        human_input_repository = get_human_input_repository()
-                        human_input_repo.create(content=initial_request, source='chat')
+                        human_input_repository.create(content=initial_request, source='chat')
-                        human_input_repo.garbage_collect()
+                        human_input_repository.garbage_collect()
                     except Exception as e:
                         logger.error(f"Failed to record initial chat input: {str(e)}")
 
@@ -552,11 +558,11 @@ def main():
                 
                 # Record CLI input in database
                 try:
-                    from ra_aid.database.repositories.human_input_repository import HumanInputRepository
+                    # Using get_human_input_repository() to access the repository from context
-                    human_input_repo = HumanInputRepository(db)
+                    human_input_repository = get_human_input_repository()
-                    human_input_repo.create(content=base_task, source='cli')
+                    human_input_repository.create(content=base_task, source='cli')
                     # Run garbage collection to ensure we don't exceed 100 inputs
-                    human_input_repo.garbage_collect()
+                    human_input_repository.garbage_collect()
                     logger.debug(f"Recorded CLI input: {base_task}")
                 except Exception as e:
                     logger.error(f"Failed to record CLI input: {str(e)}")

ra_aid/agent_utils.py

@@ -86,7 +86,7 @@ from ra_aid.tool_configs import (
 from ra_aid.tools.handle_user_defined_test_cmd_execution import execute_test_command
 from ra_aid.database.repositories.key_fact_repository import get_key_fact_repository
 from ra_aid.database.repositories.key_snippet_repository import get_key_snippet_repository
-from ra_aid.database.repositories.human_input_repository import HumanInputRepository
+from ra_aid.database.repositories.human_input_repository import get_human_input_repository
 from ra_aid.model_formatters import format_key_facts_dict
 from ra_aid.model_formatters.key_snippets_formatter import format_key_snippets_dict
 from ra_aid.tools.memory import (
@@ -402,11 +402,15 @@ def run_research_agent(
 
     # Get the last human input, if it exists
     base_task = base_task_or_query
-    human_input_repository = HumanInputRepository()
+    try:
-    recent_inputs = human_input_repository.get_recent(1)
+        human_input_repository = get_human_input_repository()
-    if recent_inputs and len(recent_inputs) > 0:
+        recent_inputs = human_input_repository.get_recent(1)
-        last_human_input = recent_inputs[0].content
+        if recent_inputs and len(recent_inputs) > 0:
-        base_task = f"<last human input>{last_human_input}</last human input>\n{base_task}"
+            last_human_input = recent_inputs[0].content
+            base_task = f"<last human input>{last_human_input}</last human input>\n{base_task}"
+    except RuntimeError as e:
+        logger.error(f"Failed to access human input repository: {str(e)}")
+        # Continue without appending last human input
 
     try:
         project_info = get_project_info(".", file_limit=2000)

ra_aid/agents/key_facts_gc_agent.py

@@ -18,14 +18,13 @@ logger = logging.getLogger(__name__)
 
 from ra_aid.agent_utils import create_agent, run_agent_with_retry
 from ra_aid.database.repositories.key_fact_repository import get_key_fact_repository
-from ra_aid.database.repositories.human_input_repository import HumanInputRepository
+from ra_aid.database.repositories.human_input_repository import get_human_input_repository
 from ra_aid.llm import initialize_llm
 from ra_aid.prompts.key_facts_gc_prompts import KEY_FACTS_GC_PROMPT
 from ra_aid.tools.memory import log_work_event, _global_memory
 
 
 console = Console()
-human_input_repository = HumanInputRepository()
 
 
 @tool
@@ -46,7 +45,7 @@ def delete_key_facts(fact_ids: List[int]) -> str:
     # Try to get the current human input to protect its facts
     current_human_input_id = None
     try:
-        recent_inputs = human_input_repository.get_recent(1)
+        recent_inputs = get_human_input_repository().get_recent(1)
         if recent_inputs and len(recent_inputs) > 0:
             current_human_input_id = recent_inputs[0].id
     except Exception as e:
@@ -128,7 +127,7 @@ def run_key_facts_gc_agent() -> None:
         # Try to get the current human input ID to exclude its facts
         current_human_input_id = None
         try:
-            recent_inputs = human_input_repository.get_recent(1)
+            recent_inputs = get_human_input_repository().get_recent(1)
             if recent_inputs and len(recent_inputs) > 0:
                 current_human_input_id = recent_inputs[0].id
         except Exception as e:

ra_aid/agents/key_snippets_gc_agent.py

@@ -15,14 +15,13 @@ from rich.panel import Panel
 
 from ra_aid.agent_utils import create_agent, run_agent_with_retry
 from ra_aid.database.repositories.key_snippet_repository import get_key_snippet_repository
-from ra_aid.database.repositories.human_input_repository import HumanInputRepository
+from ra_aid.database.repositories.human_input_repository import get_human_input_repository
 from ra_aid.llm import initialize_llm
 from ra_aid.prompts.key_snippets_gc_prompts import KEY_SNIPPETS_GC_PROMPT
 from ra_aid.tools.memory import log_work_event, _global_memory
 
 
 console = Console()
-human_input_repository = HumanInputRepository()
 
 
 @tool
@@ -44,7 +43,7 @@ def delete_key_snippets(snippet_ids: List[int]) -> str:
     # Try to get the current human input to protect its snippets
     current_human_input_id = None
     try:
-        recent_inputs = human_input_repository.get_recent(1)
+        recent_inputs = get_human_input_repository().get_recent(1)
         if recent_inputs and len(recent_inputs) > 0:
             current_human_input_id = recent_inputs[0].id
     except Exception as e:
@@ -120,7 +119,7 @@ def run_key_snippets_gc_agent() -> None:
         # Try to get the current human input ID to exclude its snippets
         current_human_input_id = None
         try:
-            recent_inputs = human_input_repository.get_recent(1)
+            recent_inputs = get_human_input_repository().get_recent(1)
             if recent_inputs and len(recent_inputs) > 0:
                 current_human_input_id = recent_inputs[0].id
         except Exception as e:

ra_aid/database/repositories/human_input_repository.py

@@ -6,15 +6,93 @@ following the repository pattern for data access abstraction.
 """
 
 from typing import Dict, List, Optional
+import contextvars
 
 import peewee
 
-from ra_aid.database.connection import get_db
+from ra_aid.database.models import HumanInput
-from ra_aid.database.models import HumanInput, initialize_database
 from ra_aid.logging_config import get_logger
 
 logger = get_logger(__name__)
 
+# Create contextvar to hold the HumanInputRepository instance
+human_input_repo_var = contextvars.ContextVar("human_input_repo", default=None)
+
+
+class HumanInputRepositoryManager:
+    """
+    Context manager for HumanInputRepository.
+
+    This class provides a context manager interface for HumanInputRepository,
+    using the contextvars approach for thread safety.
+
+    Example:
+        with DatabaseManager() as db:
+            with HumanInputRepositoryManager(db) as repo:
+                # Use the repository
+                input_record = repo.create(content="User input", source="chat")
+                recent_inputs = repo.get_recent(5)
+    """
+
+    def __init__(self, db):
+        """
+        Initialize the HumanInputRepositoryManager.
+
+        Args:
+            db: Database connection to use (required)
+        """
+        self.db = db
+
+    def __enter__(self) -> 'HumanInputRepository':
+        """
+        Initialize the HumanInputRepository and return it.
+
+        Returns:
+            HumanInputRepository: The initialized repository
+        """
+        repo = HumanInputRepository(self.db)
+        human_input_repo_var.set(repo)
+        return repo
+
+    def __exit__(
+        self,
+        exc_type: Optional[type],
+        exc_val: Optional[Exception],
+        exc_tb: Optional[object],
+    ) -> None:
+        """
+        Reset the repository when exiting the context.
+
+        Args:
+            exc_type: The exception type if an exception was raised
+            exc_val: The exception value if an exception was raised
+            exc_tb: The traceback if an exception was raised
+        """
+        # Reset the contextvar to None
+        human_input_repo_var.set(None)
+
+        # Don't suppress exceptions
+        return False
+
+
+def get_human_input_repository() -> 'HumanInputRepository':
+    """
+    Get the current HumanInputRepository instance.
+
+    Returns:
+        HumanInputRepository: The current repository instance
+        
+    Raises:
+        RuntimeError: If no repository has been initialized with HumanInputRepositoryManager
+    """
+    repo = human_input_repo_var.get()
+    if repo is None:
+        raise RuntimeError(
+            "No HumanInputRepository available. "
+            "Make sure to initialize one with HumanInputRepositoryManager first."
+        )
+    return repo
+
 
 class HumanInputRepository:
     """
@@ -24,18 +102,21 @@ class HumanInputRepository:
     abstracting the database access details from the business logic.
     
     Example:
-        repo = HumanInputRepository()
+        with DatabaseManager() as db:
-        input = repo.create("User's message", "chat")
+            with HumanInputRepositoryManager(db) as repo:
-        recent_inputs = repo.get_recent(5)
+                input_record = repo.create("User's message", "chat")
+                recent_inputs = repo.get_recent(5)
     """
     
-    def __init__(self, db=None):
+    def __init__(self, db):
         """
-        Initialize the repository with an optional database connection.
+        Initialize the repository with a database connection.
         
         Args:
-            db: Optional database connection to use. If None, will use initialize_database()
+            db: Database connection to use (required)
         """
+        if db is None:
+            raise ValueError("Database connection is required for HumanInputRepository")
         self.db = db
     
     def create(self, content: str, source: str) -> HumanInput:
@@ -53,7 +134,6 @@ class HumanInputRepository:
             peewee.DatabaseError: If there's an error creating the record
         """
         try:
-            db = self.db if self.db is not None else initialize_database()
             input_record = HumanInput.create(content=content, source=source)
             logger.debug(f"Created human input ID {input_record.id} from {source}")
             return input_record
@@ -75,7 +155,6 @@ class HumanInputRepository:
             peewee.DatabaseError: If there's an error accessing the database
         """
         try:
-            db = self.db if self.db is not None else initialize_database()
             return HumanInput.get_or_none(HumanInput.id == input_id)
         except peewee.DatabaseError as e:
             logger.error(f"Failed to fetch human input {input_id}: {str(e)}")
@@ -97,7 +176,6 @@ class HumanInputRepository:
             peewee.DatabaseError: If there's an error updating the record
         """
         try:
-            db = self.db if self.db is not None else initialize_database()
             # First check if the record exists
             input_record = self.get(input_id)
             if not input_record:
@@ -131,7 +209,6 @@ class HumanInputRepository:
             peewee.DatabaseError: If there's an error deleting the record
         """
         try:
-            db = self.db if self.db is not None else initialize_database()
             # First check if the record exists
             input_record = self.get(input_id)
             if not input_record:
@@ -157,7 +234,6 @@ class HumanInputRepository:
             peewee.DatabaseError: If there's an error accessing the database
         """
         try:
-            db = self.db if self.db is not None else initialize_database()
             return list(HumanInput.select().order_by(HumanInput.created_at.desc()))
         except peewee.DatabaseError as e:
             logger.error(f"Failed to fetch all human inputs: {str(e)}")
@@ -177,7 +253,6 @@ class HumanInputRepository:
             peewee.DatabaseError: If there's an error accessing the database
         """
         try:
-            db = self.db if self.db is not None else initialize_database()
             return list(HumanInput.select().order_by(HumanInput.created_at.desc()).limit(limit))
         except peewee.DatabaseError as e:
             logger.error(f"Failed to fetch recent human inputs: {str(e)}")
@@ -197,7 +272,6 @@ class HumanInputRepository:
             peewee.DatabaseError: If there's an error accessing the database
         """
         try:
-            db = self.db if self.db is not None else initialize_database()
             return list(HumanInput.select().where(HumanInput.source == source).order_by(HumanInput.created_at.desc()))
         except peewee.DatabaseError as e:
             logger.error(f"Failed to fetch human inputs by source {source}: {str(e)}")
@@ -216,7 +290,6 @@ class HumanInputRepository:
             peewee.DatabaseError: If there's an error accessing the database
         """
         try:
-            db = self.db if self.db is not None else initialize_database()
             # Get the count of records
             record_count = HumanInput.select().count()
             

ra_aid/tools/human.py

@@ -57,7 +57,7 @@ def ask_human(question: str) -> str:
     
     # Record human response in database
     try:
-        from ra_aid.database.repositories.human_input_repository import HumanInputRepository
+        from ra_aid.database.repositories.human_input_repository import get_human_input_repository
         from ra_aid.tools.memory import _global_memory
         
         # Determine the source based on context
@@ -70,15 +70,15 @@ def ask_human(question: str) -> str:
         else:
             source = "chat"  # Default fallback
             
-        # Store the input
+        # Get the repository from context and store the input
-        human_input_repo = HumanInputRepository()
+        human_input_repo = get_human_input_repository()
         human_input_repo.create(content=response, source=source)
         
         # Run garbage collection to ensure we don't exceed 100 inputs
         human_input_repo.garbage_collect()
+    except RuntimeError as e:
+        logger.error(f"Failed to record human input: No HumanInputRepository available in context. {str(e)}")
     except Exception as e:
-        from ra_aid.logging_config import get_logger
-        logger = get_logger(__name__)
         logger.error(f"Failed to record human input: {str(e)}")
     
     return response

ra_aid/tools/memory.py

@@ -17,9 +17,9 @@ from ra_aid.agent_context import (
     mark_should_exit,
     mark_task_completed,
 )
-from ra_aid.database.repositories.key_fact_repository import KeyFactRepository, get_key_fact_repository
+from ra_aid.database.repositories.key_fact_repository import get_key_fact_repository
-from ra_aid.database.repositories.key_snippet_repository import KeySnippetRepository, get_key_snippet_repository
+from ra_aid.database.repositories.key_snippet_repository import get_key_snippet_repository
-from ra_aid.database.repositories.human_input_repository import HumanInputRepository
+from ra_aid.database.repositories.human_input_repository import get_human_input_repository
 from ra_aid.model_formatters import key_snippets_formatter
 from ra_aid.logging_config import get_logger
 
@@ -114,11 +114,13 @@ def emit_key_facts(facts: List[str]) -> str:
     
     # Try to get the latest human input
     human_input_id = None
-    human_input_repo = HumanInputRepository()
     try:
+        human_input_repo = get_human_input_repository()
         recent_inputs = human_input_repo.get_recent(1)
         if recent_inputs and len(recent_inputs) > 0:
             human_input_id = recent_inputs[0].id
+    except RuntimeError as e:
+        logger.warning(f"No HumanInputRepository available: {str(e)}")
     except Exception as e:
         logger.warning(f"Failed to get recent human input: {str(e)}")
     
@@ -225,10 +227,12 @@ def emit_key_snippet(snippet_info: SnippetInfo) -> str:
     # Try to get the latest human input
     human_input_id = None
     try:
-        human_input_repo = HumanInputRepository()
+        human_input_repo = get_human_input_repository()
         recent_inputs = human_input_repo.get_recent(1)
         if recent_inputs and len(recent_inputs) > 0:
             human_input_id = recent_inputs[0].id
+    except RuntimeError as e:
+        logger.warning(f"No HumanInputRepository available: {str(e)}")
     except Exception as e:
         logger.warning(f"Failed to get recent human input: {str(e)}")