flowcept 0.9.17__py3-none-any.whl → 0.9.18__py3-none-any.whl

flowcept/agents/tools/general_tools.py

@@ -1,7 +1,7 @@
 import json
 from typing import List
 
-from flowcept.agents.agents_utils import build_llm_model, ToolResult
+from flowcept.agents.agents_utils import build_llm_model, ToolResult, normalize_message
 from flowcept.agents.flowcept_ctx_manager import mcp_flowcept
 from flowcept.agents.prompts.general_prompts import ROUTING_PROMPT, SMALL_TALK_PROMPT
 
@@ -105,6 +105,19 @@ def reset_records() -> ToolResult:
         return ToolResult(code=499, result=str(e))
 
 
+@mcp_flowcept.tool()
+def reset_context() -> ToolResult:
+    """
+    Resets all context.
+    """
+    try:
+        ctx = mcp_flowcept.get_context()
+        ctx.request_context.lifespan_context.reset_context()
+        return ToolResult(code=201, result="Context reset.")
+    except Exception as e:
+        return ToolResult(code=499, result=str(e))
+
+
 @mcp_flowcept.tool()
 def prompt_handler(message: str) -> ToolResult:
     """
@@ -120,20 +133,24 @@ def prompt_handler(message: str) -> ToolResult:
     TextContent
         The AI response or routing feedback.
     """
-    df_key_words = ["df", "save", "result = df", "reset context"]
+    df_key_words = ["df", "save", "result = df"]
     for key in df_key_words:
         if key in message:
             return run_df_query(llm=None, query=message, plot=False)
 
+    if "reset context" in message:
+        return reset_context()
     if "@record" in message:
         return record_guidance(message)
     if "@show records" in message:
         return show_records()
     if "@reset records" in message:
-        return reset_records(message)
+        return reset_records()
 
     llm = build_llm_model()
 
+    message = normalize_message(message)
+
     prompt = ROUTING_PROMPT + message
     route = llm.invoke(prompt)
 

flowcept/agents/tools/in_memory_queries/in_memory_queries_tools.py

@@ -63,11 +63,6 @@ def run_df_query(llm, query: str, plot=False) -> ToolResult:
 
     Examples
     --------
-    Reset the context:
-
-    >>> run_df_query(llm, "reset context")
-    ToolResult(code=201, result="Context Reset!")
-
     Save the current DataFrame:
 
     >>> run_df_query(llm, "save")
@@ -90,10 +85,6 @@ def run_df_query(llm, query: str, plot=False) -> ToolResult:
     custom_user_guidance = ctx.request_context.lifespan_context.custom_guidance
     if df is None or not len(df):
         return ToolResult(code=404, result="Current df is empty or null.")
-
-    if "reset context" in query:
-        ctx.request_context.lifespan_context.df = pd.DataFrame()
-        return ToolResult(code=201, result="Context Reset!")
     elif "save" in query:
         return save_df(df, schema, value_examples)
     elif "result = df" in query:
