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

flowcept/agents/agent_client.py

@@ -1,4 +1,6 @@
 import asyncio
+import json
+import re
 from typing import Dict, List, Callable
 
 from flowcept.configs import AGENT_HOST, AGENT_PORT
@@ -48,10 +50,14 @@ def run_tool(
                 result: List[TextContent] = await session.call_tool(tool_name, arguments=kwargs)
                 actual_result = []
                 for r in result.content:
-                    if isinstance(r, str):
-                        actual_result.append(r)
-                    else:
-                        actual_result.append(r.text)
+                    text = r if isinstance(r, str) else r.text
+                    try:
+                        json.loads(text)
+                        actual_result.append(text)
+                    except Exception:
+                        match = re.search(r"Error code:\\s*(\\d+)", text)
+                        code = int(match.group(1)) if match else 400
+                        actual_result.append(json.dumps({"code": code, "result": text, "tool_name": tool_name}))
 
                 return actual_result
 

flowcept/agents/agents_utils.py

@@ -139,8 +139,8 @@ def build_llm_model(
     if _service_provider == "sambanova":
         from langchain_community.llms.sambanova import SambaStudio
 
-        os.environ["SAMBASTUDIO_URL"] = AGENT.get("llm_server_url")
-        os.environ["SAMBASTUDIO_API_KEY"] = AGENT.get("api_key")
+        os.environ["SAMBASTUDIO_URL"] = os.environ.get("SAMBASTUDIO_URL", AGENT.get("llm_server_url"))
+        os.environ["SAMBASTUDIO_API_KEY"] = os.environ.get("SAMBASTUDIO_API_KEY", AGENT.get("api_key"))
 
         llm = SambaStudio(model_kwargs=_model_kwargs)
     elif _service_provider == "azure":
@@ -155,7 +155,16 @@ def build_llm_model(
         from langchain_openai import ChatOpenAI
 
         api_key = os.environ.get("OPENAI_API_KEY", AGENT.get("api_key", None))
-        llm = ChatOpenAI(openai_api_key=api_key, **model_kwargs)
+        base_url = os.environ.get("OPENAI_BASE_URL", AGENT.get("llm_server_url") or None)
+        org = os.environ.get("OPENAI_ORG_ID", AGENT.get("organization", None))
+
+        init_kwargs = {"api_key": api_key}
+        if base_url:
+            init_kwargs["base_url"] = base_url
+        if org:
+            init_kwargs["organization"] = org
+
+        llm = ChatOpenAI(**init_kwargs, **_model_kwargs)
     elif _service_provider == "google":
         if "claude" in _model_kwargs["model"]:
             api_key = os.environ.get("GOOGLE_API_KEY", AGENT.get("api_key", None))
@@ -168,22 +177,6 @@ def build_llm_model(
             from flowcept.agents.llms.gemini25 import Gemini25LLM
 
             llm = Gemini25LLM(**_model_kwargs)
-    elif _service_provider == "openai":
-        from langchain_openai import ChatOpenAI
-
-        api_key = os.environ.get("OPENAI_API_KEY", AGENT.get("api_key"))
-        base_url = os.environ.get("OPENAI_BASE_URL", AGENT.get("llm_server_url") or None)  # optional
-        org = os.environ.get("OPENAI_ORG_ID", AGENT.get("organization", None))  # optional
-
-        init_kwargs = {"api_key": api_key}
-        if base_url:
-            init_kwargs["base_url"] = base_url
-        if org:
-            init_kwargs["organization"] = org
-
-        # IMPORTANT: use the merged kwargs so `model` and temps flow through
-        llm = ChatOpenAI(**init_kwargs, **_model_kwargs)
-
     else:
         raise Exception("Currently supported providers are sambanova, openai, azure, and google.")
     if track_tools:

flowcept/agents/flowcept_agent.py

@@ -1,32 +1,133 @@
+import json
+import os
 from threading import Thread
-from time import sleep
 
 from flowcept.agents import check_liveness
+from flowcept.agents.agents_utils import ToolResult
+from flowcept.agents.tools.general_tools import prompt_handler
 from flowcept.agents.agent_client import run_tool
-from flowcept.agents.flowcept_ctx_manager import mcp_flowcept
-from flowcept.configs import AGENT_HOST, AGENT_PORT
-from flowcept.flowcept_api.flowcept_controller import Flowcept
+from flowcept.agents.flowcept_ctx_manager import mcp_flowcept, ctx_manager
+from flowcept.commons.flowcept_logger import FlowceptLogger
+from flowcept.configs import AGENT_HOST, AGENT_PORT, DUMP_BUFFER_PATH, MQ_ENABLED
+from flowcept.flowceptor.consumers.agent.base_agent_context_manager import BaseAgentContextManager
+from uuid import uuid4
 
 import uvicorn
 
 
-def main():
+class FlowceptAgent:
     """
-    Start the MCP server.
+    Flowcept agent server wrapper with optional offline buffer loading.
     """
-    f = Flowcept(start_persistence=False, save_workflow=False, check_safe_stops=False).start()
-    f.logger.info(f"This section's workflow_id={Flowcept.current_workflow_id}")
 
-    def run():
-        uvicorn.run(mcp_flowcept.streamable_http_app, host=AGENT_HOST, port=AGENT_PORT, lifespan="on")
+    def __init__(self, buffer_path: str | None = None):
+        """
+        Initialize a FlowceptAgent.
+
+        Parameters
+        ----------
+        buffer_path : str or None
+            Optional path to a JSONL buffer file. When MQ is disabled, the agent
+            loads this file once at startup.
+        """
+        self.buffer_path = buffer_path
+        self.logger = FlowceptLogger()
+        self._server_thread: Thread | None = None
+        self._server = None
+
+    def _load_buffer_once(self) -> int:
+        """
+        Load messages from a JSONL buffer file into the agent context.
+
+        Returns
+        -------
+        int
+            Number of messages loaded.
+        """
+        path = self.buffer_path or DUMP_BUFFER_PATH
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"Buffer file not found: {path}")
+
+        count = 0
+        self.logger.info(f"Loading agent buffer from {path}")
+        if ctx_manager.agent_id is None:
+            agent_id = str(uuid4())
+            BaseAgentContextManager.agent_id = agent_id
+            ctx_manager.agent_id = agent_id
+        with open(path, "r") as handle:
+            for line in handle:
+                line = line.strip()
+                if not line:
+                    continue
+                msg_obj = json.loads(line)
+                ctx_manager.message_handler(msg_obj)
+                count += 1
+        self.logger.info(f"Loaded {count} messages from buffer.")
+        return count
+
+    def _run_server(self):
+        """Run the MCP server (blocking call)."""
+        config = uvicorn.Config(mcp_flowcept.streamable_http_app, host=AGENT_HOST, port=AGENT_PORT, lifespan="on")
+        self._server = uvicorn.Server(config)
+        self._server.run()
+
+    def start(self):
+        """
+        Start the agent server in a background thread.
+
+        Returns
+        -------
+        FlowceptAgent
+            The current instance.
+        """
+        if not MQ_ENABLED:
+            self._load_buffer_once()
+
+        self._server_thread = Thread(target=self._run_server, daemon=False)
+        self._server_thread.start()
+        self.logger.info(f"Flowcept agent server started on {AGENT_HOST}:{AGENT_PORT}")
+        return self
 
-    server_thread = Thread(target=run, daemon=False)
-    server_thread.start()
-    sleep(2)
+    def stop(self):
+        """Stop the agent server and wait briefly for shutdown."""
+        if self._server is not None:
+            self._server.should_exit = True
+        if self._server_thread is not None:
+            self._server_thread.join(timeout=5)
+
+    def wait(self):
+        """Block until the server thread exits."""
+        if self._server_thread is not None:
+            self._server_thread.join()
+
+    def query(self, message: str) -> ToolResult:
+        """
+        Send a prompt to the agent's main router tool and return the response.
+        """
+        try:
+            resp = run_tool(tool_name=prompt_handler, kwargs={"message": message})[0]
+        except Exception as e:
+            return ToolResult(code=400, result=f"Error executing tool prompt_handler: {e}", tool_name="prompt_handler")
+
+        try:
+            return ToolResult(**json.loads(resp))
+        except Exception as e:
+            return ToolResult(
+                code=499,
+                result=f"Could not parse tool response as JSON: {resp}",
+                extra=str(e),
+                tool_name="prompt_handler",
+            )
+
+
+def main():
+    """
+    Start the MCP server.
+    """
+    agent = FlowceptAgent().start()
     # Wake up tool call
     print(run_tool(check_liveness, host=AGENT_HOST, port=AGENT_PORT)[0])
-
-    server_thread.join()
+    agent.wait()
 
 
 if __name__ == "__main__":

flowcept/agents/flowcept_ctx_manager.py

@@ -103,7 +103,7 @@ class FlowceptAgentContextManager(BaseAgentContextManager):
         self.schema_tracker = DynamicSchemaTracker(**self.tracker_config)
         self.msgs_counter = 0
         self.context_chunk_size = 1  # Should be in the settings
-        super().__init__()
+        super().__init__(allow_mq_disabled=True)
 
     def message_handler(self, msg_obj: Dict):
         """
@@ -133,12 +133,15 @@ class FlowceptAgentContextManager(BaseAgentContextManager):
                 if task_msg.activity_id == "reset_user_context":
                     self.context.reset_context()
                     self.msgs_counter = 0
-                    FlowceptTask(
-                        agent_id=self.agent_id,
-                        generated={"msg": "Provenance Agent reset context."},
-                        subtype="agent_task",
-                        activity_id="reset_user_context",
-                    ).send()
+                    if self._mq_dao is None:
+                        self.logger.warning("MQ is disabled; skipping reset_user_context response message.")
+                    else:
+                        FlowceptTask(
+                            agent_id=self.agent_id,
+                            generated={"msg": "Provenance Agent reset context."},
+                            subtype="agent_task",
+                            activity_id="reset_user_context",
+                        ).send()
                     return True
                 elif task_msg.activity_id == "provenance_query":
                     self.logger.info("Received a prov query message!")
@@ -161,14 +164,17 @@ class FlowceptAgentContextManager(BaseAgentContextManager):
                         status = Status.ERROR
                         error = f"Could not convert the following into a ToolResult:\n{resp}\nException: {e}"
                         generated = {"text": str(resp)}
-                    FlowceptTask(
-                        agent_id=self.agent_id,
-                        generated=generated,
-                        stderr=error,
-                        status=status,
-                        subtype="agent_task",
-                        activity_id="provenance_query_response",
-                    ).send()
+                    if self._mq_dao is None:
+                        self.logger.warning("MQ is disabled; skipping provenance_query response message.")
+                    else:
+                        FlowceptTask(
+                            agent_id=self.agent_id,
+                            generated=generated,
+                            stderr=error,
+                            status=status,
+                            subtype="agent_task",
+                            activity_id="provenance_query_response",
+                        ).send()
 
                     return True
 
@@ -200,12 +206,10 @@ class FlowceptAgentContextManager(BaseAgentContextManager):
                         ]
                     )
                 except Exception as e:
-                    self.logger.error(
-                        f"Could not add these tasks to buffer!\n"
-                        f"{
-                            self.context.task_summaries[self.msgs_counter - self.context_chunk_size : self.msgs_counter]
-                        }"
-                    )
+                    task_slice = self.context.task_summaries[
+                        self.msgs_counter - self.context_chunk_size : self.msgs_counter
+                    ]
+                    self.logger.error(f"Could not add these tasks to buffer!\n{task_slice}")
                     self.logger.exception(e)
 
                 # self.monitor_chunk()
@@ -232,9 +236,12 @@ class FlowceptAgentContextManager(BaseAgentContextManager):
         if len(result):
             content = result[0].text
             if content != "Error executing tool":
-                msg = {"type": "flowcept_agent", "info": "monitor", "content": content}
-                self._mq_dao.send_message(msg)
-                self.logger.debug(str(content))
+                if self._mq_dao is None:
+                    self.logger.warning("MQ is disabled; skipping monitor message.")
+                else:
+                    msg = {"type": "flowcept_agent", "info": "monitor", "content": content}
+                    self._mq_dao.send_message(msg)
+                    self.logger.debug(str(content))
             else:
                 self.logger.error(content)
 

flowcept/commons/daos/keyvalue_dao.py

@@ -1,7 +1,5 @@
 """Key value module."""
 
-from flowcept.commons.daos.redis_conn import RedisConn
-
 from flowcept.commons.flowcept_logger import FlowceptLogger
 from flowcept.configs import (
     KVDB_HOST,
@@ -26,12 +24,23 @@ class KeyValueDAO:
 
     def __init__(self):
         if not hasattr(self, "_initialized"):
-            self._initialized = True
             self.logger = FlowceptLogger()
+            from flowcept.commons.daos.redis_conn import RedisConn
+
             self.redis_conn = RedisConn.build_redis_conn_pool(
                 host=KVDB_HOST, port=KVDB_PORT, password=KVDB_PASSWORD, uri=KVDB_URI
             )
 
+            self._initialized = True
+
+    @staticmethod
+    def get_set_name(set_id: str, exec_bundle_id=None) -> str:
+        """Return a consistent set name for KVDB sets."""
+        set_name = set_id
+        if exec_bundle_id is not None:
+            set_name += "_" + str(exec_bundle_id)
+        return set_name
+
     def delete_set(self, set_name: str):
         """Delete it."""
         self.redis_conn.delete(set_name)

flowcept/commons/daos/mq_dao/mq_dao_base.py

@@ -7,6 +7,7 @@ import msgpack
 from time import time
 import flowcept.commons
 from flowcept.commons.autoflush_buffer import AutoflushBuffer
+from flowcept.commons.daos.keyvalue_dao import KeyValueDAO
 from flowcept.commons.utils import chunked
 from flowcept.commons.flowcept_logger import FlowceptLogger
 from flowcept.configs import (
@@ -29,6 +30,8 @@ class MQDao(object):
 
     ENCODER = GenericJSONEncoder if JSON_SERIALIZER == "complex" else None
     # TODO we don't have a unit test to cover complex dict!
+    MQ_THREAD_SET_ID = "started_mq_thread_execution"
+    MQ_FLUSH_COMPLETE_SET_ID = "pending_mq_flush_complete"
 
     @staticmethod
     def build(*args, **kwargs) -> "MQDao":
@@ -51,20 +54,6 @@ class MQDao(object):
         else:
             raise NotImplementedError
 
-    @staticmethod
-    def _get_set_name(exec_bundle_id=None):
-        """Get the set name.
-
-        :param exec_bundle_id: A way to group one or many interceptors, and
-         treat each group as a bundle to control when their time_based
-         threads started and ended.
-        :return:
-        """
-        set_id = "started_mq_thread_execution"
-        if exec_bundle_id is not None:
-            set_id += "_" + str(exec_bundle_id)
-        return set_id
-
     def __init__(self, adapter_settings=None):
         self.logger = FlowceptLogger()
         self.started = False
@@ -103,22 +92,36 @@ class MQDao(object):
 
     def register_time_based_thread_init(self, interceptor_instance_id: str, exec_bundle_id=None):
         """Register the time."""
-        set_name = MQDao._get_set_name(exec_bundle_id)
+        set_name = KeyValueDAO.get_set_name(MQDao.MQ_THREAD_SET_ID, exec_bundle_id)
         # self.logger.info(
         #     f"Register start of time_based MQ flush thread {set_name}.{interceptor_instance_id}"
         # )
         self._keyvalue_dao.add_key_into_set(set_name, interceptor_instance_id)
+        flush_set_name = KeyValueDAO.get_set_name(MQDao.MQ_FLUSH_COMPLETE_SET_ID, exec_bundle_id)
+        self._keyvalue_dao.add_key_into_set(flush_set_name, interceptor_instance_id)
 
     def register_time_based_thread_end(self, interceptor_instance_id: str, exec_bundle_id=None):
         """Register time."""
-        set_name = MQDao._get_set_name(exec_bundle_id)
+        set_name = KeyValueDAO.get_set_name(MQDao.MQ_THREAD_SET_ID, exec_bundle_id)
         self.logger.info(f"Registering end of time_based MQ flush thread {set_name}.{interceptor_instance_id}")
         self._keyvalue_dao.remove_key_from_set(set_name, interceptor_instance_id)
         self.logger.info(f"Done registering time_based MQ flush thread {set_name}.{interceptor_instance_id}")
 
     def all_time_based_threads_ended(self, exec_bundle_id=None):
         """Get all time."""
-        set_name = MQDao._get_set_name(exec_bundle_id)
+        set_name = KeyValueDAO.get_set_name(MQDao.MQ_THREAD_SET_ID, exec_bundle_id)
+        return self._keyvalue_dao.set_is_empty(set_name)
+
+    def register_flush_complete(self, interceptor_instance_id: str, exec_bundle_id=None):
+        """Register a flush-complete signal for an interceptor."""
+        set_name = KeyValueDAO.get_set_name(MQDao.MQ_FLUSH_COMPLETE_SET_ID, exec_bundle_id)
+        self.logger.info(f"Registering flush completion {set_name}.{interceptor_instance_id}")
+        self._keyvalue_dao.remove_key_from_set(set_name, interceptor_instance_id)
+        self.logger.info(f"Done registering flush completion {set_name}.{interceptor_instance_id}")
+
+    def all_flush_complete_received(self, exec_bundle_id=None):
+        """Return True when all interceptors in the bundle reported flush completion."""
+        set_name = KeyValueDAO.get_set_name(MQDao.MQ_FLUSH_COMPLETE_SET_ID, exec_bundle_id)
         return self._keyvalue_dao.set_is_empty(set_name)
 
     def set_campaign_id(self, campaign_id=None):
@@ -172,11 +175,14 @@ class MQDao(object):
             if self._time_based_flushing_started:
                 self.buffer.stop()
                 self._time_based_flushing_started = False
+                self.logger.debug("MQ time-based flushed for the last time!")
             else:
                 self.logger.error("MQ time-based flushing is not started")
         else:
             self.buffer = list()
 
+        self.logger.debug("Buffer closed.")
+
     def _stop_timed(self, interceptor_instance_id: str, check_safe_stops: bool = True, bundle_exec_id: int = None):
         t1 = time()
         self._stop(interceptor_instance_id, check_safe_stops, bundle_exec_id)
@@ -190,10 +196,12 @@ class MQDao(object):
 
     def _stop(self, interceptor_instance_id: str = None, check_safe_stops: bool = True, bundle_exec_id: int = None):
         """Stop MQ publisher."""
-        self.logger.debug(f"MQ pub received stop sign: bundle={bundle_exec_id}, interceptor={interceptor_instance_id}")
         self._close_buffer()
-        self.logger.debug("Flushed MQ for the last time!")
-        if check_safe_stops:
+        if check_safe_stops and MQ_ENABLED:
+            self.logger.debug(
+                f"Sending flush-complete msg. Bundle: {bundle_exec_id}; interceptor id: {interceptor_instance_id}"
+            )
+            self._send_mq_dao_flush_complete(interceptor_instance_id, bundle_exec_id)
             self.logger.debug(f"Sending stop msg. Bundle: {bundle_exec_id}; interceptor id: {interceptor_instance_id}")
             self._send_mq_dao_time_thread_stop(interceptor_instance_id, bundle_exec_id)
         self.started = False
@@ -210,6 +218,15 @@ class MQDao(object):
         # self.logger.info("Control msg sent: " + str(msg))
         self.send_message(msg)
 
+    def _send_mq_dao_flush_complete(self, interceptor_instance_id, exec_bundle_id=None):
+        msg = {
+            "type": "flowcept_control",
+            "info": "mq_flush_complete",
+            "interceptor_instance_id": interceptor_instance_id,
+            "exec_bundle_id": exec_bundle_id,
+        }
+        self.send_message(msg)
+
     def send_document_inserter_stop(self, exec_bundle_id=None):
         """Send the document."""
         # These control_messages are handled by the document inserter

flowcept/configs.py

@@ -9,7 +9,7 @@ from flowcept.version import __version__
 PROJECT_NAME = "flowcept"
 
 DEFAULT_SETTINGS = {
-    "version": __version__,
+    "flowcept_version": __version__,
     "log": {"log_file_level": "disable", "log_stream_level": "disable"},
     "project": {"dump_buffer": {"enabled": True}},
     "telemetry_capture": {},
@@ -81,7 +81,7 @@ FLOWCEPT_USER = settings["experiment"].get("user", "blank_user")
 
 MQ_INSTANCES = settings["mq"].get("instances", None)
 MQ_SETTINGS = settings["mq"]
-MQ_ENABLED = os.getenv("MQ_ENABLED", settings["mq"].get("enabled", True))
+MQ_ENABLED = os.getenv("MQ_ENABLED", str(settings["mq"].get("enabled", True))).strip().lower() in _TRUE_VALUES
 MQ_TYPE = os.getenv("MQ_TYPE", settings["mq"].get("type", "redis"))
 MQ_CHANNEL = os.getenv("MQ_CHANNEL", settings["mq"].get("channel", "interception"))
 MQ_PASSWORD = settings["mq"].get("password", None)
@@ -103,6 +103,11 @@ KVDB_PORT = int(os.getenv("KVDB_PORT", settings["kv_db"].get("port", "6379")))
 KVDB_URI = os.getenv("KVDB_URI", settings["kv_db"].get("uri", None))
 KVDB_ENABLED = settings["kv_db"].get("enabled", False)
 
+if MQ_ENABLED and not KVDB_ENABLED:
+    raise ValueError(
+        "Invalid configuration: MQ is enabled but kv_db is disabled. "
+        "Enable kv_db.enabled (and KVDB) when MQ is enabled."
+    )
 
 DATABASES = settings.get("databases", {})
 
@@ -160,6 +165,12 @@ 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)
 
+if DB_FLUSH_MODE == "online" and not MQ_ENABLED:
+    raise ValueError(
+        "Invalid configuration: project.db_flush_mode is 'online' but MQ is disabled. "
+        "Enable mq.enabled (or MQ_ENABLED=true) or set project.db_flush_mode to 'offline'."
+    )
+
 # Default: enable dump buffer only when running in offline flush mode.
 _DEFAULT_DUMP_BUFFER_ENABLED = DB_FLUSH_MODE == "offline"
 DUMP_BUFFER_ENABLED = (

flowcept/flowceptor/consumers/agent/base_agent_context_manager.py

@@ -45,12 +45,12 @@ class BaseAgentContextManager(BaseConsumer):
 
     agent_id = None
 
-    def __init__(self):
+    def __init__(self, allow_mq_disabled: bool = False):
         """
         Initializes the agent and resets its context state.
         """
         self._started = False
-        super().__init__()
+        super().__init__(allow_mq_disabled=allow_mq_disabled)
         # self.context = BaseAppContext(tasks=[])
         self.agent_id = BaseAgentContextManager.agent_id
 

flowcept/flowceptor/consumers/base_consumer.py

@@ -13,18 +13,28 @@ class BaseConsumer(object):
 
     This class provides a standard interface and shared logic for subscribing to
     message queues and dispatching messages to a handler.
+
+    Note
+    ----
+    The MQ-disabled path is only intended for agent consumers that can operate
+    from an offline buffer file. General consumers that require MQ should keep
+    the default behavior (raise when MQ_ENABLED is False).
     """
 
-    def __init__(self):
+    def __init__(self, allow_mq_disabled: bool = False):
         """Initialize the message queue DAO and logger."""
+        self.logger = FlowceptLogger()
+        self._main_thread: Optional[Thread] = None
+
         if not MQ_ENABLED:
+            if allow_mq_disabled:
+                self._mq_dao = None
+                self.logger.warning("MQ is disabled; starting consumer without a message queue.")
+                return
             raise Exception("MQ is disabled in the settings. You cannot consume messages.")
 
         self._mq_dao = MQDao.build()
 
-        self.logger = FlowceptLogger()
-        self._main_thread: Optional[Thread] = None
-
     @abstractmethod
     def message_handler(self, msg_obj: Dict) -> bool:
         """
@@ -62,6 +72,9 @@ class BaseConsumer(object):
         BaseConsumer
             The current instance (to allow chaining).
         """
+        if self._mq_dao is None:
+            self.logger.warning("MQ is disabled; skipping message consumption start.")
+            return self
         if target is None:
             target = self.default_thread_target
         self._mq_dao.subscribe()
@@ -85,6 +98,9 @@ class BaseConsumer(object):
         --------
         start : Starts the consumer and optionally spawns a background thread to run this method.
         """
+        if self._mq_dao is None:
+            self.logger.warning("MQ is disabled; no message listener will run.")
+            return
         self.logger.debug("Going to wait for new messages!")
         self._mq_dao.message_listener(self.message_handler)
         self.logger.debug("Broke main message listening loop!")
@@ -96,4 +112,6 @@ class BaseConsumer(object):
         """
         Stop consuming messages by unsubscribing from the message queue.
         """
+        if self._mq_dao is None:
+            return
         self._mq_dao.unsubscribe()

flowcept/flowceptor/consumers/document_inserter.py

@@ -197,6 +197,24 @@ class DocumentInserter(BaseConsumer):
                     f"{'' if exec_bundle_id is None else exec_bundle_id}_{interceptor_instance_id}!"
                 )
             return "continue"
+        elif message["info"] == "mq_flush_complete":
+            exec_bundle_id = message.get("exec_bundle_id", None)
+            interceptor_instance_id = message.get("interceptor_instance_id")
+            self.logger.info(
+                f"DocInserter id {id(self)}. Received mq_flush_complete message "
+                f"from the interceptor {'' if exec_bundle_id is None else exec_bundle_id}_{interceptor_instance_id}!"
+            )
+            if self.check_safe_stops:
+                self.logger.info(
+                    f"Begin register_flush_complete "
+                    f"{'' if exec_bundle_id is None else exec_bundle_id}_{interceptor_instance_id}!"
+                )
+                self._mq_dao.register_flush_complete(interceptor_instance_id, exec_bundle_id)
+                self.logger.info(
+                    f"Done register_flush_complete "
+                    f"{'' if exec_bundle_id is None else exec_bundle_id}_{interceptor_instance_id}!"
+                )
+            return "continue"
         elif message["info"] == "stop_document_inserter":
             exec_bundle_id = message.get("exec_bundle_id", None)
             if self._bundle_exec_id == exec_bundle_id:
@@ -297,7 +315,10 @@ class DocumentInserter(BaseConsumer):
             return self
         if self.check_safe_stops:
             trial = 0
-            while not self._mq_dao.all_time_based_threads_ended(bundle_exec_id):
+            while not (
+                self._mq_dao.all_time_based_threads_ended(bundle_exec_id)
+                and self._mq_dao.all_flush_complete_received(bundle_exec_id)
+            ):
                 self.logger.debug(
                     f"# time_based_threads for bundle_exec_id {bundle_exec_id} is"
                     f"{self._mq_dao._keyvalue_dao.set_count(bundle_exec_id)}"

flowcept/version.py

@@ -4,4 +4,4 @@
 # The expected format is: <Major>.<Minor>.<Patch>
 # This file is supposed to be automatically modified by the CI Bot.
 # See .github/workflows/version_bumper.py
-__version__ = "0.9.18"
+__version__ = "0.9.19"

{flowcept-0.9.18.dist-info → flowcept-0.9.19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flowcept
-Version: 0.9.18
+Version: 0.9.19
 Summary: Capture and query workflow provenance data using data observability
 Author: Oak Ridge National Laboratory
 License-Expression: MIT

{flowcept-0.9.18.dist-info → flowcept-0.9.19.dist-info}/RECORD

@@ -1,13 +1,13 @@
 flowcept/__init__.py,sha256=tvVZKyymdqv3qOsgpAyDppBlUiBc0ag4QF21IcS-mVk,2449
 flowcept/cli.py,sha256=d3hogRpuMwQFIqw_fsYD4y074o3AFslBjlkRvTthOVM,25702
-flowcept/configs.py,sha256=JN_YFuNFeFyYe1ieiPhSlEc8PVnU8OXoMcKgf7c46ko,9011
-flowcept/version.py,sha256=FVbVkAZsGLJF4VKbw0v8G0N_UTTbGRqm1dWVoX3mjfk,307
+flowcept/configs.py,sha256=hqHdaXbVaXF0IPrUjqPPV_ey5ol31Gx2Gmjh43eQuik,9512
+flowcept/version.py,sha256=cmAij7Yt87_OWquErsjUo2KO0u1npqmoClVqCUwAw90,307
 flowcept/agents/__init__.py,sha256=8eeD2CiKBtHiDsWdrHK_UreIkKlTq4dUbhHDyzw372o,175
-flowcept/agents/agent_client.py,sha256=UiBQkC9WE2weLZR2OTkEOEQt9-zqQOkPwRA17HfI-jk,2027
-flowcept/agents/agents_utils.py,sha256=aFJ_RVqE4XlXTG7e6SH6WPqfr7hjT4J3SjSgM6sgY60,8630
+flowcept/agents/agent_client.py,sha256=sfXZfF48pOaznRVgzPbTqkEgFnSrqn7SkgRsNE82-cg,2353
+flowcept/agents/agents_utils.py,sha256=oc_ExBIKYGsPPS-FkZEjuL6RAx75IlYE_ykbIzxHYGg,8375
 flowcept/agents/dynamic_schema_tracker.py,sha256=TsmXRRkyUkqB-0bEgmeqSms8xj1tMMJeYvjoaO2mtwI,6829
-flowcept/agents/flowcept_agent.py,sha256=-PM7sGalZnUk-NVCekfHIY0I7oEZD9mnYGXMxFpo0EM,945
-flowcept/agents/flowcept_ctx_manager.py,sha256=OiDzie1qp2ZlSphA-4b5kjUHlzP0yjJ4XcFSiNEGgiU,9717
+flowcept/agents/flowcept_agent.py,sha256=iF9ZucrKU7uz35_W1UdXwXybQPxLpr_MkwSZLU0vAbo,4330
+flowcept/agents/flowcept_ctx_manager.py,sha256=WIA2K5WJ71z-wwvyaHmHir8eIlJUVvuE6UqNrsNLZAk,10275
 flowcept/agents/gui/__init__.py,sha256=Qw9YKbAzgZqBjMQGnF7XWmfUo0fivtkDISQRK3LA3gU,113
 flowcept/agents/gui/agent_gui.py,sha256=VpwhQamzFKBfrmibxOIc-8wXtZnd2Cq7tbKahZZOp7c,2995
 flowcept/agents/gui/audio_utils.py,sha256=piA_dc36io1sYqLF6QArS4AMl-cfDa001jGhYz5LkB4,4279
@@ -36,14 +36,14 @@ flowcept/commons/task_data_preprocess.py,sha256=-ceLexv2ZfZOAYF43DPagGwQPgt_L_lN
 flowcept/commons/utils.py,sha256=okCShkcuWhzznBtADDDusTdfPXO0W041b2f4Aog-7SE,9831
 flowcept/commons/vocabulary.py,sha256=0psC4NulNFn88mjTcoT_aT4QxX8ljMFgTOF3FxzM40A,1118
 flowcept/commons/daos/__init__.py,sha256=RO51svfHOg9naN676zuQwbj_RQ6IFHu-RALeefvtwwk,23
-flowcept/commons/daos/keyvalue_dao.py,sha256=g7zgC9hVC1NTllwUAqGt44YqdqYUgAKgPlX8_G4BRGw,3599
+flowcept/commons/daos/keyvalue_dao.py,sha256=_tLnq2Se9deemYmQ9QjGYB_J16ZdAfa2pO3zzd3cQHU,3891
 flowcept/commons/daos/redis_conn.py,sha256=gFyW-5yf6B8ExEYopCmbap8ki-iEwuIw-KH9f6o7UGQ,1495
 flowcept/commons/daos/docdb_dao/__init__.py,sha256=qRvXREeUJ4mkhxdC9bzpOsVX6M2FB5hDyLFxhMxTGhs,30
 flowcept/commons/daos/docdb_dao/docdb_dao_base.py,sha256=YbfSVJPwZGK2GBYkeapRC83HkmP0c6Msv5TriD88RcI,11812
 flowcept/commons/daos/docdb_dao/lmdb_dao.py,sha256=5FV11hjIpFUZTP4HJwPTswd4TT27EJ3it82TpY5Z2RI,12187
 flowcept/commons/daos/docdb_dao/mongodb_dao.py,sha256=5x0un15uCDTcnuITOyOhvF9mKj_bUmF2du0AHQfjN9k,40055
 flowcept/commons/daos/mq_dao/__init__.py,sha256=Xxm4FmbBUZDQ7XIAmSFbeKE_AdHsbgFmSuftvMWSykQ,21
-flowcept/commons/daos/mq_dao/mq_dao_base.py,sha256=VXqXzesU01dCHE5i0urnYQppixUNGZbJMRmm4jSAcgM,9424
+flowcept/commons/daos/mq_dao/mq_dao_base.py,sha256=BAegPRL9yozNFJnlkpyc2pW06ZiLXHzGITKFhY2ID7A,10777
 flowcept/commons/daos/mq_dao/mq_dao_kafka.py,sha256=mWoY9RvViHegzXXynuegU_jg-S55YSV5lfgNqaPuMlg,5085
 flowcept/commons/daos/mq_dao/mq_dao_mofka.py,sha256=tRdMGYDzdeIJxad-B4-DE6u8Wzs61eTzOW4ojZrnTxs,4057
 flowcept/commons/daos/mq_dao/mq_dao_redis.py,sha256=be1ejbQofeJhf2qw7AWT9JF8z-a3kk-DbSCMU9wAwuI,6773
@@ -82,11 +82,11 @@ flowcept/flowceptor/adapters/tensorboard/__init__.py,sha256=LrcR4WCIlBwwHIUSteQ8
 flowcept/flowceptor/adapters/tensorboard/tensorboard_dataclasses.py,sha256=lSfDd6TucVNzGxbm69BYyCVgMr2p9iUEQjnsS4jIfeI,554
 flowcept/flowceptor/adapters/tensorboard/tensorboard_interceptor.py,sha256=GcRiY93MjuEmdhh37PAyj4ZtwBovA7FPiEM2TjVxsPw,5123
 flowcept/flowceptor/consumers/__init__.py,sha256=foxtVEb2ZEe9g1slfYIKM4tIFv-He1l7XS--SYs7nlQ,28
-flowcept/flowceptor/consumers/base_consumer.py,sha256=hrZ3VFV7pJBMXZsvh7Q2Y36b_ifcnbJkgwe2MiuZL70,3324
+flowcept/flowceptor/consumers/base_consumer.py,sha256=B_oo-FI222PsHErlTiW0IALMPqu2d3cSSRA8nMW5hDk,4113
 flowcept/flowceptor/consumers/consumer_utils.py,sha256=E6R07zIKNXJTCxvL-OCrCKNYRpqtwRiXiZx0D2BKidk,5893
-flowcept/flowceptor/consumers/document_inserter.py,sha256=IeVl6Y4Q1KlpYGvE7uDI0vKQf-MGf2pgnIpxCYtyzKE,13392
+flowcept/flowceptor/consumers/document_inserter.py,sha256=f9EnQfBbMx27UKzCIymNvsyTxuv8mP4ftY4njxvyjo0,14515
 flowcept/flowceptor/consumers/agent/__init__.py,sha256=R1uvjBPeTLw9SpYgyUc6Qmo16pE84PFHcELTTFvyTWU,56
-flowcept/flowceptor/consumers/agent/base_agent_context_manager.py,sha256=sAoMtbfB63Ys7AxUiARlTm9z1QX9HcVNYE-mpjyYo58,4116
+flowcept/flowceptor/consumers/agent/base_agent_context_manager.py,sha256=Hi2fBgyWLks3mamWNoT0dKT5vDspYIURq4eflq-JKeU,4184
 flowcept/instrumentation/__init__.py,sha256=M5bTmg80E4QyN91gUX3qfw_nbtJSXwGWcKxdZP3vJz0,34
 flowcept/instrumentation/flowcept_agent_task.py,sha256=XN9JU4LODca0SgojUm4F5iU_V8tuWkOt1fAKcoOAG34,10757
 flowcept/instrumentation/flowcept_decorator.py,sha256=X4Lp_FSsoL08K8ZhRM4mC0OjKupbQtbMQR8zxy3ezDY,1350
@@ -94,9 +94,9 @@ flowcept/instrumentation/flowcept_loop.py,sha256=nF7Sov-DCDapyYvS8zx-1ZFrnjc3CPg
 flowcept/instrumentation/flowcept_task.py,sha256=_G1e5SOMQ1y8RhO_rVexe7icQbgFF4TpaHlpW_MxERc,11135
 flowcept/instrumentation/flowcept_torch.py,sha256=kkZQRYq6cDBpdBU6J39_4oKRVkhyF3ODlz8ydV5WGKw,23455
 flowcept/instrumentation/task_capture.py,sha256=p_Cj9_cHVMhgzqDXhKqdpk01-88VAAVhbgk5IDa-7sk,8576
-resources/sample_settings.yaml,sha256=AYFqPC0im0ia5TiR2Ri_d5C11gLSDaqXX5j5fa79D6Q,6895
-flowcept-0.9.18.dist-info/METADATA,sha256=jBMMehf112ABa6qRQhSBU3q-2Aho9KvD02tx5aPEYRU,33386
-flowcept-0.9.18.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-flowcept-0.9.18.dist-info/entry_points.txt,sha256=i8q67WE0201rVxYI2lyBtS52shvgl93x2Szp4q8zMlw,47
-flowcept-0.9.18.dist-info/licenses/LICENSE,sha256=r5-2P6tFTuRGWT5TiX32s1y0tnp4cIqBEC1QjTaXe2k,1086
-flowcept-0.9.18.dist-info/RECORD,,
+resources/sample_settings.yaml,sha256=5k1e81rNmBs7MEW-FyX1DPnrZcQQUrF7imAPKr4Mtys,6895
+flowcept-0.9.19.dist-info/METADATA,sha256=bcyJrlLE1SGz8M96BA6SpVIZkjsCxrgYV6NlYOpxVvM,33386
+flowcept-0.9.19.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+flowcept-0.9.19.dist-info/entry_points.txt,sha256=i8q67WE0201rVxYI2lyBtS52shvgl93x2Szp4q8zMlw,47
+flowcept-0.9.19.dist-info/licenses/LICENSE,sha256=r5-2P6tFTuRGWT5TiX32s1y0tnp4cIqBEC1QjTaXe2k,1086
+flowcept-0.9.19.dist-info/RECORD,,

resources/sample_settings.yaml

@@ -1,4 +1,4 @@
-flowcept_version: 0.9.18 # Version of the Flowcept package. This setting file is compatible with this version.
+flowcept_version: 0.9.19 # Version of the Flowcept package. This setting file is compatible with this version.
 
 project:
   debug: true # Toggle debug mode. This will add a property `debug: true` to all saved data, making it easier to retrieve/delete them later.