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

flowcept/flowceptor/adapters/brokers/mqtt_interceptor.py

@@ -0,0 +1,132 @@
+"""Zambeze interceptor module."""
+
+import uuid
+from threading import Thread
+from time import sleep
+import paho.mqtt.client as mqtt
+import json
+from typing import Dict
+from flowcept.commons.flowcept_dataclasses.task_object import TaskObject
+from flowcept.flowcept_api.flowcept_controller import Flowcept
+from flowcept.flowceptor.adapters.base_interceptor import (
+    BaseInterceptor,
+)
+
+
+class MQTTBrokerInterceptor(BaseInterceptor):
+    """Zambeze interceptor."""
+
+    def __init__(self, plugin_key="broker_mqtt"):
+        super().__init__(plugin_key)
+
+        assert self.settings.get("protocol") == "mqtt3.1.1", "We only support mqtt3.1.1 for this interceptor."
+        self._host = self.settings.get("host", "localhost")
+        self._port = self.settings.get("port", 1883)
+        self._username = self.settings.get("username", "username")
+        self._password = self.settings.get("password", None)
+        self._queues = self.settings.get("queues")
+        self._qos = self.settings.get("qos", 2)
+        self._id = str(id(self))
+
+        self._tracked_keys = self.settings.get("tracked_keys")
+        self._task_subtype = self.settings.get("task_subtype", None)
+        self._client: mqtt.Client = None
+
+        self._observer_thread: Thread = None
+
+    def _connect(self):
+        """Establish a connection to the MQTT broker."""
+        try:
+            self._client = mqtt.Client(client_id=self._id, clean_session=False, protocol=mqtt.MQTTv311)
+            self._client.username_pw_set(self._username, self._password)
+
+            self._client.on_message = self.callback
+            self._client.on_connect = self._on_connect
+            self._client.on_disconnect = self._on_disconnect
+
+            self.logger.debug("Connecting to MQTT broker...")
+            self._client.connect(self._host, self._port, 60)
+            self.logger.debug("Connected.")
+        except Exception as e:
+            self.logger.error(f"Connection failed: {e}")
+            raise e
+
+    def _on_connect(self, *_):
+        """Handle connection events and subscribe to the topic."""
+        for q in self._queues:
+            self.logger.debug(f"Client {self._id} connected to MQTT queue {q}. Waiting for messages...")
+            self._client.subscribe(q, qos=self._qos)
+
+    def callback(self, _, __, msg):
+        """Implement the callback."""
+        msg_str = msg.payload.decode()
+        topic = msg.topic
+        self.logger.debug(f"Received message: '{msg_str}' on topic '{topic}'")
+
+        msg_dict = json.loads(msg_str)
+        msg_dict["topic"] = topic
+
+        task_msg = self.prepare_task_msg(msg_dict)
+        self.intercept(task_msg.to_dict())
+
+    def _on_disconnect(self, *_):
+        """Handle disconnections and attempt reconnection."""
+        self.logger.warning("MQTT Observer Client Disconnected.")
+
+    def start(self, bundle_exec_id) -> "MQTTBrokerInterceptor":
+        """Start it."""
+        super().start(bundle_exec_id)
+        self._observer_thread = Thread(target=self.observe, daemon=True)
+        self._observer_thread.start()
+        return self
+
+    def observe(self):
+        """Start the MQTT loop."""
+        self._connect()
+        self._client.loop_forever()
+
+    def prepare_task_msg(self, msg: Dict) -> TaskObject:
+        """Prepare a task."""
+        task_dict = {}
+        custom_metadata = {"topic": msg.get("topic", None)}
+        for key in self._tracked_keys:
+            if key != "custom_metadata":
+                if self._tracked_keys.get(key):
+                    task_dict[key] = msg.get(self._tracked_keys.get(key), None)
+            else:
+                cm = self._tracked_keys.get("custom_metadata", None)
+                if cm and len(cm):
+                    for k in cm:
+                        custom_metadata[k] = msg[k]
+                    task_dict["custom_metadata"] = custom_metadata
+
+        if isinstance(task_dict.get("used"), str):
+            task_dict["used"] = {"payload": task_dict.get("used")}
+
+        if "task_id" not in task_dict:
+            task_dict["task_id"] = str(uuid.uuid4())
+
+        task_obj = TaskObject.from_dict(task_dict)
+        task_obj.subtype = self._task_subtype
+
+        if task_obj.campaign_id is None:
+            task_obj.campaign_id = Flowcept.campaign_id
+
+        if task_obj.workflow_id is None:
+            task_obj.workflow_id = Flowcept.current_workflow_id
+
+        print(task_obj)
+        return task_obj
+
+    def stop(self) -> bool:
+        """Stop it."""
+        self.logger.debug("Interceptor stopping...")
+        super().stop()
+        try:
+            self._client.disconnect()
+        except Exception as e:
+            self.logger.warning(f"This exception is expected to occur after channel.basic_cancel: {e}")
+        sleep(2)
+        self._observer_thread.join()
+        self.logger.debug("Interceptor stopped.")
+        return True

