zach
/
RA.Aid
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

use context vars for human input repo

This commit is contained in:
AI Christianson 2025-03-03 17:58:55 -05:00
parent fc58aa0b77
commit 5202d2e7f3
7 changed files with 136 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
ra_aid/__main__.py
View File

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

16
ra_aid/agent_utils.py
View File

@ -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>{last_human_input}</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>{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)

7
ra_aid/agents/key_facts_gc_agent.py
View File

@ -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:

7
ra_aid/agents/key_snippets_gc_agent.py
View File

@ -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:

105
ra_aid/database/repositories/human_input_repository.py
View File

@ -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()

12
ra_aid/tools/human.py
View File

@ -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

14
ra_aid/tools/memory.py
View File

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