diff --git a/ra_aid/__main__.py b/ra_aid/__main__.py index fe4c048..d91ecd5 100644 --- a/ra_aid/__main__.py +++ b/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: - # This initializes both repositories and makes them available via their respective get methods + with KeyFactRepositoryManager(db) as key_fact_repo, \ + 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 - human_input_repo = HumanInputRepository(db) - human_input_repo.create(content=initial_request, source='chat') - human_input_repo.garbage_collect() + # Using get_human_input_repository() to access the repository from context + human_input_repository = get_human_input_repository() + human_input_repository.create(content=initial_request, source='chat') + 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 - human_input_repo = HumanInputRepository(db) - human_input_repo.create(content=base_task, source='cli') + # Using get_human_input_repository() to access the repository from context + human_input_repository = get_human_input_repository() + 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)}") diff --git a/ra_aid/agent_utils.py b/ra_aid/agent_utils.py index a3af9e8..8f86519 100644 --- a/ra_aid/agent_utils.py +++ b/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() - recent_inputs = human_input_repository.get_recent(1) - if recent_inputs and len(recent_inputs) > 0: - last_human_input = recent_inputs[0].content - base_task = f"{last_human_input}\n{base_task}" + try: + human_input_repository = get_human_input_repository() + recent_inputs = human_input_repository.get_recent(1) + if recent_inputs and len(recent_inputs) > 0: + last_human_input = recent_inputs[0].content + base_task = f"{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) diff --git a/ra_aid/agents/key_facts_gc_agent.py b/ra_aid/agents/key_facts_gc_agent.py index b58adbc..e8eae2c 100644 --- a/ra_aid/agents/key_facts_gc_agent.py +++ b/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: diff --git a/ra_aid/agents/key_snippets_gc_agent.py b/ra_aid/agents/key_snippets_gc_agent.py index 2e13538..120d6e6 100644 --- a/ra_aid/agents/key_snippets_gc_agent.py +++ b/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: diff --git a/ra_aid/database/repositories/human_input_repository.py b/ra_aid/database/repositories/human_input_repository.py index 871d8a0..f8f89ff 100644 --- a/ra_aid/database/repositories/human_input_repository.py +++ b/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, initialize_database +from ra_aid.database.models import HumanInput 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() - input = repo.create("User's message", "chat") - recent_inputs = repo.get_recent(5) + with DatabaseManager() as db: + with HumanInputRepositoryManager(db) as repo: + 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() diff --git a/ra_aid/tools/human.py b/ra_aid/tools/human.py index 39f8769..dc77937 100644 --- a/ra_aid/tools/human.py +++ b/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 - human_input_repo = HumanInputRepository() + # Get the repository from context and store the input + 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 + return response \ No newline at end of file diff --git a/ra_aid/tools/memory.py b/ra_aid/tools/memory.py index ece46fa..4ae0120 100644 --- a/ra_aid/tools/memory.py +++ b/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_snippet_repository import KeySnippetRepository, get_key_snippet_repository -from ra_aid.database.repositories.human_input_repository import HumanInputRepository +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 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)}")