From f88ad5bc7a7684cdaec861de711f6a67739cbdbe Mon Sep 17 00:00:00 2001 From: AI Christianson Date: Sun, 2 Mar 2025 20:18:00 -0500 Subject: [PATCH] associate key facts and snippets with latest human input --- ra_aid/database/models.py | 36 +++++----- .../repositories/key_fact_repository.py | 5 +- .../repositories/key_snippet_repository.py | 7 +- ...250302_201611_add_human_input_reference.py | 67 +++++++++++++++++++ ra_aid/tools/memory.py | 26 ++++++- 5 files changed, 119 insertions(+), 22 deletions(-) create mode 100644 ra_aid/migrations/005_20250302_201611_add_human_input_reference.py diff --git a/ra_aid/database/models.py b/ra_aid/database/models.py index 8f2dacc..2fb2f06 100644 --- a/ra_aid/database/models.py +++ b/ra_aid/database/models.py @@ -99,6 +99,22 @@ class BaseModel(peewee.Model): raise +class HumanInput(BaseModel): + """ + Model representing human input stored in the database. + + Human inputs are text inputs provided by users through various interfaces + such as CLI, chat, or HIL (human-in-the-loop). This model tracks these inputs + along with their source for analysis and reference. + """ + content = peewee.TextField() + source = peewee.TextField() # 'cli', 'chat', or 'hil' + # created_at and updated_at are inherited from BaseModel + + class Meta: + table_name = "human_input" + + class KeyFact(BaseModel): """ Model representing a key fact stored in the database. @@ -107,6 +123,7 @@ class KeyFact(BaseModel): that need to be referenced later. """ content = peewee.TextField() + human_input = peewee.ForeignKeyField(HumanInput, backref='key_facts', null=True) # created_at and updated_at are inherited from BaseModel class Meta: @@ -125,23 +142,8 @@ class KeySnippet(BaseModel): line_number = peewee.IntegerField() snippet = peewee.TextField() description = peewee.TextField(null=True) + human_input = peewee.ForeignKeyField(HumanInput, backref='key_snippets', null=True) # created_at and updated_at are inherited from BaseModel class Meta: - table_name = "key_snippet" - - -class HumanInput(BaseModel): - """ - Model representing human input stored in the database. - - Human inputs are text inputs provided by users through various interfaces - such as CLI, chat, or HIL (human-in-the-loop). This model tracks these inputs - along with their source for analysis and reference. - """ - content = peewee.TextField() - source = peewee.TextField() # 'cli', 'chat', or 'hil' - # created_at and updated_at are inherited from BaseModel - - class Meta: - table_name = "human_input" \ No newline at end of file + table_name = "key_snippet" \ No newline at end of file diff --git a/ra_aid/database/repositories/key_fact_repository.py b/ra_aid/database/repositories/key_fact_repository.py index 1dfa824..1ad75dc 100644 --- a/ra_aid/database/repositories/key_fact_repository.py +++ b/ra_aid/database/repositories/key_fact_repository.py @@ -38,12 +38,13 @@ class KeyFactRepository: """ self.db = db - def create(self, content: str) -> KeyFact: + def create(self, content: str, human_input_id: Optional[int] = None) -> KeyFact: """ Create a new key fact in the database. Args: content: The text content of the key fact + human_input_id: Optional ID of the associated human input Returns: KeyFact: The newly created key fact instance @@ -53,7 +54,7 @@ class KeyFactRepository: """ try: db = self.db if self.db is not None else initialize_database() - fact = KeyFact.create(content=content) + fact = KeyFact.create(content=content, human_input_id=human_input_id) logger.debug(f"Created key fact ID {fact.id}: {content}") return fact except peewee.DatabaseError as e: diff --git a/ra_aid/database/repositories/key_snippet_repository.py b/ra_aid/database/repositories/key_snippet_repository.py index f4d6f2c..24bbdaa 100644 --- a/ra_aid/database/repositories/key_snippet_repository.py +++ b/ra_aid/database/repositories/key_snippet_repository.py @@ -44,7 +44,8 @@ class KeySnippetRepository: self.db = db def create( - self, filepath: str, line_number: int, snippet: str, description: Optional[str] = None + self, filepath: str, line_number: int, snippet: str, description: Optional[str] = None, + human_input_id: Optional[int] = None ) -> KeySnippet: """ Create a new key snippet in the database. @@ -54,6 +55,7 @@ class KeySnippetRepository: line_number: Line number where the snippet starts snippet: The source code snippet text description: Optional description of the significance + human_input_id: Optional ID of the associated human input Returns: KeySnippet: The newly created key snippet instance @@ -67,7 +69,8 @@ class KeySnippetRepository: filepath=filepath, line_number=line_number, snippet=snippet, - description=description + description=description, + human_input_id=human_input_id ) logger.debug(f"Created key snippet ID {key_snippet.id}: {filepath}:{line_number}") return key_snippet diff --git a/ra_aid/migrations/005_20250302_201611_add_human_input_reference.py b/ra_aid/migrations/005_20250302_201611_add_human_input_reference.py new file mode 100644 index 0000000..8827dcc --- /dev/null +++ b/ra_aid/migrations/005_20250302_201611_add_human_input_reference.py @@ -0,0 +1,67 @@ +"""Peewee migrations -- 005_20250302_201611_add_human_input_reference.py. + +Some examples (model - class or model name):: + + > Model = migrator.orm['table_name'] # Return model in current state by name + > Model = migrator.ModelClass # Return model in current state by name + + > migrator.sql(sql) # Run custom SQL + > migrator.run(func, *args, **kwargs) # Run python function with the given args + > migrator.create_model(Model) # Create a model (could be used as decorator) + > migrator.remove_model(model, cascade=True) # Remove a model + > migrator.add_fields(model, **fields) # Add fields to a model + > migrator.change_fields(model, **fields) # Change fields + > migrator.remove_fields(model, *field_names, cascade=True) + > migrator.rename_field(model, old_field_name, new_field_name) + > migrator.rename_table(model, new_table_name) + > migrator.add_index(model, *col_names, unique=False) + > migrator.add_not_null(model, *field_names) + > migrator.add_default(model, field_name, default) + > migrator.add_constraint(model, name, sql) + > migrator.drop_index(model, *col_names) + > migrator.drop_not_null(model, *field_names) + > migrator.drop_constraints(model, *constraints) + +""" + +from contextlib import suppress + +import peewee as pw +from peewee_migrate import Migrator + + +with suppress(ImportError): + import playhouse.postgres_ext as pw_pext + + +def migrate(migrator: Migrator, database: pw.Database, *, fake=False): + """Add human_input_id foreign key field to KeyFact and KeySnippet tables.""" + + # Add human_input_id field to KeyFact model + migrator.add_fields( + 'key_fact', + human_input_id=pw.ForeignKeyField( + 'human_input', + null=True, + field='id', + on_delete='SET NULL' + ) + ) + + # Add human_input_id field to KeySnippet model + migrator.add_fields( + 'key_snippet', + human_input_id=pw.ForeignKeyField( + 'human_input', + null=True, + field='id', + on_delete='SET NULL' + ) + ) + + +def rollback(migrator: Migrator, database: pw.Database, *, fake=False): + """Remove human_input_id field from KeyFact and KeySnippet tables.""" + + migrator.remove_fields('key_fact', 'human_input_id') + migrator.remove_fields('key_snippet', 'human_input_id') \ No newline at end of file diff --git a/ra_aid/tools/memory.py b/ra_aid/tools/memory.py index c2f1c88..3e52bad 100644 --- a/ra_aid/tools/memory.py +++ b/ra_aid/tools/memory.py @@ -19,6 +19,7 @@ from ra_aid.agent_context import ( ) from ra_aid.database.repositories.key_fact_repository import KeyFactRepository from ra_aid.database.repositories.key_snippet_repository import KeySnippetRepository +from ra_aid.database.repositories.human_input_repository import HumanInputRepository from ra_aid.model_formatters import key_snippets_formatter from ra_aid.logging_config import get_logger @@ -45,6 +46,9 @@ key_fact_repository = KeyFactRepository() # Initialize repository for key snippets key_snippet_repository = KeySnippetRepository() +# Initialize repository for human inputs +human_input_repository = HumanInputRepository() + # Global memory store _global_memory: Dict[str, Any] = { "research_notes": [], @@ -115,9 +119,19 @@ def emit_key_facts(facts: List[str]) -> str: facts: List of key facts to store """ results = [] + + # Try to get the latest human input + human_input_id = None + try: + recent_inputs = human_input_repository.get_recent(1) + if recent_inputs and len(recent_inputs) > 0: + human_input_id = recent_inputs[0].id + except Exception as e: + logger.warning(f"Failed to get recent human input: {str(e)}") + for fact in facts: # Create fact in database using repository - created_fact = key_fact_repository.create(fact) + created_fact = key_fact_repository.create(fact, human_input_id=human_input_id) fact_id = created_fact.id # Display panel with ID @@ -207,12 +221,22 @@ def emit_key_snippet(snippet_info: SnippetInfo) -> str: # Add filepath to related files emit_related_files.invoke({"files": [snippet_info["filepath"]]}) + # Try to get the latest human input + human_input_id = None + try: + recent_inputs = human_input_repository.get_recent(1) + if recent_inputs and len(recent_inputs) > 0: + human_input_id = recent_inputs[0].id + except Exception as e: + logger.warning(f"Failed to get recent human input: {str(e)}") + # Create a new key snippet in the database key_snippet = key_snippet_repository.create( filepath=snippet_info["filepath"], line_number=snippet_info["line_number"], snippet=snippet_info["snippet"], description=snippet_info["description"], + human_input_id=human_input_id, ) # Get the snippet ID from the database record