flock-core 0.1.1__py3-none-any.whl

flock/core/serializable.py

@@ -0,0 +1,206 @@
+"""Module for serializable objects in the system."""
+
+import json
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, TypeVar
+
+import cloudpickle
+import msgpack
+
+from flock.core.logging import flock_logger, performance_handler
+
+T = TypeVar("T", bound="Serializable")
+
+
+class Serializable(ABC):
+    """Base class for all serializable objects in the system.
+
+    Provides methods for serializing/deserializing objects to various formats
+    with comprehensive logging and performance tracking.
+    """
+
+    @abstractmethod
+    def to_dict(self) -> dict[str, Any]:
+        """Convert instance to dictionary representation."""
+        pass
+
+    @classmethod
+    @abstractmethod
+    def from_dict(cls: type[T], data: dict[str, Any]) -> T:
+        """Create instance from dictionary representation."""
+        pass
+
+    def to_json(self) -> str:
+        """Serialize to JSON string."""
+        try:
+            with performance_handler.track_time("json_serialization"):
+                flock_logger.debug(
+                    "Serializing to JSON",
+                    class_name=self.__class__.__name__,
+                )
+                json_str = json.dumps(self.to_dict())
+                flock_logger.debug(
+                    "JSON serialization complete",
+                    size_bytes=len(json_str),
+                )
+                return json_str
+        except Exception as e:
+            flock_logger.error(
+                "JSON serialization failed",
+                class_name=self.__class__.__name__,
+                error=str(e),
+            )
+            raise
+
+    @classmethod
+    def from_json(cls: type[T], json_str: str) -> T:
+        """Create instance from JSON string."""
+        try:
+            with performance_handler.track_time("json_deserialization"):
+                flock_logger.debug(
+                    "Deserializing from JSON",
+                    class_name=cls.__name__,
+                    size_bytes=len(json_str),
+                )
+                instance = cls.from_dict(json.loads(json_str))
+                flock_logger.debug("JSON deserialization complete")
+                return instance
+        except Exception as e:
+            flock_logger.error(
+                "JSON deserialization failed",
+                class_name=cls.__name__,
+                error=str(e),
+            )
+            raise
+
+    def to_msgpack(self, path: Path | None = None) -> bytes:
+        """Serialize to msgpack bytes."""
+        try:
+            with performance_handler.track_time("msgpack_serialization"):
+                flock_logger.debug(
+                    "Serializing to msgpack",
+                    class_name=self.__class__.__name__,
+                )
+                msgpack_bytes = msgpack.packb(self.to_dict())
+
+                if path:
+                    flock_logger.debug(f"Writing msgpack to file: {path}")
+                    path.write_bytes(msgpack_bytes)
+
+                flock_logger.debug(
+                    "Msgpack serialization complete",
+                    size_bytes=len(msgpack_bytes),
+                    file_path=str(path) if path else None,
+                )
+                return msgpack_bytes
+        except Exception as e:
+            flock_logger.error(
+                "Msgpack serialization failed",
+                class_name=self.__class__.__name__,
+                file_path=str(path) if path else None,
+                error=str(e),
+            )
+            raise
+
+    @classmethod
+    def from_msgpack(cls: type[T], msgpack_bytes: bytes) -> T:
+        """Create instance from msgpack bytes."""
+        try:
+            with performance_handler.track_time("msgpack_deserialization"):
+                flock_logger.debug(
+                    "Deserializing from msgpack",
+                    class_name=cls.__name__,
+                    size_bytes=len(msgpack_bytes),
+                )
+                instance = cls.from_dict(msgpack.unpackb(msgpack_bytes))
+                flock_logger.debug("Msgpack deserialization complete")
+                return instance
+        except Exception as e:
+            flock_logger.error(
+                "Msgpack deserialization failed",
+                class_name=cls.__name__,
+                error=str(e),
+            )
+            raise
+
+    @classmethod
+    def from_msgpack_file(cls: type[T], path: Path) -> T:
+        """Create instance from msgpack file."""
+        try:
+            with performance_handler.track_time("msgpack_file_read"):
+                flock_logger.debug(
+                    f"Reading msgpack from file: {path}",
+                    class_name=cls.__name__,
+                )
+                return cls.from_msgpack(path.read_bytes())
+        except Exception as e:
+            flock_logger.error(
+                "Msgpack file read failed",
+                class_name=cls.__name__,
+                file_path=str(path),
+                error=str(e),
+            )
+            raise
+
+    def to_pickle(self) -> bytes:
+        """Serialize to pickle bytes."""
+        try:
+            with performance_handler.track_time("pickle_serialization"):
+                flock_logger.debug(
+                    "Serializing to pickle",
+                    class_name=self.__class__.__name__,
+                )
+                pickle_bytes = cloudpickle.dumps(self)
+                flock_logger.debug(
+                    "Pickle serialization complete",
+                    size_bytes=len(pickle_bytes),
+                )
+                return pickle_bytes
+        except Exception as e:
+            flock_logger.error(
+                "Pickle serialization failed",
+                class_name=self.__class__.__name__,
+                error=str(e),
+            )
+            raise
+
+    @classmethod
+    def from_pickle(cls, pickle_bytes: bytes) -> T:
+        """Create instance from pickle bytes."""
+        try:
+            with performance_handler.track_time("pickle_deserialization"):
+                flock_logger.debug(
+                    "Deserializing from pickle",
+                    class_name=cls.__name__,
+                    size_bytes=len(pickle_bytes),
+                )
+                instance = cloudpickle.loads(pickle_bytes)
+                flock_logger.debug("Pickle deserialization complete")
+                return instance
+        except Exception as e:
+            flock_logger.error(
+                "Pickle deserialization failed",
+                class_name=cls.__name__,
+                error=str(e),
+            )
+            raise
+
+    @classmethod
+    def from_pickle_file(cls: type[T], path: Path) -> T:
+        """Create instance from pickle file."""
+        try:
+            with performance_handler.track_time("pickle_file_read"):
+                flock_logger.debug(
+                    f"Reading pickle from file: {path}",
+                    class_name=cls.__name__,
+                )
+                return cls.from_pickle(path.read_bytes())
+        except Exception as e:
+            flock_logger.error(
+                "Pickle file read failed",
+                class_name=cls.__name__,
+                file_path=str(path),
+                error=str(e),
+            )
+            raise

