flowcept 0.8.10__py3-none-any.whl → 0.8.11__py3-none-any.whl

flowcept/cli.py

@@ -14,6 +14,9 @@ Supports:
 - `flowcept --help --command` for command-specific help
 """
 
+import subprocess
+from time import sleep
+from typing import Dict, Optional
 import argparse
 import os
 import sys
@@ -21,6 +24,8 @@ import json
 import textwrap
 import inspect
 from functools import wraps
+from importlib import resources
+from pathlib import Path
 from typing import List
 
 from flowcept import Flowcept, configs
@@ -46,11 +51,31 @@ def show_config():
     }
     print(f"This is the settings path in this session: {configs.SETTINGS_PATH}")
     print(
-        f"This is your FLOWCEPT_SETTINGS_PATH environment variable value: "
-        f"{config_data['env_FLOWCEPT_SETTINGS_PATH']}"
+        f"This is your FLOWCEPT_SETTINGS_PATH environment variable value: {config_data['env_FLOWCEPT_SETTINGS_PATH']}"
     )
 
 
+def init_settings():
+    """
+    Create a new settings.yaml file in your home directory under ~/.flowcept.
+    """
+    dest_path = Path(os.path.join(configs._SETTINGS_DIR, "settings.yaml"))
+
+    if dest_path.exists():
+        overwrite = input(f"{dest_path} already exists. Overwrite? (y/N): ").strip().lower()
+        if overwrite != "y":
+            print("Operation aborted.")
+            return
+
+    os.makedirs(configs._SETTINGS_DIR, exist_ok=True)
+
+    SAMPLE_SETTINGS_PATH = str(resources.files("resources").joinpath("sample_settings.yaml"))
+
+    with open(SAMPLE_SETTINGS_PATH, "rb") as src_file, open(dest_path, "wb") as dst_file:
+        dst_file.write(src_file.read())
+    print(f"Copied {configs.SETTINGS_PATH} to {dest_path}")
+
+
 def start_consumption_services(bundle_exec_id: str = None, check_safe_stops: bool = False, consumers: List[str] = None):
     """
     Start services that consume data from a queue or other source.
@@ -121,28 +146,202 @@ def workflow_count(workflow_id: str):
     print(json.dumps(result, indent=2))
 
 
-def query(query_str: str):
+def query(filter: str, project: str = None, sort: str = None, limit: int = 0):
+    """
+    Query the MongoDB task collection with an optional projection, sort, and limit.
+
+    Parameters
+    ----------
+    filter : str
+        A JSON string representing the MongoDB filter query.
+    project : str, optional
+        A JSON string specifying fields to include or exclude in the result (MongoDB projection).
+    sort : str, optional
+        A JSON string specifying sorting criteria (e.g., '[["started_at", -1]]').
+    limit : int, optional
+        Maximum number of documents to return. Default is 0 (no limit).
+
+    Returns
+    -------
+    List[dict]
+        A list of task documents matching the query.
+    """
+    _filter = json.loads(filter)
+    _project = json.loads(project) or None
+    _sort = list(sort) or None
+    print(
+        json.dumps(Flowcept.db.query(filter=_filter, project=_project, sort=_sort, limit=limit), indent=2, default=str)
+    )
+
+
+def get_task(task_id: str):
     """
-    Query the Document DB.
+    Query the Document DB to retrieve a task.
 
     Parameters
     ----------
-    query_str : str
-        A JSON string representing the Mongo query.
+    task_id : str
+        The identifier of the task.
+    """
+    _query = {"task_id": task_id}
+    print(json.dumps(Flowcept.db.query(_query), indent=2, default=str))
+
+
+def start_agent():
+    """Start Flowcept agent."""
+    from flowcept.flowceptor.adapters.agents.flowcept_agent import main
+
+    main()
+
+
+def agent_client(tool_name: str, kwargs: str = None):
+    """Agent Client.
+
+    Parameters.
+    ----------
+    tool_name : str
+        Name of the tool
+    kwargs : str, optional
+        A stringfied JSON containing the kwargs for the tool, if needed.
+    """
+    print(kwargs)
+    if kwargs is not None:
+        kwargs = json.loads(kwargs)
+
+    print(f"Going to run agent tool '{tool_name}'.")
+    if kwargs:
+        print(f"Using kwargs: {kwargs}")
+    print("-----------------")
+    from flowcept.flowceptor.consumers.agent.client_agent import run_tool
+
+    result = run_tool(tool_name, kwargs)[0]
+
+    print(result.text)
+
+
+def check_services():
+    """
+    Run a full diagnostic test on the Flowcept system and its dependencies.
+
+    This function:
+    - Prints the current configuration path.
+    - Checks if required services (e.g., MongoDB, agent) are alive.
+    - Runs a test function wrapped with Flowcept instrumentation.
+    - Verifies MongoDB insertion (if enabled).
+    - Verifies agent communication and LLM connectivity (if enabled).
+
+    Returns
+    -------
+    None
+        Prints diagnostics to stdout; returns nothing.
     """