flowcept/flowceptor/adapters/mlflow/mlflow_interceptor.py

@@ -64,17 +64,17 @@ class MLFlowInterceptor(BaseInterceptor):
                 task_msg = self.prepare_task_msg(run_data).to_dict()
                 self.intercept(task_msg)
 
-    def start(self, bundle_exec_id) -> "MLFlowInterceptor":
+    def start(self, bundle_exec_id, check_safe_stops) -> "MLFlowInterceptor":
         """Start it."""
         super().start(bundle_exec_id)
         self._observer_thread = Thread(target=self.observe, daemon=True)
         self._observer_thread.start()
         return self
 
-    def stop(self) -> bool:
+    def stop(self, check_safe_stops: bool = True) -> bool:
         """Stop it."""
         sleep(1)
-        super().stop()
+        super().stop(check_safe_stops)
         self.logger.debug("Interceptor stopping...")
         self._observer.stop()
         self._observer_thread.join()

flowcept/flowceptor/adapters/tensorboard/tensorboard_interceptor.py

@@ -91,17 +91,17 @@ class TensorboardInterceptor(BaseInterceptor):
                 self.intercept(task_msg.to_dict())
                 self.state_manager.add_element_id(child_event.log_path)
 
-    def start(self, bundle_exec_id) -> "TensorboardInterceptor":
+    def start(self, bundle_exec_id, check_safe_stops: bool = True) -> "TensorboardInterceptor":
         """Start it."""
         super().start(bundle_exec_id)
         self.observe()
         return self
 
-    def stop(self) -> bool:
+    def stop(self, check_safe_stops: bool = True) -> bool:
         """Stop it."""
         sleep(1)
         self.logger.debug("Interceptor stopping...")
-        super().stop()
+        super().stop(check_safe_stops)
         self._observer.stop()
         self.logger.debug("Interceptor stopped.")
         return True

flowcept/flowceptor/consumers/agent/__init__.py

@@ -0,0 +1 @@
+"""Flowcept agent and Flowcept-enabled agent module."""

flowcept/flowceptor/consumers/agent/base_agent_context_manager.py

@@ -0,0 +1,101 @@
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import Dict, List
+
+from flowcept.flowceptor.consumers.base_consumer import BaseConsumer
+
+
+@dataclass
+class BaseAppContext:
+    """
+    Container for storing agent context data during the lifespan of an application session.
+
+    Attributes
+    ----------
+    tasks : list of dict
+        A list of task messages received from the message queue. Each task message is stored as a dictionary.
+    """
+
+    tasks: List[Dict]
+
+
+class BaseAgentContextManager(BaseConsumer):
+    """
+    Base class for any MCP Agent that wants to participate in the Flowcept ecosystem.
+
+    Agents inheriting from this class can:
+    - Subscribe to and consume messages from the Flowcept-compatible message queue (MQ)
+    - Handle task-related messages and accumulate them in context
+    - Gracefully manage their lifecycle using an async context manager
+    - Interact with Flowcept’s provenance system to read/write messages, query databases, and store chat history
+
+    To integrate with Flowcept:
+    - Inherit from `BaseAgentContextManager`
+    - Override `message_handler()` if custom message handling is needed
+    - Access shared state via `self.context` during execution
+    """
+
+    def __init__(self):
+        """
+        Initializes the agent and resets its context state.
+        """
+        super().__init__()
+        self.context = None
+        self.reset_context()
+
+    def message_handler(self, msg_obj: Dict) -> bool:
+        """
+        Handles a single message received from the message queue.
+
+        Parameters
+        ----------
+        msg_obj : dict
+            The message received, typically structured with a "type" field.
+
+        Returns
+        -------
+        bool
+            Return True to continue listening for messages, or False to stop the loop.
+
+        Notes
+        -----
+        This default implementation stores messages of type 'task' in the internal context.
+        Override this method in a subclass to handle other message types or implement custom logic.
+        """
+        msg_type = msg_obj.get("type", None)
+        msg_subtype = msg_obj.get("subtype", "")
+        if msg_type == "task":
+            self.logger.debug("Received task msg!")
+            if msg_subtype not in {"llm_query"}:
+                self.context.tasks.append(msg_obj)
+
+        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):
+        """
+        Async context manager to handle the agent’s lifecycle within an application.
+
+        Starts the message consumption when the context is entered and stops it when exited.
+
+        Parameters
+        ----------
+        app : Any
+            The application instance using this context (typically unused but included for compatibility).
+
+        Yields
+        ------
+        BaseAppContext
+            The current application context, including collected tasks.
+        """
+        self.start()
+        try:
+            yield self.context
+        finally:
+            self.stop_consumption()