flock/core/tools/basic_tools.py

@@ -0,0 +1,98 @@
+import os
+
+import httpx
+from markdownify import markdownify as md
+from tavily import TavilyClient
+
+
+def web_search_tavily(query: str):
+    client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
+    try:
+        response = client.search(query, include_answer=True)  # type: ignore
+        return response
+    except Exception:
+        raise
+
+
+def get_web_content_as_markdown(url: str):
+    try:
+        response = httpx.get(url)
+        response.raise_for_status()
+        markdown = md(response.text)
+        return markdown
+    except Exception:
+        raise
+
+
+def evaluate_math(expression: str) -> float:
+    import dspy
+
+    try:
+        result = dspy.PythonInterpreter({}).execute(expression)
+        return result
+    except Exception:
+        raise
+
+
+def code_eval(python_code: str) -> float:
+    import dspy
+
+    try:
+        result = dspy.PythonInterpreter({}).execute(python_code)
+        return result
+    except Exception:
+        raise
+
+
+def get_current_time() -> str:
+    import datetime
+
+    time = datetime.datetime.now().isoformat()
+    return time
+
+
+def count_words(text: str) -> int:
+    count = len(text.split())
+    return count
+
+
+def extract_urls(text: str) -> list[str]:
+    import re
+
+    url_pattern = r"https?://(?:[-\w.]|(?:%[\da-fA-F]{2}))+"
+    urls = re.findall(url_pattern, text)
+    return urls
+
+
+def extract_numbers(text: str) -> list[float]:
+    import re
+
+    numbers = [float(x) for x in re.findall(r"-?\d*\.?\d+", text)]
+    return numbers
+
+
+def json_parse_safe(text: str) -> dict:
+    import json
+
+    try:
+        result = json.loads(text)
+        return result
+    except Exception:
+        return {}
+
+
+def save_to_file(content: str, filename: str):
+    try:
+        with open(filename, "w") as f:
+            f.write(content)
+    except Exception:
+        raise
+
+
+def read_from_file(filename: str) -> str:
+    try:
+        with open(filename) as f:
+            content = f.read()
+        return content
+    except Exception:
+        raise

flock/workflow/activities.py

