flock-core 0.1.2__py3-none-any.whl → 0.2.1__py3-none-any.whl

flock/agents/user_agent.py

@@ -1,230 +1,145 @@
-from collections.abc import Callable
-from typing import Any
-
-from pydantic import Field
-from temporalio import activity
-
-from flock.core.agent import Agent
-from flock.core.context import FlockContext
-from flock.core.logging import flock_logger, live_update_handler, performance_handler
-
-
-@activity.defn
-async def run_user_agent_activity(context: dict[str, Any]) -> dict[str, Any]:
-    """Temporal activity to process a user agent task.
-
-    Expects context to contain:
-      - "model": the model name
-      - "agent_input": the key used for the input
-      - "output": the agent output specification
-      - "init_input": the initial input
-      - "tools": (optional) list of tools
-    """
-    try:
-        import dspy
-
-        flock_logger.info("Starting user agent activity", context=context)
-
-        with performance_handler.track_time("model_configuration"):
-            model = context.get("model")
-            agent_input = context.get("agent_input")
-            output = context.get("output")
-            init_input = context.get("init_input")
-            tools = context.get("tools")
-
-            flock_logger.debug(
-                "Configuring model",
-                model=model,
-                input=agent_input,
-                output=output,
-            )
-
-            lm = dspy.LM(model)
-            dspy.configure(lm=lm)
-
-        with performance_handler.track_time("task_execution"):
-            if tools:
-                flock_logger.info("Creating ReAct task with tools", num_tools=len(tools))
-                agent_task = dspy.ReAct(f"{agent_input} -> {output}", tools=tools)
-            else:
-                flock_logger.info("Creating Predict task")
-                agent_task = dspy.Predict(f"{agent_input} -> {output}")
-
-            kwargs = {agent_input: init_input}
-            flock_logger.info("Executing task", kwargs=kwargs)
-            result = agent_task(**kwargs).toDict()
-            result[agent_input] = init_input
-
-            flock_logger.success("Task completed successfully")
-            return result
-
-    except Exception as e:
-        flock_logger.error(
-            "User agent activity failed",
-            error=str(e),
-            context=context,
-        )
-        raise
-
-
-class UserAgent(Agent):
-    """An agent that evaluates declarative inputs with user interaction capabilities.
-
-    This agent extends the base Agent class with the ability to interact with users
-    during execution, while maintaining compatibility with both local and Temporal
-    execution modes.
-
-    Attributes:
-        input: Input domain for the agent
-        output: Output types for the agent
-        tools: Tools the agent is allowed to use
-        require_confirmation: Whether to require user confirmation before proceeding
-    """
-
-    input: str = Field(default="", description="Input domain for the agent")
-    output: str = Field(default="", description="Output types for the agent")
-    tools: list[Callable] | None = Field(default=None, description="Tools the agent is allowed to use")
-    require_confirmation: bool = Field(default=False, description="Whether to require user confirmation")
-
-    async def _configure_model(self) -> tuple[Any, Any]:
-        """Configure the model and create the appropriate task."""
-        try:
-            with performance_handler.track_time("model_configuration"):
-                import dspy
-
-                flock_logger.debug(
-                    "Configuring model",
-                    model=self.model,
-                    input=self.input,
-                    output=self.output,
-                )
-
-                lm = dspy.LM(self.model)
-                dspy.configure(lm=lm)
-
-                if self.tools:
-                    flock_logger.info("Creating ReAct task with tools", num_tools=len(self.tools))
-                    agent_task = dspy.ReAct(f"{self.input} -> {self.output}", tools=self.tools)
-                else:
-                    flock_logger.info("Creating Predict task")
-                    agent_task = dspy.Predict(f"{self.input} -> {self.output}")
-
-                return lm, agent_task
-
-        except Exception as e:
-            flock_logger.error(
-                "Model configuration failed",
-                error=str(e),
-                agent=self.name,
-            )
-            raise
-
-    async def _execute_task(self, task: Any, context: FlockContext) -> dict[str, Any]:
-        """Execute the configured task."""
-        try:
-            with performance_handler.track_time("task_execution"):
-                kwargs = {self.input: context.get_variable("init_input")}
-                flock_logger.info("Executing task", kwargs=kwargs)
-
-                with live_update_handler.update_activity_status(
-                    self.name,
-                    "Executing Task",
-                    "Running",
-                    {"input": kwargs},
-                ):
-                    result = task(**kwargs).toDict()
-                    result[self.input] = kwargs[self.input]
-
-                flock_logger.success("Task executed successfully")
-                return result
-
-        except Exception as e:
-            flock_logger.error(
-                "Task execution failed",
-                error=str(e),
-                agent=self.name,
-                kwargs=kwargs,
-            )
-            raise
-
-    async def run(self, context: FlockContext) -> dict[str, Any]:
-        """Run the agent on a task with optional user interaction."""
-        try:
-            flock_logger.info(f"Starting user agent: {self.name}")
-
-            # Configure model and task
-            _, task = await self._configure_model()
-
-            # Execute with user confirmation if required
-            if self.require_confirmation:
-                with live_update_handler.update_activity_status(
-                    self.name,
-                    "Awaiting User Confirmation",
-                    "Paused",
-                    {"agent": self.name},
-                ):
-                    flock_logger.info("Waiting for user confirmation")
-                    # Here you would implement user confirmation logic
-                    # For now, we'll just proceed
-                    pass
-
-            # Execute the task
-            result = await self._execute_task(task, context)
-            return result
-
-        except Exception as e:
-            flock_logger.error(
-                "User agent execution failed",
-                error=str(e),
-                agent=self.name,
-            )
-            raise
-
-    async def run_temporal(self, context: FlockContext) -> dict[str, Any]:
-        """Run the user agent via Temporal."""
-        try:
-            from temporalio.client import Client
-
-            from flock.workflow.temporal_setup import run_activity
-
-            with performance_handler.track_time("temporal_setup"):
-                flock_logger.info(f"Starting temporal user agent: {self.name}")
-                client = await Client.connect("localhost:7233", namespace="default")
-
-                # Prepare context for temporal activity
-                activity_context = {
-                    "model": self.model,
-                    "agent_input": self.input,
-                    "output": self.output,
-                    "init_input": context.get_variable("init_input"),
-                    "tools": self.tools,
-                }
-
-            with live_update_handler.update_workflow_status(
-                self.name,
-                "Running",
-                {"phase": "user_agent_execution"},
-            ):
-                # Execute the activity
-                with performance_handler.track_time("temporal_execution"):
-                    flock_logger.info("Starting user agent activity")
-                    result = await run_activity(
-                        client,
-                        self.name,
-                        run_user_agent_activity,
-                        activity_context,
-                    )
-
-                flock_logger.success(
-                    "Temporal user agent completed successfully",
-                    agent=self.name,
-                )
-                return result
-
-        except Exception as e:
-            flock_logger.error(
-                "Temporal user agent execution failed",
-                error=str(e),
-                agent=self.name,
-            )
-            raise
+from collections.abc import Callable
+from typing import Any
+
+from pydantic import Field
+from temporalio import activity
+
+from flock.core.context.context import FlockContext
+from flock.core.flock_agent import FlockAgent
+
+
+@activity.defn
+async def run_user_agent_activity(context: dict[str, Any]) -> dict[str, Any]:
+    """Temporal activity to process a user agent task.
+
+    Expects context to contain:
+      - "model": the model name
+      - "agent_input": the key used for the input
+      - "output": the agent output specification
+      - "init_input": the initial input
+      - "tools": (optional) list of tools
+    """
+    try:
+        import dspy
+
+        model = context.get("model")
+        agent_input = context.get("agent_input")
+        output = context.get("output")
+        init_input = context.get("init_input")
+        tools = context.get("tools")
+
+        lm = dspy.LM(model)
+        dspy.configure(lm=lm)
+
+        if tools:
+            agent_task = dspy.ReAct(f"{agent_input} -> {output}", tools=tools)
+        else:
+            agent_task = dspy.Predict(f"{agent_input} -> {output}")
+
+        kwargs = {agent_input: init_input}
+        result = agent_task(**kwargs).toDict()
+        result[agent_input] = init_input
+
+        return result
+
+    except Exception:
+        raise
+
+
+class UserAgent(FlockAgent):
+    """An agent that evaluates declarative inputs with user interaction capabilities.
+
+    This agent extends the base Agent class with the ability to interact with users
+    during execution, while maintaining compatibility with both local and Temporal
+    execution modes.
+
+    Attributes:
+        input: Input domain for the agent
+        output: Output types for the agent
+        tools: Tools the agent is allowed to use
+        require_confirmation: Whether to require user confirmation before proceeding
+    """
+
+    input: str = Field(default="", description="Input domain for the agent")
+    output: str = Field(default="", description="Output types for the agent")
+    tools: list[Callable] | None = Field(default=None, description="Tools the agent is allowed to use")
+    require_confirmation: bool = Field(default=False, description="Whether to require user confirmation")
+
+    async def _configure_model(self) -> tuple[Any, Any]:
+        """Configure the model and create the appropriate task."""
+        try:
+            import dspy
+
+            lm = dspy.LM(self.model)
+            dspy.configure(lm=lm)
+
+            if self.tools:
+                agent_task = dspy.ReAct(f"{self.input} -> {self.output}", tools=self.tools)
+            else:
+                agent_task = dspy.Predict(f"{self.input} -> {self.output}")
+
+            return lm, agent_task
+
+        except Exception:
+            raise
+
+    async def _execute_task(self, task: Any, context: FlockContext) -> dict[str, Any]:
+        """Execute the configured task."""
+        try:
+            kwargs = {self.input: context.get_variable("init_input")}
+            result = task(**kwargs).toDict()
+            result[self.input] = kwargs[self.input]
+            return result
+
+        except Exception:
+            raise
+
+    async def run(self, context: FlockContext) -> dict[str, Any]:
+        """Run the agent on a task with optional user interaction."""
+        try:
+            # Configure model and task
+            _, task = await self._configure_model()
+
+            # Execute with user confirmation if required
+            if self.require_confirmation:
+                # Here you would implement user confirmation logic
+                # For now, we'll just proceed
+                pass
+
+            # Execute the task
+            result = await self._execute_task(task, context)
+            return result
+
+        except Exception:
+            raise
+
+    async def run_temporal(self, context: FlockContext) -> dict[str, Any]:
+        """Run the user agent via Temporal."""
+        try:
+            from temporalio.client import Client
+
+            from flock.workflow.temporal_setup import run_activity
+
+            client = await Client.connect("localhost:7233", namespace="default")
+
+            # Prepare context for temporal activity
+            activity_context = {
+                "model": self.model,
+                "agent_input": self.input,
+                "output": self.output,
+                "init_input": context.get_variable("init_input"),
+                "tools": self.tools,
+            }
+
+            # Execute the activity
+            result = await run_activity(
+                client,
+                self.name,
+                run_user_agent_activity,
+                activity_context,
+            )
+
+            return result
+
+        except Exception:
+            raise