-    query = json.loads(query_str)
-    print(Flowcept.db.query(query))
+    print(f"Testing with settings at: {configs.SETTINGS_PATH}")
+    from flowcept.configs import MONGO_ENABLED, AGENT, KVDB_ENABLED, INSERTION_BUFFER_TIME
+
+    if not Flowcept.services_alive():
+        print("Some of the enabled services are not alive!")
+        return
+
+    check_safe_stops = KVDB_ENABLED
+
+    from uuid import uuid4
+    from flowcept.instrumentation.flowcept_task import flowcept_task
+
+    workflow_id = str(uuid4())
+
+    @flowcept_task
+    def test_function(n: int) -> Dict[str, int]:
+        return {"output": n + 1}
+
+    with Flowcept(workflow_id=workflow_id, check_safe_stops=check_safe_stops):
+        test_function(2)
+
+    if MONGO_ENABLED:
+        print("MongoDB is enabled, so we are testing it too.")
+        tasks = Flowcept.db.query({"workflow_id": workflow_id})
+        if len(tasks) != 1:
+            print(f"The query result, {len(tasks)}, is not what we expected.")
+            return
+
+    if AGENT.get("enabled", False):
+        print("Agent is enabled, so we are testing it too.")
+        from flowcept.flowceptor.consumers.agent.client_agent import run_tool
+
+        try:
+            print(run_tool("check_liveness"))
+        except Exception as e:
+            print(e)
+            return
+
+        print("Testing LLM connectivity")
+        check_llm_result = run_tool("check_llm")[0]
+        print(check_llm_result.text)
+
+        if "error" in check_llm_result.text.lower():
+            print("There is an error with the LLM communication.")
+            return
+        elif MONGO_ENABLED:
+            print("Testing if llm chat was stored in MongoDB.")
+            response_metadata = json.loads(check_llm_result.text.split("\n")[0])
+            print(response_metadata)
+            sleep(INSERTION_BUFFER_TIME * 1.05)
+            chats = Flowcept.db.query({"workflow_id": response_metadata["agent_id"]})
+            if chats:
+                print(chats)
+            else:
+                print("Could not find chat history. Make sure that the DB Inserter service is on.")
+    print("\n\nAll expected services seem to be working properly!")
+    return
 
 
 COMMAND_GROUPS = [
-    ("Basic Commands", [show_config, start_services, stop_services]),
+    ("Basic Commands", [check_services, show_config, init_settings, start_services, stop_services]),
     ("Consumption Commands", [start_consumption_services, stop_consumption_services]),
-    ("Database Commands", [workflow_count, query]),
+    ("Database Commands", [workflow_count, query, get_task]),
+    ("Agent Commands", [start_agent, agent_client]),
 ]
 
 COMMANDS = set(f for _, fs in COMMAND_GROUPS for f in fs)
 
 