@@ -0,0 +1,115 @@
+"""Defines Temporal activities for running a chain of agents."""
+
+from datetime import datetime
+
+from temporalio import activity
+
+from flock.core.agent import Agent
+from flock.core.agent_registry import Registry
+from flock.core.context import FlockContext
+from flock.core.context_vars import FLOCK_CURRENT_AGENT, FLOCK_INITIAL_INPUT
+from flock.core.handoff.handoff_base import HandoffBase
+from flock.core.logging import flock_logger
+
+
+@activity.defn
+async def run_agent(context: FlockContext) -> dict:
+    """Runs a chain of agents using the provided context.
+
+    The context contains:
+      - A state (e.g., "init_input", "current_agent", etc.),
+      - A history of agent runs.
+
+    Each agent uses the current value of the variable specified in its `input` attribute.
+    After each run, its output is merged into the context state.
+    The default handoff behavior is to fetch the next agent's input automatically from the context.
+    """
+    registry = Registry()
+    current_agent_name = context.get_variable(FLOCK_CURRENT_AGENT)
+    flock_logger.info(f"Starting agent chain with initial agent", agent=current_agent_name)
+
+    agent = registry.get_agent(current_agent_name)
+    if not agent:
+        flock_logger.error(f"Agent not found", agent=current_agent_name)
+        return {"error": f"Agent '{current_agent_name}' not found."}
+
+    while agent:
+        # Determine the input for this agent.
+        # (Preferably, the agent's input key is used; otherwise "init_input" is assumed.)
+        agent_input = (
+            context.get_variable(agent.input)
+            if getattr(agent, "input", None)
+            else context.get_variable(FLOCK_INITIAL_INPUT)
+        )
+        flock_logger.info(
+            "Determined agent input", agent=agent.name, input_source=getattr(agent, "input", "FLOCK_INITIAL_INPUT")
+        )
+
+        # Prepare a deep copy of the state for the agent's run.
+        local_context = FlockContext.from_dict(context.to_dict())
+        local_context.set_variable(
+            f"{agent.name}.{agent.input}", agent_input
+        )  # ensure the agent's expected input is set
+
+        # Execute the agent.
+        flock_logger.info("Executing agent", agent=agent.name)
+        result = await agent.run(local_context)
+        flock_logger.info("Agent execution completed", agent=agent.name)
+
+        # If no handoff is defined, return the result.
+        if not agent.hand_off:
+            # Record the agent's execution.
+            context.record(agent.name, result, timestamp=datetime.now(), hand_off=None)
+            flock_logger.info("No handoff defined, completing chain", agent=agent.name)
+            return result
+
+        # Determine the next agent.
+        if callable(agent.hand_off):
+            # The handoff function may override the default behavior:
+            # it can explicitly return a dict with "next_agent" and optionally "input"
+            flock_logger.info("Executing handoff function", agent=agent.name)
+            handoff_data: HandoffBase = agent.hand_off(context, result)
+            if isinstance(handoff_data.next_agent, Agent):
+                next_agent_name = handoff_data.next_agent.name
+            else:
+                next_agent_name = handoff_data.next_agent
+
+            # Use the provided new input if present, otherwise let the context update be automatic.
+            if handoff_data.input:
+                context.state["init_input"] = handoff_data["input"]
+                flock_logger.debug("Using handoff-provided input", agent=agent.name)
+            if handoff_data.context_params:
+                context.state.update(handoff_data.context_params)
+                flock_logger.debug("Updated context with handoff params", agent=agent.name)
+        elif isinstance(agent.hand_off, str | Agent):
+            next_agent_name = agent.hand_off if isinstance(agent.hand_off, str) else agent.hand_off.name
+
+        else:
+            return {"error": "Unsupported hand_off type."}
+        context.record(agent.name, result, timestamp=datetime.now(), hand_off=next_agent_name)
+        # Update the current agent and prepare the next input automatically.
+        next_agent = registry.get_agent(next_agent_name)
+        if not next_agent:
+            flock_logger.error("Next agent not found", agent=next_agent_name)
+            return {"error": f"Next agent '{next_agent_name}' not found."}
+
+        context.state["current_agent"] = next_agent.name
+        # By default, the next input is determined from the current state using next_input_for().
+        context.state["init_input"] = context.next_input_for(next_agent)
+        flock_logger.info("Handing off to next agent", current=agent.name, next=next_agent.name)
+        agent = next_agent
+
+    # If the loop ever terminates without a final result, return the latest state value.
+    flock_logger.info("Chain completed, returning final state")
+    return context.get_variable("init_input")
+
+
+@activity.defn
+async def get_next_agent(name: str) -> Agent | None:
+    """Retrieves the agent with the given name from the registry."""
+    flock_logger.debug("Looking up next agent", agent=name)
+    registry = Registry()
+    agent = registry.get_agent(name)
+    if not agent:
+        flock_logger.warning("Next agent not found", agent=name)
+    return agent

