Merge pull request #292 from embersax/componentmemory

[MRG]add memory to each coponent

Author: huangyz0918

mle/agents/advisor.py

@@ -4,6 +4,7 @@
 
 from mle.function import *
 from mle.utils import get_config, print_in_box, clean_json_string
+from mle.utils.component_memory import trace_component
 
 
 def process_report(requirement: str, suggestions: dict):
@@ -136,6 +137,7 @@ def __init__(self, model, console=None, mode='normal'):
         self.sys_prompt += self.json_mode_prompt
         self.chat_history.append({"role": 'system', "content": self.sys_prompt})
 
+    @trace_component("advisor")
     def suggest(self, requirement, return_raw=False):
         """
         Handle the query from the model query response.
@@ -163,6 +165,7 @@ def suggest(self, requirement, return_raw=False):
 
         return process_report(requirement, suggestions)
 
+    @trace_component("advisor")
     def interact(self, requirement):
         """
         Interact with the user to ask and suggest.

mle/agents/chat.py

@@ -2,7 +2,7 @@
 
 from mle.function import *
 from mle.utils import get_config, WorkflowCache
-
+from mle.utils.component_memory import trace_component
 
 class ChatAgent:
 
@@ -107,7 +107,7 @@ def greet(self):
 
         self.chat_history.append({"role": "assistant", "content": greets})
         return greets
-
+    @trace_component("chat")
     def chat(self, user_prompt):
         """
         Handle the response from the model streaming.

mle/agents/coder.py

@@ -4,7 +4,7 @@
 
 from mle.function import *
 from mle.utils import get_config, print_in_box, clean_json_string
-
+from mle.utils.component_memory import trace_component
 
 def process_summary(summary_dict: dict):
     """
@@ -148,7 +148,8 @@ def __init__(self, model, working_dir='.', console=None, single_file=False):
 
         self.sys_prompt += self.json_mode_prompt
         self.chat_history.append({"role": 'system', "content": self.sys_prompt})
-
+        
+    @trace_component("coder")
     def read_requirement(self, advisor_report: str):
         """
         Read the user requirement and the advisor report.
@@ -157,6 +158,7 @@ def read_requirement(self, advisor_report: str):
         """
         self.chat_history.append({"role": "system", "content": advisor_report})
 
+    @trace_component("coder")
     def code(self, task_dict: dict):
         """
         Handle the query from the model query response.
@@ -181,7 +183,8 @@ def code(self, task_dict: dict):
             code_summary = clean_json_string(text)
             code_summary.update({'task': task_dict.get('task'), 'task_description': task_dict.get('description')})
         return code_summary
-
+    
+    @trace_component("coder")
     def debug(self, task_dict: dict, debug_report: dict):
         """
         Handle the query from the model query response.

mle/agents/debugger.py

@@ -4,7 +4,7 @@
 from mle.utils import get_config, print_in_box
 
 from rich.console import Console
-
+from mle.utils.component_memory import trace_component
 
 def process_debug_report(debug_report):
     """
@@ -143,7 +143,7 @@ def analyze_with_log(self, commands, logs):
         report_dict = json.loads(text)
         print_in_box(process_debug_report(report_dict), self.console, title="MLE Debugger", color="yellow")
         return report_dict
-
+    @trace_component("debugger")
     def analyze(self, code_report):
         """
         Handle the query from the model query response.
@@ -177,4 +177,4 @@ def analyze(self, code_report):
         self.chat_history.append({"role": "assistant", "content": text})
         report_dict = json.loads(text)
         print_in_box(process_debug_report(report_dict), self.console, title="MLE Debugger", color="yellow")
-        return report_dict
+        return report_dict

mle/agents/planner.py

@@ -4,6 +4,7 @@
 from rich.console import Console
 
 from mle.utils import print_in_box, clean_json_string
+from mle.utils.component_memory import trace_component
 
 
 def process_plan(plan_dict: dict):
@@ -82,6 +83,7 @@ def __init__(self, model, console=None):
         self.sys_prompt += self.json_mode_prompt
         self.chat_history.append({"role": 'system', "content": self.sys_prompt})
 
+    @trace_component("planner")
     def plan(self, user_prompt):
         """
         Handle the query from the model query response.
@@ -102,6 +104,7 @@ def plan(self, user_prompt):
         except json.JSONDecodeError as e:
             return clean_json_string(text)
 
+    @trace_component("planner")
     def interact(self, user_prompt):
         """
         Handle the query from the model query response.

mle/agents/reporter.py

@@ -1,7 +1,7 @@
 import json
 from rich.console import Console
 from time import gmtime, strftime
-
+from mle.utils.component_memory import trace_component
 
 class ReportAgent:
 