flowcept/flowceptor/consumers/agent/client_agent.py

@@ -0,0 +1,48 @@
+from typing import Dict, List
+
+from flowcept.configs import AGENT
+from mcp import ClientSession
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.types import TextContent
+
+MCP_HOST = AGENT.get("mcp_host", "0.0.0.0")
+MCP_PORT = AGENT.get("mcp_port", 8000)
+
+MCP_URL = f"http://{MCP_HOST}:{MCP_PORT}/mcp"
+
+
+def run_tool(tool_name: str, kwargs: Dict = None) -> List[TextContent]:
+    """
+    Run a tool using an MCP client session via a local streamable HTTP connection.
+
+    This function opens an asynchronous connection to a local MCP server,
+    initializes a session, and invokes a specified tool with optional arguments.
+    The tool's response content is returned as a list of `TextContent` objects.
+
+    Parameters
+    ----------
+    tool_name : str
+        The name of the tool to call within the MCP framework.
+    kwargs : Dict, optional
+        A dictionary of keyword arguments to pass as input to the tool. Defaults to None.
+
+    Returns
+    -------
+    List[TextContent]
+        A list of `TextContent` objects returned by the tool execution.
+
+    Notes
+    -----
+    This function uses `asyncio.run`, so it must not be called from an already-running
+    event loop (e.g., inside another async function in environments like Jupyter).
+    """
+    import asyncio
+
+    async def _run():
+        async with streamablehttp_client(MCP_URL) as (read, write, _):
+            async with ClientSession(read, write) as session:
+                await session.initialize()
+                result = await session.call_tool(tool_name, arguments=kwargs)
+                return result.content
+
+    return asyncio.run(_run())

flowcept/flowceptor/consumers/agent/flowcept_agent_context_manager.py