+def _run_command(cmd_str: str, check_output: bool = True, popen_kwargs: Optional[Dict] = None) -> Optional[str]:
+    """
+    Run a shell command with optional output capture.
+
+    Parameters
+    ----------
+    cmd_str : str
+        The command to execute.
+    check_output : bool, optional
+        If True, capture and return the command's standard output.
+        If False, run interactively (stdout/stderr goes to terminal).
+    popen_kwargs : dict, optional
+        Extra keyword arguments to pass to subprocess.run.
+
+    Returns
+    -------
+    output : str or None
+        The standard output of the command if check_output is True, else None.
+
+    Raises
+    ------
+    subprocess.CalledProcessError
+        If the command exits with a non-zero status.
+    """
+    if popen_kwargs is None:
+        popen_kwargs = {}
+
+    kwargs = {"shell": True, "check": True, **popen_kwargs}
+
+    if check_output:
+        kwargs.update({"capture_output": True, "text": True})
+        result = subprocess.run(cmd_str, **kwargs)
+        return result.stdout.strip()
+    else:
+        subprocess.run(cmd_str, **kwargs)
+        return None
+
+
 def _parse_numpy_doc(docstring: str):
     parsed = {}
     lines = docstring.splitlines() if docstring else []
@@ -258,3 +457,4 @@ def main():  # noqa: D103
 
 if __name__ == "__main__":
     main()
+    # check_services()

flowcept/commons/daos/keyvalue_dao.py

@@ -1,6 +1,6 @@
 """Key value module."""
 
-from redis import Redis, ConnectionPool
+from flowcept.commons.daos.redis_conn import RedisConn
 
 from flowcept.commons.flowcept_logger import FlowceptLogger
 from flowcept.configs import (
@@ -24,32 +24,13 @@ class KeyValueDAO:
             cls._instance = super(KeyValueDAO, cls).__new__(cls)
         return cls._instance
 
-    @staticmethod
-    def build_redis_conn_pool():
-        """Utility function to build Redis connection."""
-        pool = ConnectionPool(
-            host=KVDB_HOST,
-            port=KVDB_PORT,
-            db=0,
-            password=KVDB_PASSWORD,
-            decode_responses=False,
-            max_connections=10000,  # TODO: Config file
-            socket_keepalive=True,
-            retry_on_timeout=True,
-        )
-        return Redis(connection_pool=pool)
-        # return Redis()
-
     def __init__(self):
         if not hasattr(self, "_initialized"):
             self._initialized = True
             self.logger = FlowceptLogger()
-            if KVDB_URI is not None:
-                # If a URI is provided, use it for connection
-                self.redis_conn = Redis.from_url(KVDB_URI)
-            else:
-                # Otherwise, use the host, port, and password settings
-                self.redis_conn = KeyValueDAO.build_redis_conn_pool()
+            self.redis_conn = RedisConn.build_redis_conn_pool(
+                host=KVDB_HOST, port=KVDB_PORT, password=KVDB_PASSWORD, uri=KVDB_URI
+            )
 
     def delete_set(self, set_name: str):
         """Delete it."""
@@ -133,3 +114,18 @@ class KeyValueDAO:
         None
         """
         self.redis_conn.delete(key)
+
+    def liveness_test(self):
+        """Get the livelyness of it."""
+        try:
+            response = self.redis_conn.ping()
+            if response:
+                return True
+            else:
+                return False
+        except ConnectionError as e:
+            self.logger.exception(e)
+            return False
+        except Exception as e:
+            self.logger.exception(e)
+            return False

flowcept/commons/daos/mq_dao/mq_dao_base.py

@@ -20,6 +20,7 @@ from flowcept.configs import (
     MQ_CHUNK_SIZE,
     MQ_TYPE,
     MQ_TIMING,
+    KVDB_ENABLED,
 )
 
 from flowcept.commons.utils import GenericJSONEncoder
@@ -67,7 +68,13 @@ class MQDao(ABC):
         self.logger = FlowceptLogger()
         self.started = False
         self._adapter_settings = adapter_settings
-        self._keyvalue_dao = KeyValueDAO()
+        if KVDB_ENABLED:
+            self._keyvalue_dao = KeyValueDAO()
+        else:
+            self._keyvalue_dao = None
+            self.logger.warning(
+                "We are going to run without KVDB. If you are running a workflow, this may lead to errors."
+            )
         self._time_based_flushing_started = False
         self.buffer: Union[AutoflushBuffer, List] = None
         if MQ_TIMING:
@@ -138,7 +145,7 @@ class MQDao(ABC):
         """
         self._keyvalue_dao.delete_key("current_campaign_id")
 
-    def init_buffer(self, interceptor_instance_id: str, exec_bundle_id=None):
+    def init_buffer(self, interceptor_instance_id: str, exec_bundle_id=None, check_safe_stops=True):
         """Create the buffer."""
         if not self.started:
             if flowcept.configs.DB_FLUSH_MODE == "online":
@@ -147,7 +154,8 @@ class MQDao(ABC):
                     max_size=MQ_BUFFER_SIZE,
                     flush_interval=MQ_INSERTION_BUFFER_TIME,
                 )