@@ -138,7 +138,8 @@ def process_knowledge(self, github_summary: dict, calendar_events: list = None,
 
         self.knowledge = info_prompt
         return info_prompt
-
+    
+    @trace_component("reporter")
     def gen_report(self, github_summary: dict, calendar_events: list = None, okr: str = None):
         """
         Handle the query from the model query response.
@@ -165,4 +166,4 @@ def gen_report(self, github_summary: dict, calendar_events: list = None, okr: st
             result_dict = json.loads(text)
             with open(f'progress_report_{today}.json', 'w') as f:
                 json.dump(result_dict, f)
-            return result_dict
+            return result_dict

mle/cli.py

@@ -8,6 +8,7 @@
 from pathlib import Path
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn, TextColumn, BarColumn
+import json
 
 import mle
 from mle.server import app
@@ -21,6 +22,7 @@
 )
 from mle.utils import CodeChunker
 from mle.utils import LanceDBMemory, list_files, read_file
+from mle.utils.component_memory import ComponentMemory
 
 console = Console()
 
@@ -303,7 +305,7 @@ def new(name):
         base_url = questionary.text(
             "What is your vLLM server URL? (default: http://localhost:8000/v1)"
         ).ask() or "http://localhost:8000/v1"
-        
+
         model_name = questionary.text(
             "What is the model name loaded in your vLLM server? (default: mistralai/Mistral-7B-Instruct-v0.3B)"
         ).ask() or "mistralai/Mistral-7B-Instruct-v0.3"
@@ -433,3 +435,79 @@ def memory(add, rm, update):
                     table_name=table_name,
                     metadata=[{'file': file_path, 'chunk_key': k} for k, _ in chunks.items()]
                 )
+
+
+@cli.command()
+@click.option('--component', type=click.Choice([
+    'advisor', 'planner', 'coder', 'debugger', 'reporter', 'chat',
+    'github_summarizer', 'git_summarizer'
+]), help='Component to view traces for')
+@click.option('--limit', default=5, help='Maximum number of traces to show')
+@click.option('--full-output', is_flag=True, help='Show complete output (not truncated)')
+def traces(component, limit, full_output):
+    """View execution traces for components."""
+    if not component:
+        console.print("[yellow]Please specify a component to view traces for.[/yellow]")
+        return
+
+    memory = ComponentMemory(os.getcwd())
+    traces = memory.get_recent_traces(component, limit)
+
+    if not traces:
+        console.print(f"[yellow]No traces found for component: {component}[/yellow]")
+        return
+
+    console.print(f"[green]Recent {component} traces:[/green]")
+
+    for i, trace in enumerate(traces):
+        console.print(f"\n[bold cyan]Trace #{i+1}[/bold cyan] ({trace['timestamp']})")
+        console.print(f"Status: {trace['status']}")
+
+        if trace['execution_time']:
+            console.print(f"Execution Time: {trace['execution_time']:.2f} seconds")
+
+        # Show context information
+        if trace['context']:
+            context = trace['context']
+            if isinstance(context, dict) and context:
+                console.print("\n[bold]Context:[/bold]")
+                for key, value in context.items():
+                    console.print(f"  {key}: {value}")
+
+        # Show full input
+        console.print("\n[bold]Input:[/bold]")
+        if isinstance(trace['input_data'], str):
+            if full_output:
+                console.print(trace['input_data'])
+            else:
+                console.print(trace['input_data'][:500] + ("..." if len(trace['input_data']) > 500 else ""))
+        else:
+            # Handle dictionary or other structured data
+            input_str = json.dumps(trace['input_data'], indent=2) if isinstance(trace['input_data'], (dict, list)) else str(trace['input_data'])
+            if full_output:
+                console.print(input_str)
+            else:
+                console.print(input_str[:500] + ("..." if len(input_str) > 500 else ""))
+
+        # Show full output
+        console.print("\n[bold]Output:[/bold]")
+        if isinstance(trace['output_data'], str):
+            if full_output:
+                console.print(trace['output_data'])
+            else:
+                console.print(trace['output_data'][:500] + ("..." if len(trace['output_data']) > 500 else ""))
+        else:
+            # Handle dictionary or other structured data
+            output_str = json.dumps(trace['output_data'], indent=2) if isinstance(trace['output_data'], (dict, list)) else str(trace['output_data'])
+            if full_output:
+                console.print(output_str)
+            else:
+                console.print(output_str[:500] + ("..." if len(output_str) > 500 else ""))
+
+        console.print("-" * 50)
+
+    # Show command for seeing full output
+    if not full_output and traces:
+        console.print("[yellow]Tip: Use --full-output flag to see complete trace data[/yellow]")
+
+    memory.close()