use helper.log also for pentesttree

src/helper/log.py

@@ -107,4 +107,7 @@ def print_message(self, message):
         elif isinstance(message, ToolMessage):
             self.console.print(Panel(message.content, title="Answer from Tool"))
         else:
-            self.console.log(Pretty(message))
+            self.console.log(Pretty(message))
+
+    def print_box(self, title, text):
+        self.console.print(Panel(text, title=title))

src/pentest_task_tree.py

@@ -2,23 +2,18 @@
 import operator
 
 from dotenv import load_dotenv
-from rich.console import Console
-from rich.panel import Panel
-
-from langchain_core.prompts import PromptTemplate
-from langchain_openai import ChatOpenAI
-
-from helper.common import get_or_fail
-from helper.ui import print_event
-from tools.ssh import get_ssh_connection_from_env, SshTestCredentialsTool, SshExecuteTool
-from graphs.initial_version import create_chat_tool_agent_graph
-
 from typing import Annotated, List, Tuple, Union
 from typing_extensions import TypedDict
 from pydantic import BaseModel, Field
 
-from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
 from langgraph.graph import StateGraph, START, END
+from langchain_openai import ChatOpenAI
+
+from graphs.initial_version import create_chat_tool_agent_graph
+from helper.common import get_or_fail
+from helper.log import RichLogger
+from tools.ssh import get_ssh_connection_from_env, SshTestCredentialsTool, SshExecuteTool
 
 # setup configuration from environment variables
 load_dotenv()
@@ -27,7 +22,7 @@
 conn.connect()
 
 # prepare console for debug output
-console = Console()
+logger = RichLogger()
 
 # the shared graph data structure
 class PlanExecute(TypedDict):
@@ -140,9 +135,10 @@ def execute_step(state: PlanExecute):
         stream_mode='values'
     )
 
+
     agent_response = None
     for event in events:
-        print_event(console, event)
+        logger.capture_event(event)
         agent_response = event
 
     return {
@@ -165,15 +161,15 @@ def plan_step(state: PlanExecute):
 
     def replan_step(state: PlanExecute):
 
-        console.print(Panel(state["plan"], title="Plan"))
+        logger.print_box('Plan', state["plan"])
 
         replanner = replanner_prompt | llm.with_structured_output(Act)
         output = replanner.invoke(state)
         if isinstance(output.action, Response):
             return {"response": output.action.response}
         else:
-            console.print(Panel(output.action.next_step, title="Next Step"))
-            console.print(Panel(output.action.steps, title="New Plan"))
+            logger.print_box('Next Step', output.action.next_step)
+            logger.print_box('New Plan', output.action.steps)
             return {"plan": output.action.steps, "next_step": output.action.next_step}
 
     workflow = StateGraph(PlanExecute)
@@ -215,4 +211,4 @@ def replan_step(state: PlanExecute):
 
 # output all occurring events 
 for event in events:
-    print_event(console, event)
+    logger.capture_event(event)