improving researcher

Author: Grigorij-Dudnik

manager.py

@@ -19,7 +19,7 @@
 from src.tools.tools_project_manager import add_task, modify_task, finish_project_planning, reorder_tasks
 from src.tools.tools_coder_pipeline import prepare_list_dir_tool, prepare_see_file_tool, ask_human_tool, retrieve_files_by_semantic_query
 from src.tools.rag.index_file_descriptions import prompt_index_project_files
-from src.utilities.util_functions import save_state_history_to_disk
+from src.utilities.util_functions import save_state_history_to_disk, join_paths
 from src.utilities.manager_utils import (
     actualize_tasks_list_and_progress_description,
     setup_todoist_project_if_needed,
@@ -32,7 +32,6 @@
     no_tools_msg,
 )
 from src.utilities.start_project_functions import set_up_dot_clean_coder_dir
-from src.utilities.util_functions import join_paths
 from src.utilities.llms import init_llms_medium_intelligence
 from src.utilities.print_formatters import print_formatted
 from src.tools.rag.retrieval import vdb_available

single_task_coder.py

@@ -26,15 +26,11 @@
 use_frontend_feedback = bool(os.getenv("FRONTEND_URL"))
 
 
-def run_clean_coder_pipeline(task: str, work_dir: str, doc_harvest: bool = False):
-    researcher = Researcher()
+def run_clean_coder_pipeline(task: str, work_dir: str, task_id: str=None):
+    researcher = Researcher(task_id=task_id)
     files, image_paths = researcher.research_task(task)
-    # documentation = None
-    # if doc_harvest:
-    #     harvester = Doc_harvester()
-    #     documentation = harvester.find_documentation(task, work_dir)
 
-    plan = planning(task, files, image_paths, work_dir) #, documentation=documentation)
+    plan = planning(task, files, image_paths, work_dir)
 
     executor = Executor(files, work_dir)
 

src/agents/researcher_agent.py

@@ -11,7 +11,14 @@
     retrieve_files_by_semantic_query,
 )
 from src.tools.rag.retrieval import vdb_available
-from src.utilities.util_functions import list_directory_tree, read_coderrules, load_prompt
+from src.utilities.util_functions import (
+    list_directory_tree,
+    read_coderrules,
+    load_prompt,
+    save_state_history_to_disk,
+    load_state_history_from_disk,
+    join_paths,
+)
 from src.utilities.langgraph_common_functions import (
     call_model,
     call_tool,
@@ -21,7 +28,7 @@
 )
 from src.utilities.print_formatters import print_formatted
 from src.utilities.llms import init_llms_medium_intelligence
-from src.utilities.util_functions import save_state_history_to_disk
+
 import os
 
 
@@ -63,6 +70,12 @@ def __init__(self, silent=False, task_id=None):
         if vdb_available():
             self.tools.append(retrieve_files_by_semantic_query)
         self.llms = init_llms_medium_intelligence(self.tools, "Researcher")
+        # Try to load previous research session for this task (if any)
+        self.prev_messages: List[BaseMessage] = []
+        if task_id:
+            history_file = join_paths(work_dir, ".clean_coder", f"research_history_task_{task_id}.json")
+            if os.path.exists(history_file):
+                self.prev_messages = load_state_history_from_disk(history_file)
 
         # workflow definition
         researcher_workflow = StateGraph(AgentState)
@@ -117,10 +130,16 @@ def research_task(self, task):
             print_formatted("👋 Hey! I'm looking for files on which we will work on together!", color="light_blue")
 
         system_prompt_template = load_prompt("researcher_system")
-        system_message = system_prompt_template.format(task=task, project_rules=read_coderrules())
-        inputs = {
-            "messages": [SystemMessage(content=system_message), HumanMessage(content=list_directory_tree(work_dir))]
-        }
+        system_message = SystemMessage(
+            content=system_prompt_template.format(task=task, project_rules=read_coderrules())
+        )
+
+        # Continue previous dialogue if available; otherwise start fresh
+        if self.prev_messages:
+            messages = [system_message] + self.prev_messages
+        else:
+            messages = [system_message, HumanMessage(content=list_directory_tree(work_dir))]
+        inputs = {"messages": messages}
         researcher_response = self.researcher.invoke(inputs, {"recursion_limit": 100})["messages"][-3]
         response_args = researcher_response.tool_calls[0]["args"]
         text_files = set(CodeFile(f) for f in response_args["files_to_work_on"] + response_args["reference_files"])

src/tools/tools_project_manager.py

@@ -135,13 +135,14 @@ def finish_project_planning(dummy: Annotated[str, "Type 'ok' to proceed."]):
     # Start background research of second task if available and not researched yet
     if len(tasks) >= 2:
         second_task = tasks[1]
-        if "<researched_files>" not in second_task.description:
+        history_file = join_paths(work_dir, ".clean_coder", f"research_history_task_{second_task.id}.json")
+        if not os.path.exists(history_file):
             background_executor.submit(research_second_task, second_task)
 
     # Execute the main pipeline to implement the task
     print_formatted("Asked programmer to execute task:", color="light_blue")
     print_text_snippet(first_task.description, title=first_task.content)
-    run_clean_coder_pipeline(task_name_description, work_dir)
+    run_clean_coder_pipeline(task_name_description, work_dir, task_id=first_task.id)
 
     actualize_progress_description_file(task_name_description)
 

src/utilities/util_functions.py

@@ -6,7 +6,7 @@
 from dotenv import load_dotenv, find_dotenv
 from todoist_api_python.api import TodoistAPI
 from langchain_core.messages import HumanMessage, ToolMessage
-from langchain_core.load import dumps
+from langchain_core.load import dumps, loads
 import json
 import click
 
@@ -257,5 +257,17 @@ def save_state_history_to_disk(state, messages_path):
         json.dump(messages_string, f)
 
 
+def load_state_history_from_disk(messages_path):
+    """Read messages saved by `save_state_history_to_disk` and recreate them.
+
+    Returns an empty list when file does not exist.
+    """
+    if not os.path.exists(messages_path):
+        return []
+    with open(messages_path, "r") as f:
+        messages_string = json.load(f)
+    return loads(messages_string)
+
+
 if __name__ == "__main__":
     print(list_directory_tree(Work.dir()))