flock/workflow/agent_activities.py

@@ -0,0 +1,26 @@
+from temporalio import activity
+
+from flock.agents.declarative_agent import DeclarativeAgent
+from flock.core.context import FlockContext
+
+
+@activity.defn
+async def run_declarative_agent_activity(params: dict) -> dict:
+    """Temporal activity to run a declarative (or batch) agent.
+
+    Expects a dictionary with:
+      - "agent_data": a dict representation of the agent (as produced by .dict()),
+      - "context_data": a dict containing the FlockContext state and optionally other fields.
+
+    The activity reconstructs the agent and a FlockContext, then calls the agent’s _evaluate() method.
+    """
+    agent_data = params.get("agent_data")
+    context_data = params.get("context_data", {})
+    # Reconstruct the agent from its serialized representation.
+    agent = DeclarativeAgent.parse_obj(agent_data)
+    # Reconstruct the FlockContext from the state.
+    state = context_data.get("state", {})
+    # (For simplicity, we ignore history and agent_definitions here. You can extend this if needed.)
+    context = FlockContext(state=state)
+    result = await agent._evaluate(context)
+    return result

flock/workflow/temporal_setup.py

@@ -0,0 +1,37 @@
+import asyncio
+import uuid
+
+from temporalio.client import Client
+from temporalio.worker import Worker
+
+
+async def create_temporal_client() -> Client:
+    client = await Client.connect("localhost:7233")
+    return client
+
+async def setup_worker(workflow,activity) -> Client:
+    worker_client = await create_temporal_client()
+    worker = Worker(worker_client, task_queue="flock-queue", workflows=[workflow], activities=[activity])
+    asyncio.create_task(worker.run())
+    await asyncio.sleep(1)
+
+
+async def run_worker(client: Client, task_queue: str, workflows, activities):
+    worker = Worker(client, task_queue=task_queue, workflows=workflows, activities=activities)
+    await worker.run()
+
+
+async def run_activity(client: Client, name: str, func, param):
+    run_id = f"{name}_{uuid.uuid4().hex[:4]}"
+
+    try:
+        result = await client.execute_activity(
+            func,
+            param,
+            id=run_id,
+            task_queue="flock-queue",
+            start_to_close_timeout=300,  # e.g., 5 minutes
+        )
+        return result
+    except Exception:
+        raise

flock/workflow/workflow.py

@@ -0,0 +1,53 @@
+from datetime import timedelta
+
+from temporalio import workflow
+
+from flock.core.context import FlockContext
+from flock.core.logging import flock_logger
+from flock.workflow.activities import run_agent
+
+# Import activity, passing it through the sandbox without reloading the module
+with workflow.unsafe.imports_passed_through():
+    from flock.workflow.activities import run_agent
+
+
+@workflow.defn
+class FlockWorkflow:
+    def __init__(self) -> None:
+        self.context = None
+
+    @workflow.run
+    async def run(self, context_dict: dict) -> dict:
+        self.context = FlockContext.from_dict(context_dict)
+        self.context.workflow_id = workflow.info().workflow_id
+        self.context.workflow_timestamp = workflow.info().start_time.strftime("%Y-%m-%d %H:%M:%S")
+
+        try:
+            flock_logger.set_context(workflow_id=self.context.workflow_id)
+            flock_logger.workflow_event(f"Starting workflow execution at {self.context.workflow_timestamp}")
+
+            result = await workflow.execute_activity(
+                run_agent, self.context, start_to_close_timeout=timedelta(minutes=5)
+            )
+
+            self.context.set_variable(
+                "flock.result",
+                {
+                    "result": result,
+                    "success": True,
+                },
+            )
+
+            flock_logger.workflow_event("Workflow completed successfully")
+            return result
+
+        except Exception as e:
+            flock_logger.error(f"Workflow execution failed: {e}")
+            self.context.set_variable(
+                "flock.result",
+                {
+                    "result": f"Failed: {e}",
+                    "success": False,
+                },
+            )
+            return self.context