@@ -173,7 +164,7 @@ def generate_plot_code(llm, query, dynamic_schema, value_examples, df, custom_us
     >>> print(result.result["plot_code"])
     plt.bar(result_df["region"], result_df["total_sales"])
     """
-    plot_prompt = generate_plot_code_prompt(query, dynamic_schema, value_examples)
+    plot_prompt = generate_plot_code_prompt(query, dynamic_schema, value_examples, list(df.columns))
     try:
         response = llm(plot_prompt)
     except Exception as e:
@@ -300,7 +291,9 @@ def generate_result_df(
     if llm is None:
         llm = build_llm_model()
     try:
-        prompt = generate_pandas_code_prompt(query, dynamic_schema, example_values, custom_user_guidance)
+        prompt = generate_pandas_code_prompt(
+            query, dynamic_schema, example_values, custom_user_guidance, list(df.columns)
+        )
         response = llm(prompt)
     except Exception as e:
         return ToolResult(code=400, result=str(e), extra=prompt)
@@ -317,9 +310,10 @@ def generate_result_df(
                 extra={"generated_code": result_code, "exception": str(e), "prompt": prompt},
             )
         else:
-            tool_result = extract_or_fix_python_code(llm, result_code)
+            tool_result = extract_or_fix_python_code(llm, result_code, list(df.columns))
             if tool_result.code == 201:
                 new_result_code = tool_result.result
+                result_code = new_result_code
                 try:
                     result_df = safe_execute(df, new_result_code)
                 except Exception as e:
@@ -357,12 +351,7 @@ def generate_result_df(
     if summarize:
         try:
             tool_result = summarize_result(
-                llm,
-                result_code,
-                result_df,
-                query,
-                dynamic_schema,
-                example_values,
+                llm, result_code, result_df, query, dynamic_schema, example_values, list(df.columns)
             )
             if tool_result.is_success():
                 return_code = 301
@@ -377,7 +366,7 @@ def generate_result_df(
             return_code = 303
 
     try:
-        result_df = format_result_df(result_df)
+        result_df_str = format_result_df(result_df)
     except Exception as e:
         return ToolResult(
             code=405,
@@ -387,7 +376,8 @@ def generate_result_df(
 
     this_result = {
         "result_code": result_code,
-        "result_df": result_df,
+        "result_df": result_df_str,
+        "result_df_markdown": result_df.to_markdown(index=False),
         "summary": summary,
         "summary_error": summary_error,
     }
@@ -473,7 +463,7 @@ def run_df_code(user_code: str, df):
 
 
 @mcp_flowcept.tool()
-def extract_or_fix_python_code(llm, raw_text):
+def extract_or_fix_python_code(llm, raw_text, current_fields):
     """
     Extract or repair JSON code from raw text using an LLM.
 
@@ -523,7 +513,7 @@ def extract_or_fix_python_code(llm, raw_text):
     >>> print(res)
     ToolResult(code=499, result='LLM service unavailable')
     """
-    prompt = extract_or_fix_python_code_prompt(raw_text)
+    prompt = extract_or_fix_python_code_prompt(raw_text, current_fields)
     try:
         response = llm(prompt)
         return ToolResult(code=201, result=response)
@@ -582,14 +572,7 @@ def extract_or_fix_json_code(llm, raw_text) -> ToolResult:
 
 
 @mcp_flowcept.tool()
-def summarize_result(
-    llm,
-    code,
-    result,
-    query: str,
-    dynamic_schema,
-    example_values,
-) -> ToolResult:
+def summarize_result(llm, code, result, query: str, dynamic_schema, example_values, current_fields) -> ToolResult:
     """
     Summarize the pandas result with local reduction for large DataFrames.
     - For wide DataFrames, selects top columns based on variance and uniqueness.
@@ -597,7 +580,7 @@ def summarize_result(
     - Constructs a detailed prompt for the LLM with original column context.
     """
     summarized_df = summarize_df(result, code)
-    prompt = dataframe_summarizer_context(code, summarized_df, dynamic_schema, example_values, query)
+    prompt = dataframe_summarizer_context(code, summarized_df, dynamic_schema, example_values, query, current_fields)
     try:
         response = llm(prompt)
         return ToolResult(code=201, result=response)

flowcept/commons/daos/docdb_dao/lmdb_dao.py

@@ -115,6 +115,54 @@ class LMDBDAO(DocumentDBDAO):
             self.logger.exception(e)
             return False
 
+    def delete_task_keys(self, key_name, keys_list: List[str]) -> bool:
+        """Delete task documents by a key value list.
+
+        When deleting by task_id, deletes keys directly. Otherwise, scans
+        tasks and deletes matching entries.
+        """
+        if self._is_closed:
+            self._open()
+        if type(keys_list) is not list:
+            keys_list = [keys_list]
+        try:
+            with self._env.begin(write=True, db=self._tasks_db) as txn:
+                if key_name == "task_id":
+                    for key in keys_list:
+                        if key is None:
+                            continue
+                        txn.delete(str(key).encode())
+                else:
+                    cursor = txn.cursor()
+                    for key, value in cursor:
+                        entry = json.loads(value.decode())
+                        if entry.get(key_name) in keys_list:
+                            cursor.delete()
+            return True
+        except Exception as e:
+            self.logger.exception(e)
+            return False
+
+    def count_tasks(self) -> int:
+        """Count number of docs in tasks collection."""
+        if self._is_closed:
+            self._open()
+        try:
+            return self._env.stat(db=self._tasks_db).get("entries", 0)
+        except Exception as e:
+            self.logger.exception(e)
+            return -1
+
+    def count_workflows(self) -> int:
+        """Count number of docs in workflows collection."""
+        if self._is_closed:
+            self._open()
+        try:
+            return self._env.stat(db=self._workflows_db).get("entries", 0)
+        except Exception as e:
+            self.logger.exception(e)
+            return -1
+
     @staticmethod
     def _match_filter(entry, filter):
         """

flowcept/commons/daos/mq_dao/mq_dao_kafka.py

@@ -42,7 +42,7 @@ class MQDaoKafka(MQDao):
     def message_listener(self, message_handler: Callable):
         """Get message listener."""
         try:
-            while True:
+            while self._consumer is not None:
                 msg = self._consumer.poll(1.0)
                 if msg is None:
                     continue
@@ -59,7 +59,7 @@ class MQDaoKafka(MQDao):
         except Exception as e:
             self.logger.exception(e)
         finally:
-            self._consumer.close()
+            self.unsubscribe()
 
     def send_message(self, message: dict, channel=MQ_CHANNEL, serializer=msgpack.dumps):
         """Send the message."""

flowcept/commons/daos/mq_dao/mq_dao_redis.py

@@ -1,5 +1,6 @@
 """MQ redis module."""
 
+from threading import Thread
 from typing import Callable
 import redis
 
@@ -14,12 +15,15 @@ from flowcept.configs import MQ_CHANNEL, MQ_HOST, MQ_PORT, MQ_PASSWORD, MQ_URI,
 class MQDaoRedis(MQDao):
     """MQ redis class."""
 
-    MESSAGE_TYPES_IGNORE = {"psubscribe"}
+    MESSAGE_TYPES_IGNORE = {"psubscribe", "subscribe", "pong"}
 
     def __init__(self, adapter_settings=None):
         super().__init__(adapter_settings)
 
         self._consumer = None
+        self._ping_thread = None
+        self._ping_stop = False
+
         use_same_as_kv = MQ_SETTINGS.get("same_as_kvdb", False)
         if use_same_as_kv:
             if KVDB_ENABLED:
@@ -37,6 +41,26 @@ class MQDaoRedis(MQDao):
         """
         self._consumer = self._producer.pubsub()
         self._consumer.psubscribe(MQ_CHANNEL)
+        self._start_ping_thread()
+
+    def _start_ping_thread(self, interval: int = 30):
+        """Start a background thread to ping Redis pubsub periodically."""
+        if self._ping_thread and self._ping_thread.is_alive():
+            return
+
+        self._ping_stop = False
+
+        def _pinger():
+            while not self._ping_stop:
+                try:
+                    if self._consumer is not None:
+                        self._consumer.ping()
+                except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError) as e:
+                    self.logger.critical(f"Redis PubSub ping failed: {e}")
+                sleep(interval)
+
+        self._ping_thread = Thread(target=_pinger, daemon=True)
+        self._ping_thread.start()
 
     def unsubscribe(self):
         """
@@ -75,8 +99,15 @@ class MQDaoRedis(MQDao):
                     current_trials = 0
             except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError) as e:
                 current_trials += 1
-                self.logger.critical(f"Redis connection lost: {e}. Reconnecting in 3 seconds...")
+                self.logger.critical(f"Redis connection lost: {e}. Trying to reconnect in 3 seconds...")
                 sleep(3)
+                try:
+                    self.subscribe()
+                    self.logger.warning(f"Redis reconnected after {current_trials} trials.")
+                    current_trials = 0
+                except Exception as e:
+                    self.logger.critical(f"Redis error when trying to reconnect: {e}.")
+
             except Exception as e:
                 self.logger.exception(e)
                 continue

flowcept/commons/flowcept_dataclasses/task_object.py

@@ -134,7 +134,10 @@ class TaskObject:
     """User-defined tags attached to the task."""
 
     agent_id: str = None
-    """Identifier of the agent responsible for executing this task (if any)."""
+    """Identifier of the agent that executed (or is going to execute) this task."""
+
+    source_agent_id: str = None
+    """Identifier of the agent that sent this task to be executed (if any)."""
 
     _DEFAULT_ENRICH_VALUES = {
         "node_name": NODE_NAME,

flowcept/configs.py

@@ -155,14 +155,16 @@ DB_INSERTER_SLEEP_TRIALS_STOP = db_buffer_settings.get("stop_trials_sleep", 0.01
 ###########################
 
 DB_FLUSH_MODE = settings["project"].get("db_flush_mode", "offline")
-# DEBUG_MODE = settings["project"].get("debug", False)
 PERF_LOG = settings["project"].get("performance_logging", False)
 JSON_SERIALIZER = settings["project"].get("json_serializer", "default")
 REPLACE_NON_JSON_SERIALIZABLE = settings["project"].get("replace_non_json_serializable", True)
 ENRICH_MESSAGES = settings["project"].get("enrich_messages", True)
 
+# Default: enable dump buffer only when running in offline flush mode.
 _DEFAULT_DUMP_BUFFER_ENABLED = DB_FLUSH_MODE == "offline"
 DUMP_BUFFER_ENABLED = (
+    # Env var "DUMP_BUFFER" overrides settings.yaml.
+    # Falls back to settings project.dump_buffer.enabled, then to the default above.
     os.getenv(
         "DUMP_BUFFER", str(settings["project"].get("dump_buffer", {}).get("enabled", _DEFAULT_DUMP_BUFFER_ENABLED))
     )
@@ -170,6 +172,7 @@ DUMP_BUFFER_ENABLED = (
     .lower()
     in _TRUE_VALUES
 )
+# Path is only read from settings.yaml; env override is not supported here.
 DUMP_BUFFER_PATH = settings["project"].get("dump_buffer", {}).get("path", "flowcept_buffer.jsonl")
 
 TELEMETRY_CAPTURE = settings.get("telemetry_capture", None)

flowcept/flowcept_api/flowcept_controller.py

@@ -320,7 +320,11 @@ class Flowcept(object):
             file_path = DUMP_BUFFER_PATH
         assert file_path is not None, "Please indicate file_path either in the argument or in the config file."
         if not os.path.exists(file_path):
-            raise FileNotFoundError(f"Flowcept buffer file '{file_path}' was not found.")
+            raise FileNotFoundError(
+                f"Flowcept buffer file '{file_path}' was not found. "
+                f"Check your settings to see if you're dumping the data to a file and check if you"
+                f"have started Flowcept."
+            )
 
         with open(file_path, "rb") as f:
             lines = [ln for ln in f.read().splitlines() if ln]

flowcept/flowceptor/adapters/mlflow/interception_event_handler.py

@@ -13,7 +13,38 @@ class InterceptionEventHandler(FileSystemEventHandler):
         self.callback_function = callback_function
         self.interceptor_instance = interceptor_instance
 
+    def _matches_watch_target(self, path):
+        """Return True when a path matches the watch target.
+
+        If the target is a directory, any event under that directory matches.
+        If the target is a file, only the exact file path matches.
+        """
+        if not path:
+            return False
+        target = Path(self.file_path_to_watch).resolve()
+        candidate = Path(path).resolve()
+        if target.is_dir():
+            try:
+                candidate.relative_to(target)
+                return True
+            except ValueError:
+                return False
+        return candidate == target
+
+    def _maybe_callback(self, event):
+        """Invoke the callback when an event matches the watch target."""
+        paths = [getattr(event, "src_path", None), getattr(event, "dest_path", None)]
+        if any(self._matches_watch_target(path) for path in paths):
+            self.callback_function(self.interceptor_instance)
+
     def on_modified(self, event):
         """Get on modified."""
-        if Path(event.src_path).resolve() == Path(self.file_path_to_watch).resolve():
-            self.callback_function(self.interceptor_instance)
+        self._maybe_callback(event)
+
+    def on_created(self, event):
+        """Get on created."""
+        self._maybe_callback(event)
+
+    def on_moved(self, event):
+        """Get on moved."""
+        self._maybe_callback(event)

flowcept/flowceptor/adapters/mlflow/mlflow_interceptor.py

@@ -50,9 +50,10 @@ class MLFlowInterceptor(BaseInterceptor):
         interesting change, it calls self.intercept; otherwise, let it
         go....
         """
+        intercepted = 0
         runs = self.dao.get_finished_run_uuids()
         if not runs:
-            return
+            return intercepted
         for run_uuid_tuple in runs:
             run_uuid = run_uuid_tuple[0]
             if not self.state_manager.has_element_id(run_uuid):
@@ -63,6 +64,8 @@ class MLFlowInterceptor(BaseInterceptor):
                     continue
                 task_msg = self.prepare_task_msg(run_data).to_dict()
                 self.intercept(task_msg)
+                intercepted += 1
+        return intercepted
 
     def start(self, bundle_exec_id, check_safe_stops) -> "MLFlowInterceptor":
         """Start it."""
@@ -74,10 +77,20 @@ class MLFlowInterceptor(BaseInterceptor):
     def stop(self, check_safe_stops: bool = True) -> bool:
         """Stop it."""
         sleep(1)
-        super().stop(check_safe_stops)
         self.logger.debug("Interceptor stopping...")
-        self._observer.stop()
-        self._observer_thread.join()
+        # Flush any late writes before stopping the observer.
+        try:
+            intercepted = self.callback()
+            if intercepted == 0:
+                sleep(self.settings.watch_interval_sec)
+                self.callback()
+        except Exception as e:
+            self.logger.exception(e)
+        super().stop(check_safe_stops)
+        if self._observer is not None:
+            self._observer.stop()
+            if self._observer_thread is not None:
+                self._observer_thread.join()
         self.logger.debug("Interceptor stopped.")
         return True
 
@@ -98,4 +111,5 @@ class MLFlowInterceptor(BaseInterceptor):
         watch_dir = os.path.dirname(self.settings.file_path) or "."
         self._observer.schedule(event_handler, watch_dir, recursive=True)
         self._observer.start()
+        sleep(0.2)
         self.logger.info(f"Watching directory {watch_dir} with file {self.settings.file_path} ")

flowcept/flowceptor/adapters/tensorboard/tensorboard_interceptor.py

@@ -122,4 +122,5 @@ class TensorboardInterceptor(BaseInterceptor):
 
         self._observer.schedule(event_handler, self.settings.file_path, recursive=True)
         self._observer.start()
+        sleep(0.2)
         self.logger.debug(f"Watching {self.settings.file_path}")

flowcept/flowceptor/consumers/agent/base_agent_context_manager.py

@@ -20,6 +20,12 @@ class BaseAppContext:
 
     tasks: List[Dict]
 
+    def reset_context(self):
+        """
+        Method to reset the variables in the context.
+        """
+        self.tasks = []
+
 
 class BaseAgentContextManager(BaseConsumer):
     """
@@ -45,8 +51,7 @@ class BaseAgentContextManager(BaseConsumer):
         """
         self._started = False
         super().__init__()
-        self.context = None
-        self.reset_context()
+        # self.context = BaseAppContext(tasks=[])
         self.agent_id = BaseAgentContextManager.agent_id
 
     def message_handler(self, msg_obj: Dict) -> bool:
@@ -77,12 +82,6 @@ class BaseAgentContextManager(BaseConsumer):
 
         return True
 
-    def reset_context(self):
-        """
-        Resets the internal context, clearing all stored task data.
-        """
-        self.context = BaseAppContext(tasks=[])
-
     @asynccontextmanager
     async def lifespan(self, app):
         """