-                self.register_time_based_thread_init(interceptor_instance_id, exec_bundle_id)
+                if check_safe_stops:
+                    self.register_time_based_thread_init(interceptor_instance_id, exec_bundle_id)
                 self._time_based_flushing_started = True
             else:
                 self.buffer = list()
@@ -164,9 +172,9 @@ class MQDao(ABC):
             self.bulk_publish(self.buffer)
             self.buffer = list()
 
-    def _stop_timed(self, interceptor_instance_id: str, bundle_exec_id: int = None):
+    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, bundle_exec_id)
+        self._stop(interceptor_instance_id, check_safe_stops, bundle_exec_id)
         t2 = time()
         self._flush_events.append(["final", t1, t2, t2 - t1, "n/a"])
 
@@ -175,14 +183,14 @@ class MQDao(ABC):
             writer.writerow(["type", "start", "end", "duration", "size"])
             writer.writerows(self._flush_events)
 
-    def _stop(self, interceptor_instance_id: str, bundle_exec_id: int = None):
-        """Stop it."""
-        msg0 = "MQ publisher received stop signal! bundle: "
-        self.logger.debug(msg0 + f"{bundle_exec_id}; interceptor id: {interceptor_instance_id}")
+    def _stop(self, interceptor_instance_id: str, 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()
-        msg = "Flushed MQ for last time! Send stop msg. bundle: "
-        self.logger.debug(msg + f"{bundle_exec_id}; interceptor id: {interceptor_instance_id}")
-        self._send_mq_dao_time_thread_stop(interceptor_instance_id, bundle_exec_id)
+        self.logger.debug("Flushed MQ for the last time!")
+        if check_safe_stops:
+            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
 
     def _send_mq_dao_time_thread_stop(self, interceptor_instance_id, exec_bundle_id=None):
@@ -197,10 +205,10 @@ class MQDao(ABC):
         # self.logger.info("Control msg sent: " + str(msg))
         self.send_message(msg)
 
-    def send_document_inserter_stop(self):
+    def send_document_inserter_stop(self, exec_bundle_id=None):
         """Send the document."""
         # These control_messages are handled by the document inserter
-        msg = {"type": "flowcept_control", "info": "stop_document_inserter"}
+        msg = {"type": "flowcept_control", "info": "stop_document_inserter", "exec_bundle_id": exec_bundle_id}
         self.send_message(msg)
 
     @abstractmethod
@@ -223,20 +231,12 @@ class MQDao(ABC):
         """Subscribe to the interception channel."""
         raise NotImplementedError()
 
+    @abstractmethod
+    def unsubscribe(self):
+        """Subscribe to the interception channel."""
+        raise NotImplementedError()
+
     @abstractmethod
     def liveness_test(self) -> bool:
-        """Check whether the base KV store's connection is ready. This is enough for Redis MQ too. Other MQs need
-        further liveness implementation.
-        """
-        try:
-            response = self._keyvalue_dao.redis_conn.ping()
-            if response:
-                return True
-            else:
-                return False
-        except ConnectionError as e:
-            self.logger.exception(e)
-            return False
-        except Exception as e:
-            self.logger.exception(e)
-            return False
+        """Checks if the MQ system is alive."""
+        raise NotImplementedError()

flowcept/commons/daos/mq_dao/mq_dao_kafka.py

@@ -108,12 +108,13 @@ class MQDaoKafka(MQDao):
     def liveness_test(self):
         """Get the livelyness of it."""
         try:
-            if not super().liveness_test():
-                self.logger.error("KV Store not alive!")
-                return False
             admin_client = AdminClient(self._kafka_conf)
             kafka_metadata = admin_client.list_topics(timeout=5)
             return MQ_CHANNEL in kafka_metadata.topics
         except Exception as e:
             self.logger.exception(e)
             return False
+
+    def unsubscribe(self):
+        """Unsubscribes from Kafka topic."""
+        raise NotImplementedError()

flowcept/commons/daos/mq_dao/mq_dao_mofka.py

@@ -104,3 +104,7 @@ class MQDaoMofka(MQDao):
     def liveness_test(self):
         """Test Mofka Liveness."""
         return True
+
+    def unsubscribe(self):
+        """Unsubscribes from Mofka topic."""
+        raise NotImplementedError()

flowcept/commons/daos/mq_dao/mq_dao_redis.py

@@ -7,7 +7,8 @@ import msgpack
 from time import time, sleep
 
 from flowcept.commons.daos.mq_dao.mq_dao_base import MQDao
-from flowcept.configs import MQ_CHANNEL
+from flowcept.commons.daos.redis_conn import RedisConn
+from flowcept.configs import MQ_CHANNEL, MQ_HOST, MQ_PORT, MQ_PASSWORD, MQ_URI, MQ_SETTINGS, KVDB_ENABLED
 
 
 class MQDaoRedis(MQDao):
@@ -17,16 +18,32 @@ class MQDaoRedis(MQDao):
 
     def __init__(self, adapter_settings=None):
         super().__init__(adapter_settings)
-        self._producer = self._keyvalue_dao.redis_conn  # if MQ is redis, we use the same KV for the MQ
+
         self._consumer = None
+        use_same_as_kv = MQ_SETTINGS.get("same_as_kvdb", False)
+        if use_same_as_kv:
+            if KVDB_ENABLED:
+                self._producer = self._keyvalue_dao.redis_conn
+            else:
+                raise Exception("You have same_as_kvdb in your settings, but kvdb is disabled.")
+        else:
+            self._producer = RedisConn.build_redis_conn_pool(
+                host=MQ_HOST, port=MQ_PORT, password=MQ_PASSWORD, uri=MQ_URI
+            )
 
     def subscribe(self):
         """
         Subscribe to interception channel.
         """
-        self._consumer = self._keyvalue_dao.redis_conn.pubsub()
+        self._consumer = self._producer.pubsub()
         self._consumer.psubscribe(MQ_CHANNEL)
 
+    def unsubscribe(self):
+        """
+        Unsubscribe to interception channel.
+        """
+        self._consumer.unsubscribe(MQ_CHANNEL)
+
     def message_listener(self, message_handler: Callable):
         """Get message listener with automatic reconnection."""
         max_retrials = 10
@@ -37,13 +54,22 @@ class MQDaoRedis(MQDao):
                 for message in self._consumer.listen():
                     if message and message["type"] in MQDaoRedis.MESSAGE_TYPES_IGNORE:
                         continue
+
+                    if not isinstance(message["data"], (bytes, bytearray)):
+                        self.logger.warning(
+                            f"Skipping message with unexpected data type: {type(message['data'])} - {message['data']}"
+                        )
+                        continue
+
                     try:
                         msg_obj = msgpack.loads(message["data"], strict_map_key=False)
+                        # self.logger.debug(f"In mq dao redis, received msg!  {msg_obj}")
                         if not message_handler(msg_obj):
                             should_continue = False  # Break While loop
                             break  # Break For loop
                     except Exception as e:
-                        self.logger.error(f"Failed to process message: {e}")
+                        self.logger.error(f"Failed to process message {message}")
+                        self.logger.exception(e)
 
                     current_trials = 0
             except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError) as e:
