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)}")