@@ -0,0 +1,145 @@
+from dataclasses import dataclass
+from typing import Dict, List
+
+from flowcept.flowceptor.consumers.agent.base_agent_context_manager import BaseAgentContextManager, BaseAppContext
+from langchain.chains.retrieval_qa.base import BaseRetrievalQA
+from langchain_community.embeddings import HuggingFaceEmbeddings
+
+from flowcept.flowceptor.consumers.agent import client_agent
+from flowcept.flowceptor.consumers.agent.flowcept_qa_manager import FlowceptQAManager
+from flowcept.commons.task_data_preprocess import summarize_task
+
+
+@dataclass
+class FlowceptAppContext(BaseAppContext):
+    """
+    Context object for holding flowcept-specific state (e.g., tasks data) during the agent's lifecycle.
+
+    Attributes
+    ----------
+    task_summaries : List[Dict]
+        List of summarized task dictionaries.
+    critical_tasks : List[Dict]
+        List of critical task summaries with tags or anomalies.
+    qa_chain : BaseRetrievalQA
+        The QA chain used for question-answering over task summaries.
+    vectorstore_path : str
+        Path to the persisted vectorstore used by the QA system.
+    embedding_model : HuggingFaceEmbeddings
+        The embedding model used to generate vector representations for tasks.
+    """
+
+    task_summaries: List[Dict]
+    critical_tasks: List[Dict]
+    qa_chain: BaseRetrievalQA
+    vectorstore_path: str
+    embedding_model: HuggingFaceEmbeddings
+
+
+class FlowceptAgentContextManager(BaseAgentContextManager):
+    """
+    Manages agent context and operations for Flowcept's intelligent task monitoring.
+
+    This class extends BaseAgentContextManager and maintains a rolling buffer of task messages.
+    It summarizes and tags tasks, builds a QA index over them, and uses LLM tools to analyze
+    task batches periodically.
+
+    Attributes
+    ----------
+    context : FlowceptAppContext
+        Current application context holding task state and QA components.
+    msgs_counter : int
+        Counter tracking how many task messages have been processed.
+    context_size : int
+        Number of task messages to collect before triggering QA index building and LLM analysis.
+    qa_manager : FlowceptQAManager
+        Utility for constructing QA chains from task summaries.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.context: FlowceptAppContext = None
+        self.reset_context()
+        self.msgs_counter = 0
+        self.context_size = 5
+        self.qa_manager = FlowceptQAManager()
+
+    def message_handler(self, msg_obj: Dict):
+        """
+        Handle an incoming message and update context accordingly.
+
+        Parameters
+        ----------
+        msg_obj : Dict
+            The incoming message object.
+
+        Returns
+        -------
+        bool
+            True if the message was handled successfully.
+        """
+        msg_type = msg_obj.get("type", None)
+        if msg_type == "task":
+            self.msgs_counter += 1
+            self.logger.debug("Received task msg!")
+            self.context.tasks.append(msg_obj)
+
+            self.logger.debug(f"This is QA! {self.context.qa_chain}")
+
+            task_summary = summarize_task(msg_obj)
+            self.context.task_summaries.append(task_summary)
+            if len(task_summary.get("tags", [])):
+                self.context.critical_tasks.append(task_summary)
+
+            if self.msgs_counter > 0 and self.msgs_counter % self.context_size == 0:
+                self.build_qa_index()
+
+                self.monitor_chunk()
+
+        return True
+
+    def monitor_chunk(self):
+        """
+        Perform LLM-based analysis on the current chunk of task messages and send the results.
+        """
+        self.logger.debug(f"Going to begin LLM job! {self.msgs_counter}")
+        result = client_agent.run_tool("analyze_task_chunk")
+        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))
+            else:
+                self.logger.error(content)
+
+    def build_qa_index(self):
+        """
+        Build a new QA index from the current list of task summaries.
+        """
+        self.logger.debug(f"Going to begin QA Build! {self.msgs_counter}")
+        try:
+            qa_chain_result = self.qa_manager.build_qa(docs=self.context.task_summaries)
+
+            self.context.qa_chain = qa_chain_result.get("qa_chain")
+            self.context.vectorstore_path = qa_chain_result.get("path")
+
+            self.logger.debug(f"Built QA! {self.msgs_counter}")
+            assert self.context.qa_chain is not None
+            self.logger.debug(f"This is QA! {self.context.qa_chain}")
+            self.logger.debug(f"This is QA path! {self.context.vectorstore_path}")
+        except Exception as e:
+            self.logger.exception(e)
+
+    def reset_context(self):
+        """
+        Reset the agent's context to a clean state, initializing a new QA setup.
+        """
+        self.context = FlowceptAppContext(
+            tasks=[],
+            task_summaries=[],
+            critical_tasks=[],
+            qa_chain=None,
+            vectorstore_path=None,
+            embedding_model=FlowceptQAManager.embedding_model,
+        )

flowcept/flowceptor/consumers/agent/flowcept_qa_manager.py

@@ -0,0 +1,112 @@
+from typing import Dict, List
+
+from langchain.chains.retrieval_qa.base import RetrievalQA
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain.schema import Document
+from langchain_core.language_models import LLM
+
+from flowcept.commons.flowcept_logger import FlowceptLogger
+from flowcept.flowceptor.adapters.agents.agents_utils import build_llm_model
+
+
+# TODO If all methods are static, this doesnt need to be a class.
+class FlowceptQAManager(object):
+    """
+    Manager for building and loading question-answering (QA) chains using LangChain.
+
+    This utility constructs a `RetrievalQA` chain by converting task dictionaries into
+    `Document` objects, embedding them with HuggingFace, storing them in a FAISS vectorstore,
+    and returning a ready-to-query QA pipeline.
+
+    Attributes
+    ----------
+    embedding_model : HuggingFaceEmbeddings
+        The default embedding model used to embed documents into vector representations.
+    """
+
+    embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+    @staticmethod
+    def build_qa(docs: List[Dict] = None, llm: LLM = None):
+        """
+        Build a RetrievalQA chain from a list of task dictionaries.
+
+        Parameters
+        ----------
+        docs : List[Dict], optional
+            A list of task dictionaries to be converted into retrievable documents.
+        llm : LLM, optional
+            The language model to use for the QA chain. If None, a default model is built.
+
+        Returns
+        -------
+        dict
+            A dictionary containing:
+            - 'qa_chain': the constructed RetrievalQA chain
+            - 'path': local path where the FAISS vectorstore is saved
+
+        Notes
+        -----
+        If no documents are provided, the method returns None.
+        """
+        if not len(docs):
+            return None
+
+        if llm is None:
+            llm = build_llm_model()
+
+        documents = []
+        for d in docs:
+            content = str(d)  # convert the dict to a string
+            metadata = {"task_id": d.get("task_id", "unknown")}
+            documents.append(Document(page_content=content, metadata=metadata))
+
+        FlowceptLogger().debug(f"Number of documents to index: {len(documents)}")
+        vectorstore = FAISS.from_documents(documents=documents, embedding=FlowceptQAManager.embedding_model)
+        path = "/tmp/qa_index"
+        vectorstore.save_local(path)
+
+        retriever = vectorstore.as_retriever()
+        qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever, return_source_documents=True)
+
+        return {"qa_chain": qa_chain, "path": path}
+
+    @staticmethod
+    def _load_qa_chain(path, llm=None, embedding_model=None) -> RetrievalQA:
+        if embedding_model is None:
+            embedding_model = FlowceptQAManager.embedding_model
+        if llm is None:
+            llm = build_llm_model()
+
+        vectorstore = FAISS.load_local(path, embeddings=embedding_model, allow_dangerous_deserialization=True)
+
+        retriever = vectorstore.as_retriever()
+
+        return RetrievalQA.from_chain_type(llm=llm, retriever=retriever, return_source_documents=True)
+
+    @staticmethod
+    def build_qa_chain_from_vectorstore_path(vectorstore_path, llm=None) -> RetrievalQA:
+        """
+        Build a RetrievalQA chain from an existing vectorstore path.
+
+        Parameters
+        ----------
+        vectorstore_path : str
+            Path to the FAISS vectorstore previously saved to disk.
+        llm : LLM, optional
+            Language model to use. If None, a default model is built.
+
+        Returns
+        -------
+        RetrievalQA
+            A RetrievalQA chain constructed using the loaded vectorstore.
+        """
+        if llm is None:
+            llm = build_llm_model()  # TODO: consider making this llm instance static
+        qa_chain = FlowceptQAManager._load_qa_chain(
+            path=vectorstore_path,  # Only here we really need the QA. We might no
+            llm=llm,
+            embedding_model=FlowceptQAManager.embedding_model,
+        )
+        return qa_chain

flowcept/flowceptor/consumers/base_consumer.py

@@ -0,0 +1,90 @@
+from abc import abstractmethod
+from threading import Thread
+from typing import Callable, Dict, Tuple, Optional
+
+from flowcept.commons.daos.mq_dao.mq_dao_base import MQDao
+from flowcept.commons.flowcept_logger import FlowceptLogger
+
+
+class BaseConsumer(object):
+    """
+    Abstract base class for message consumers in a pub-sub architecture.
+
+    This class provides a standard interface and shared logic for subscribing to
+    message queues and dispatching messages to a handler.
+    """
+
+    def __init__(self):
+        """Initialize the message queue DAO and logger."""
+        self._mq_dao = MQDao.build()
+        self.logger = FlowceptLogger()
+        self._main_thread: Optional[Thread] = None
+
+    @abstractmethod
+    def message_handler(self, msg_obj: Dict) -> bool:
+        """
+        Handle a single incoming message.
+
+        Parameters
+        ----------
+        msg_obj : dict
+            The parsed message object received from the queue.
+
+        Returns
+        -------
+        bool
+            Return False to break the message listener loop.
+            Return True to continue listening.
+        """
+        pass
+
+    def start(self, target: Callable = None, args: Tuple = (), threaded: bool = True, daemon=False):
+        """
+        Start the consumer by subscribing and launching the message handler.
+
+        Parameters
+        ----------
+        target : Callable
+            The function to run for listening to messages (usually the message loop).
+        args : tuple, optional
+            Arguments to pass to the target function.
+        threaded : bool, default=True
+            Whether to run the target function in a background thread.
+        daemon: bool
+
+        Returns
+        -------
+        BaseConsumer
+            The current instance (to allow chaining).
+        """
+        if target is None:
+            target = self.default_thread_target
+        self._mq_dao.subscribe()
+        if threaded:
+            self._main_thread = Thread(target=target, args=args, daemon=daemon)
+            self._main_thread.start()
+        else:
+            target(*args)
+        return self
+
+    def default_thread_target(self):
+        """
+        The default message consumption loop.
+
+        This method is used as the default thread target when starting the consumer. It listens for
+        messages from the message queue and passes them to the consumer's `message_handler`.
+
+        Typically run in a background thread when `start()` is called without a custom target.
+
+        See Also
+        --------
+        start : Starts the consumer and optionally spawns a background thread to run this method.
+        """
+        self.logger.debug("Going to wait for new messages!")
+        self._mq_dao.message_listener(self.message_handler)
+
+    def stop_consumption(self):
+        """
+        Stop consuming messages by unsubscribing from the message queue.
+        """
+        self._mq_dao.unsubscribe()