@@ -103,7 +129,14 @@ class MQDaoRedis(MQDao):
     def liveness_test(self):
         """Get the livelyness of it."""
         try:
-            return super().liveness_test()
+            response = self._producer.ping()
+            if response:
+                return True
+            else:
+                return False
+        except ConnectionError as e:
+            self.logger.exception(e)
+            return False
         except Exception as e:
             self.logger.exception(e)
             return False

flowcept/commons/daos/redis_conn.py

@@ -0,0 +1,47 @@
+"""RedisConn module."""
+
+from redis import Redis, ConnectionPool
+
+
+class RedisConn:
+    """RedisConn DAO class."""
+
+    @staticmethod
+    def build_redis_conn_pool(host: str = None, port: str = None, password: str = None, uri: str = None) -> Redis:
+        """
+        Create a Redis connection using either a URI or host/port.
+
+        If `uri` is provided, it will be used to initialize the Redis connection.
+        Otherwise, the connection will fall back to using `host` and `port`.
+
+        Parameters
+        ----------
+        host : str, optional
+            Redis host address. Used only if `uri` is not provided.
+        port : str, optional
+            Redis port. Used only if `uri` is not provided.
+        uri : str, optional
+            Full Redis URI. Takes precedence over `host` and `port` if defined.
+        password : str, optional
+            Password for authenticating with Redis.
+
+        Returns
+        -------
+        Redis
+            An instance of the Redis client with a configured connection pool.
+        """
+        pool_kwargs = {
+            "db": 0,
+            "password": password,
+            "decode_responses": False,
+            "max_connections": 10000,
+            "socket_keepalive": True,
+            "retry_on_timeout": True,
+        }
+
+        if uri:
+            pool = ConnectionPool.from_url(uri, **pool_kwargs)
+        else:
+            pool = ConnectionPool(host=host, port=port, **pool_kwargs)
+
+        return Redis(connection_pool=pool)