flock/core/__init__.py

@@ -1,7 +1 @@
-from flock.core.agent import Agent
-from flock.core.agent_registry import Registry
-from flock.core.context import FlockContext
-from flock.core.flock import Flock
-from flock.core.serializable import Serializable
-
-__all__ = ["Agent", "Flock", "FlockContext", "Registry", "Serializable"]
+

flock/core/context/context.py

@@ -0,0 +1,211 @@
+"""A context object for storing state, history, and agent definitions."""
+
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from typing import Any, Literal
+
+from flock.core.context.context_vars import FLOCK_LAST_AGENT, FLOCK_LAST_RESULT
+from flock.core.util.serializable import Serializable
+
+
+@dataclass
+class AgentRunRecord:
+    """A record of an agent run, including the agent name, data, and timestamp."""
+
+    agent: str = field(default="")  # Agent name
+    data: dict[str, Any] = field(default_factory=dict)
+    timestamp: str = field(default="")
+    hand_off: dict = field(default=dict)  # Next agent name
+    called_from: str = field(default="")  # Previous agent name
+
+
+@dataclass
+class AgentDefinition:
+    """A serializable definition for an agent, including the agent type, name, and data."""
+
+    agent_type: str = field(default="")
+    agent_name: str = field(default="")
+    agent_data: dict = field(default=dict)
+    serializer: Literal["json", "cloudpickle", "msgpack"] = field(
+        default="cloudpickle"
+    )
+
+
+@dataclass
+class FlockContext(Serializable):
+    """A context object for storing state, history, and agent definitions."""
+
+    state: dict[str, Any] = field(default_factory=dict)
+    history: list[AgentRunRecord] = field(default_factory=list)
+    agent_definitions: dict[str, AgentDefinition] = field(default_factory=dict)
+    run_id: str = field(default="")
+    workflow_id: str = field(default="")
+    workflow_timestamp: str = field(default="")
+
+    def record(
+        self,
+        agent_name: str,
+        data: dict[str, Any],
+        timestamp: str,
+        hand_off: str,
+        called_from: str,
+    ) -> None:
+        """Record an agent run and update the state with the agent's output."""
+        try:
+            record = AgentRunRecord(
+                agent=agent_name,
+                data=data.copy(),
+                timestamp=timestamp,
+                hand_off=hand_off,
+                called_from=called_from,
+            )
+            self.history.append(record)
+
+            for key, value in data.items():
+                self.set_variable(f"{agent_name}.{key}", value)
+            self.set_variable(FLOCK_LAST_RESULT, data)
+            self.set_variable(FLOCK_LAST_AGENT, agent_name)
+        except Exception:
+            raise
+
+    def get_variable(self, key: str) -> Any:
+        """Get the current value of a state variable."""
+        try:
+            return self.state.get(key)
+        except Exception:
+            raise
+
+    def set_variable(self, key: str, value: Any) -> None:
+        """Set the value of a state variable."""
+        try:
+            self.state[key] = value
+        except Exception:
+            raise
+
+    def deepcopy(self) -> "FlockContext":
+        """Create a deep copy of the context."""
+        try:
+            return FlockContext.from_dict(self.to_dict())
+        except Exception:
+            raise
+
+    def get_agent_history(self, agent_name: str) -> list[AgentRunRecord]:
+        """Return all agent run records for a given agent name."""
+        try:
+            return [
+                record for record in self.history if record.agent == agent_name
+            ]
+        except Exception:
+            raise
+
+    def next_input_for(self, agent) -> Any:
+        """By default, the next input for an agent is taken from the context state.
+
+        If the agent.input is a comma-separated list (e.g., "input1, input2"),
+        this method will return a dictionary with keys for each of the input names,
+        fetching the latest values from the state.
+
+        If only a single input is specified, the raw value is returned.
+        """
+        try:
+            if hasattr(agent, "input") and isinstance(agent.input, str):
+                keys = [k.strip() for k in agent.input.split(",") if k.strip()]
+
+                if len(keys) == 1:
+                    return self.get_variable(keys[0])
+                else:
+                    return {key: self.get_variable(key) for key in keys}
+            else:
+                # Fallback to "init_input"
+                return self.get_variable("init_input")
+        except Exception:
+            raise
+
+    def get_most_recent_value(self, variable_name: str) -> Any:
+        """Get the definition for a specific agent."""
+        try:
+            for history_record in reversed(self.history):
+                if variable_name in history_record.data:
+                    return history_record.data[variable_name]
+        except Exception:
+            raise
+
+    def get_agent_definition(self, agent_name: str) -> AgentDefinition | None:
+        """Get the definition for a specific agent."""
+        try:
+            for definition in self.agent_definitions:
+                if definition.name == agent_name:
+                    return definition
+            return None
+        except Exception:
+            raise
+
+    def add_agent_definition(
+        self, agent_type: type, agent_name: str, agent_data: Any
+    ) -> None:
+        """Add a new agent definition to the context."""
+        try:
+            definition = AgentDefinition(
+                agent_type=agent_type.__name__,
+                agent_name=agent_name,
+                agent_data=agent_data,
+            )
+            self.agent_definitions[agent_name] = definition
+        except Exception:
+            raise
+
+    # Allow dict-like access for convenience.
+    def __getitem__(self, key: str) -> Any:
+        """Get the current value of a state variable."""
+        value = self.state[key]
+        return value
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        """Set the value of a state variable."""
+        self.state[key] = value
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert the context to a dictionary for serialization."""
+        try:
+
+            def convert(obj):
+                if isinstance(obj, datetime):
+                    return obj.isoformat()
+                if hasattr(obj, "__dataclass_fields__"):  # Is a dataclass
+                    return asdict(
+                        obj,
+                        dict_factory=lambda x: {k: convert(v) for k, v in x},
+                    )
+                return obj
+
+            return convert(asdict(self))
+        except Exception:
+            raise
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "FlockContext":
+        """Create a context instance from a dictionary."""
+        try:
+
+            def convert(obj):
+                if isinstance(obj, dict):
+                    if "timestamp" in obj:  # AgentRunRecord
+                        return AgentRunRecord(
+                            **{
+                                **obj,
+                                "timestamp": datetime.fromisoformat(
+                                    obj["timestamp"]
+                                ),
+                            }
+                        )
+                    if "agent_type" in obj:  # AgentDefinition
+                        return AgentDefinition(**obj)
+                    return {k: convert(v) for k, v in obj.items()}
+                if isinstance(obj, list):
+                    return [convert(v) for v in obj]
+                return obj
+
+            converted = convert(data)
+            return cls(**converted)
+        except Exception:
+            raise

flock/core/context/context_manager.py

@@ -0,0 +1,34 @@
+"""Module for managing the FlockContext."""
+
+from flock.core.context.context import FlockContext
+from flock.core.context.context_vars import (
+    FLOCK_CURRENT_AGENT,
+    FLOCK_INITIAL_INPUT,
+    FLOCK_LOCAL_DEBUG,
+    FLOCK_RUN_ID,
+)
+
+
+def initialize_context(
+    context: FlockContext,
+    agent_name: str,
+    input_data: dict,
+    run_id: str,
+    local_debug: bool,
+) -> None:
+    """Initialize the FlockContext with standard variables before running an agent.
+
+    Args:
+        context: The FlockContext instance.
+        agent_name: The name of the current agent.
+        input_data: A dictionary of inputs for the agent.
+        run_id: A unique identifier for the run.
+        local_debug: Flag indicating whether local debugging is enabled.
+    """
+    context.set_variable(FLOCK_CURRENT_AGENT, agent_name)
+    for key, value in input_data.items():
+        context.set_variable("flock." + key, value)
+    context.set_variable(FLOCK_INITIAL_INPUT, input_data)
+    context.set_variable(FLOCK_LOCAL_DEBUG, local_debug)
+    context.run_id = run_id
+    context.set_variable(FLOCK_RUN_ID, run_id)

flock/core/{context_vars.py → context/context_vars.py}

@@ -1,6 +1,8 @@
-FLOCK_CURRENT_AGENT = "flock.current_agent"
-FLOCK_INITIAL_INPUT = "flock.initial_input"
-FLOCK_LOCAL_DEBUG = "flock.local_debug"
-FLOCK_RUN_ID = "flock.run_id"
-FLOCK_LAST_AGENT = "flock.last_agent"
-FLOCK_LAST_RESULT = "flock.last_result"
+"""Context variables for Flock."""
+
+FLOCK_CURRENT_AGENT = "flock.current_agent"
+FLOCK_INITIAL_INPUT = "flock.initial_input"
+FLOCK_LOCAL_DEBUG = "flock.local_debug"
+FLOCK_RUN_ID = "flock.run_id"
+FLOCK_LAST_AGENT = "flock.last_agent"
+FLOCK_LAST_RESULT = "flock.last_result"