flowcept/commons/flowcept_dataclasses/task_object.py

@@ -30,7 +30,7 @@ class TaskObject:
     submitted_at: float = None
     started_at: float = None
     ended_at: float = None
-    registered_at: float = None
+    registered_at: float = None  # Leave this for dates generated at the DocInserter
     telemetry_at_start: Telemetry = None
     telemetry_at_end: Telemetry = None
     workflow_name: AnyStr = None
@@ -52,6 +52,7 @@ class TaskObject:
     address: AnyStr = None
     dependencies: List = None
     dependents: List = None
+    tags: List = None
 
     _DEFAULT_ENRICH_VALUES = {
         "node_name": NODE_NAME,
@@ -145,10 +146,37 @@ class TaskObject:
             if (key not in task_dict or task_dict[key] is None) and fallback_value is not None:
                 task_dict[key] = fallback_value
 
-    # @staticmethod
-    # def deserialize(serialized_data) -> 'TaskObject':
-    #     dict_obj = msgpack.loads(serialized_data)
-    #     obj = TaskObject()
-    #     for k, v in dict_obj.items():
-    #         setattr(obj, k, v)
-    #     return obj
+    @staticmethod
+    def from_dict(task_obj_dict: Dict[AnyStr, Any]) -> "TaskObject":
+        """Create a TaskObject from a dictionary.
+
+        Parameters
+        ----------
+        task_obj_dict : Dict[AnyStr, Any]
+            Dictionary containing task attributes.
+
+        Returns
+        -------
+        TaskObject
+            A TaskObject instance populated with available data.
+        """
+        task = TaskObject()
+
+        for key, value in task_obj_dict.items():
+            if hasattr(task, key):
+                if key == "status" and isinstance(value, str):
+                    setattr(task, key, Status(value))
+                else:
+                    setattr(task, key, value)
+
+        return task
+
+    def __str__(self):
+        """Return a user-friendly string representation of the TaskObject."""
+        return self.__repr__()
+
+    def __repr__(self):
+        """Return an unambiguous string representation of the TaskObject."""
+        attrs = ["task_id", "workflow_id", "campaign_id", "activity_id", "custom_metadata", "started_at", "ended_at"]
+        attr_str = ", ".join(f"{attr}={repr(getattr(self, attr))}" for attr in attrs)
+        return f"TaskObject({attr_str})"