flock-core 0.1.1__py3-none-any.whl

flock/__init__.py

@@ -0,0 +1,4 @@
+# Initialize error handling through our custom handler
+from flock.core.logging import error_handler
+
+error_handler.install()

flock/agents/__init__.py

@@ -0,0 +1,3 @@
+from flock.agents.declarative_agent import DeclarativeAgent
+
+__all__ = ["DeclarativeAgent"]

flock/agents/batch_agent.py

@@ -0,0 +1,175 @@
+import asyncio
+import uuid
+from typing import Any
+
+from pydantic import Field
+
+from flock.core.context import FlockContext
+from flock.core.logging import flock_logger, live_update_handler, performance_handler
+
+from .declarative_agent import DeclarativeAgent
+
+
+class BatchAgent(DeclarativeAgent):
+    """A DeclarativeAgent that processes an iterable input in batches.
+
+    Additional Attributes:
+      iter_input: The key in the FlockContext that holds the iterable (a list).
+      batch_site: The number of items per batch.
+
+    For each batch, the agent’s input dictionary is built from the FlockContext with the
+    value for the iter_input key overridden by the current batch. The outputs across batches
+    are then aggregated.
+    """
+
+    iter_input: str = Field(default="", description="Key of the iterable input (must be a list in the FlockContext)")
+    batch_size: int = Field(default=1, description="Batch size (number of items per batch)")
+
+    async def run(self, context: FlockContext) -> dict:
+        """Run the BatchAgent locally by partitioning the iterable and aggregating the results."""
+        try:
+            with performance_handler.track_time("batch_preparation"):
+                flock_logger.info(f"Starting batch processing for agent: {self.name}")
+                iterable = context.get_variable(self.iter_input)
+                if not isinstance(iterable, list):
+                    error_msg = f"Expected a list for key '{self.iter_input}' in context."
+                    flock_logger.error(error_msg)
+                    return {"error": error_msg}
+
+                # Partition the iterable into batches
+                batches: list[list[Any]] = [
+                    iterable[i : i + self.batch_size] for i in range(0, len(iterable), self.batch_size)
+                ]
+                num_batches = len(batches)
+                flock_logger.info(
+                    "Prepared batches",
+                    total_items=len(iterable),
+                    batch_size=self.batch_size,
+                    num_batches=num_batches,
+                )
+
+            # Process batches with progress tracking
+            with live_update_handler.progress_tracker(f"Processing {num_batches} batches") as update_progress:
+                tasks = []
+                for i, batch in enumerate(batches):
+                    flock_logger.debug(f"Creating task for batch {i + 1}/{num_batches}", batch_size=len(batch))
+                    tasks.append(self._evaluate(context, input_overrides={self.iter_input: batch}))
+                    update_progress((i + 1) * 100 / num_batches)
+
+                with performance_handler.track_time("batch_processing"):
+                    batch_results = await asyncio.gather(*tasks)
+                    flock_logger.success(f"Completed processing {num_batches} batches")
+
+            # Aggregate the outputs
+            with performance_handler.track_time("result_aggregation"):
+                flock_logger.info("Aggregating batch results")
+                output_keys = self._parse_keys(self.output)
+                aggregated = {key: [] for key in output_keys}
+                for i, res in enumerate(batch_results):
+                    flock_logger.debug(f"Aggregating results from batch {i + 1}/{num_batches}")
+                    for key in output_keys:
+                        aggregated[key].append(res.get(key))
+                aggregated["batch_results"] = batch_results
+                flock_logger.success("Successfully aggregated all batch results")
+                return aggregated
+
+        except Exception as e:
+            flock_logger.error(
+                "Batch processing failed",
+                error=str(e),
+                agent=self.name,
+                iter_input=self.iter_input,
+            )
+            raise
+
+    async def run_temporal(self, context: FlockContext) -> dict:
+        """Run the BatchAgent via Temporal.
+
+        For each batch, the agent's evaluation is performed as a separate Temporal activity.
+        The results are then aggregated.
+        """
+        try:
+            with performance_handler.track_time("temporal_setup"):
+                flock_logger.info(f"Starting temporal batch processing for agent: {self.name}")
+
+                from temporalio.client import Client
+
+                from flock.workflow.agent_activities import run_declarative_agent_activity
+                from flock.workflow.temporal_setup import run_activity
+
+                # Connect to Temporal
+                flock_logger.info("Connecting to Temporal service...")
+                client = await Client.connect("localhost:7233", namespace="default")
+
+                # Validate and prepare input
+                iterable = context.get_variable(self.iter_input)
+                if not isinstance(iterable, list):
+                    error_msg = f"Expected a list for key '{self.iter_input}' in context."
+                    flock_logger.error(error_msg)
+                    return {"error": error_msg}
+
+                # Partition into batches
+                batches: list[list[Any]] = [
+                    iterable[i : i + self.batch_size] for i in range(0, len(iterable), self.batch_size)
+                ]
+                num_batches = len(batches)
+                flock_logger.info(
+                    "Prepared batches for temporal processing",
+                    total_items=len(iterable),
+                    batch_size=self.batch_size,
+                    num_batches=num_batches,
+                )
+
+            # Process batches with status updates
+            with live_update_handler.update_workflow_status(
+                self.name, "Running", {"phase": "batch_processing", "total_batches": num_batches}
+            ):
+                tasks = []
+                for i, batch in enumerate(batches):
+                    flock_logger.debug(f"Creating temporal task for batch {i + 1}/{num_batches}", batch_size=len(batch))
+                    # Prepare context for this batch
+                    new_state = context.state.copy()
+                    new_state[self.iter_input] = batch
+                    context_data = {
+                        "state": new_state,
+                        "history": [],  # you might choose to pass along history if needed
+                        "agent_definitions": [],
+                    }
+                    agent_data = self.dict()
+                    task_id = f"{self.name}_{uuid.uuid4().hex[:4]}"
+
+                    # Create temporal activity task
+                    tasks.append(
+                        run_activity(
+                            client,
+                            task_id,
+                            run_declarative_agent_activity,
+                            {"agent_data": agent_data, "context_data": context_data},
+                        )
+                    )
+
+                with performance_handler.track_time("temporal_batch_processing"):
+                    batch_results = await asyncio.gather(*tasks)
+                    flock_logger.success(f"Completed temporal processing of {num_batches} batches")
+
+            # Aggregate the outputs
+            with performance_handler.track_time("temporal_result_aggregation"):
+                flock_logger.info("Aggregating temporal batch results")
+                output_keys = self._parse_keys(self.output)
+                aggregated = {key: [] for key in output_keys}
+                for i, res in enumerate(batch_results):
+                    flock_logger.debug(f"Aggregating results from temporal batch {i + 1}/{num_batches}")
+                    for key in output_keys:
+                        aggregated[key].append(res.get(key))
+                aggregated["batch_results"] = batch_results
+                flock_logger.success("Successfully aggregated all temporal batch results")
+                return aggregated
+
+        except Exception as e:
+            flock_logger.error(
+                "Temporal batch processing failed",
+                error=str(e),
+                agent=self.name,
+                iter_input=self.iter_input,
+            )
+            raise

flock/agents/declarative_agent.py

@@ -0,0 +1,166 @@
+from collections.abc import Callable
+from typing import Any
+
+import dspy  # your dspy package for LM, Predict, ReAct, etc.
+from pydantic import Field
+
+from flock.core.agent import Agent
+from flock.core.context import FlockContext
+from flock.core.logging import flock_logger, performance_handler
+
+
+class DeclarativeAgent(Agent):
+    """An agent that evaluates declarative inputs.
+
+    Attributes:
+      input: A comma‐separated list of input keys.
+             If a key is not found in the incoming FlockContext and only one input is expected,
+             then the value of "init_input" is used.
+      output: A comma‐separated list of output keys (any type annotations after ":" are ignored).
+      tools: An optional list of callables (tools) that the agent is allowed to use.
+      use_cache: Whether to cache agent results.
+    """
+
+    input: str = Field(description="Comma-separated input keys (e.g., 'blog_idea' or 'url, context')")
+    output: str = Field(description="Comma-separated output keys (e.g., 'title, headers')")
+    tools: list[Callable[..., Any]] | None = Field(default=None, description="Tools the agent is allowed to use")
+    use_cache: bool = Field(default=False, description="Whether to use the cache for this agent")
+
+    def _parse_keys(self, keys_str: str) -> list[str]:
+        """Split a comma‐separated string and strip any type annotations.
+        For example, "a, b: list[str]" becomes ["a", "b"].
+        """
+        keys = []
+        for part in keys_str.split(","):
+            part = part.strip()
+            if not part:
+                continue
+            # Remove any type annotation (everything after a colon)
+            if ":" in part:
+                key = part.split(":", 1)[0].strip()
+            else:
+                key = part
+            keys.append(key)
+        return keys
+
+    def _build_input(self, context: FlockContext) -> dict:
+        """Build the dictionary of inputs for the agent based on its input specification.
+
+        For each key in the agent's input string:
+          - If the key is "context" (case-insensitive), pass the entire FlockContext.
+          - Otherwise, use context.get_variable() to fetch the value.
+          - If the key is not found and only one key is expected, default to context["init_input"].
+        """
+        input_keys = self._parse_keys(self.input)
+        inputs = {}
+        for key in input_keys:
+            if key.lower() == "context":
+                inputs[key] = context
+            else:
+                value = context.get_variable(key)
+                if value is None and len(input_keys) == 1:
+                    value = context.get_variable("init_input")
+                inputs[key] = value
+        return inputs
+
+    def _configure_task(self):
+        """Configure the dspy language model and choose a task constructor.
+        If tools are provided, ReAct is used; otherwise, Predict.
+        """
+        with performance_handler.track_time("model_configuration"):
+            flock_logger.debug(f"Configuring {self.model} for {'ReAct' if self.tools else 'Predict'} task")
+            lm = dspy.LM(self.model)
+            dspy.configure(lm=lm)
+            return dspy.ReAct if self.tools else dspy.Predict
+
+    async def _evaluate(self, context: FlockContext, input_overrides: dict | None = None) -> dict:
+        """Evaluate the agent by:
+        1. Building an input dictionary from the FlockContext (merging any overrides),
+        2. Building a prompt like "input1, input2 -> output1, output2",
+        3. Instantiating and executing the dspy task (Predict or ReAct),
+        4. Returning the resulting dictionary.
+        """
+        try:
+            # Build inputs
+            with performance_handler.track_time("input_preparation"):
+                inputs = self._build_input(context)
+                if input_overrides:
+                    inputs.update(input_overrides)
+                input_keys = self._parse_keys(self.input)
+                output_keys = self._parse_keys(self.output)
+                prompt = f"{', '.join(input_keys)} -> {', '.join(output_keys)}"
+                flock_logger.debug("Prepared inputs", inputs=inputs, prompt=prompt)
+
+            # Configure and execute task
+            with performance_handler.track_time("task_execution"):
+                task_constructor = self._configure_task()
+                if self.tools:
+                    flock_logger.info("Creating ReAct task with tools", num_tools=len(self.tools))
+                    agent_task = task_constructor(prompt, tools=self.tools)
+                else:
+                    flock_logger.info("Creating Predict task")
+                    agent_task = task_constructor(prompt)
+
+                flock_logger.info("Executing task...")
+                result = agent_task(**inputs).toDict()
+
+            # Process result
+            for key in input_keys:
+                result.setdefault(key, inputs.get(key))
+            flock_logger.success("Task completed successfully", output_keys=list(result.keys()))
+            return result
+
+        except Exception as e:
+            flock_logger.error(
+                "Task execution failed",
+                error=str(e),
+                agent=self.name,
+                inputs=inputs,
+                prompt=prompt,
+            )
+            raise
+
+    async def run(self, context: FlockContext) -> dict:
+        """Run the agent on the provided FlockContext (locally)."""
+        flock_logger.info(f"Running agent locally: {self.name}")
+        return await self._evaluate(context)
+
+    async def run_temporal(self, context: FlockContext) -> dict:
+        """Run the agent via Temporal by serializing its parameters and the FlockContext and
+        calling a dedicated Temporal activity.
+        """
+        from temporalio.client import Client
+
+        from flock.workflow.agent_activities import run_declarative_agent_activity
+        from flock.workflow.temporal_setup import run_activity
+
+        try:
+            with performance_handler.track_time("temporal_setup"):
+                # Connect to Temporal (adjust the host/namespace as needed)
+                flock_logger.info("Connecting to Temporal...")
+                client = await Client.connect("localhost:7233", namespace="default")
+
+                # Serialize the FlockContext and agent
+                flock_logger.debug("Serializing context and agent data")
+                context_data = {
+                    "state": context.state,
+                    "history": [record.__dict__ for record in context.history],
+                    "agent_definitions": [definition.__dict__ for definition in context.agent_definitions],
+                }
+                agent_data = self.dict()
+
+            # Run the activity
+            with performance_handler.track_time("temporal_activity"):
+                flock_logger.info(f"Starting Temporal activity: {self.name}")
+                result = await run_activity(
+                    client,
+                    self.name,
+                    run_declarative_agent_activity,
+                    {"agent_data": agent_data, "context_data": context_data},
+                )
+                flock_logger.success("Temporal activity completed successfully")
+                return result
+
+        except Exception as e:
+            flock_logger.error(f"Temporal execution failed: {e}", agent=self.name)
+            raise

flock/agents/loop_agent.py

@@ -0,0 +1,178 @@
+from collections.abc import Callable
+from typing import Any
+
+from pydantic import Field
+
+from flock.core.agent import Agent
+from flock.core.context import FlockContext
+from flock.core.logging import flock_logger, live_update_handler, performance_handler
+
+
+class LoopAgent(Agent):
+    """An agent that executes its logic in a loop until a termination condition is met.
+
+    Attributes:
+        input: Input domain for the agent
+        output: Output types for the agent
+        tools: Tools the agent is allowed to use
+        max_iterations: Maximum number of iterations before forced termination
+        termination_condition: Optional callable that determines when to stop the loop
+    """
+
+    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")
+    max_iterations: int = Field(default=10, description="Maximum number of iterations")
+    termination_condition: Callable[[dict[str, Any]], bool] | None = Field(
+        default=None, description="Optional function to determine loop termination"
+    )
+
+    async def _process_iteration(self, context: FlockContext, iteration: int) -> dict[str, Any]:
+        """Process a single iteration of the loop."""
+        try:
+            with performance_handler.track_time(f"iteration_{iteration}"):
+                flock_logger.debug(f"Processing iteration {iteration}", agent=self.name)
+                # Here you would implement the actual iteration logic
+                # For now, we'll just return a simple result
+                return {"iteration": iteration, "status": "completed"}
+        except Exception as e:
+            flock_logger.error(
+                f"Error in iteration {iteration}",
+                error=str(e),
+                agent=self.name,
+            )
+            raise
+
+    def _should_continue(self, result: dict[str, Any], iteration: int) -> bool:
+        """Determine if the loop should continue."""
+        if iteration >= self.max_iterations:
+            flock_logger.warning(
+                "Maximum iterations reached",
+                max_iterations=self.max_iterations,
+                agent=self.name,
+            )
+            return False
+
+        if self.termination_condition:
+            should_terminate = self.termination_condition(result)
+            if should_terminate:
+                flock_logger.info(
+                    "Termination condition met",
+                    iteration=iteration,
+                    agent=self.name,
+                )
+            return not should_terminate
+
+        return True
+
+    async def run(self, context: FlockContext) -> dict[str, Any]:
+        """Run the agent in a loop until the termination condition is met or max iterations reached."""
+        try:
+            flock_logger.info(f"Starting loop agent: {self.name}")
+            results = []
+            iteration = 0
+
+            with live_update_handler.progress_tracker("Loop Progress") as update_progress:
+                while True:
+                    # Update progress based on iteration count
+                    progress = min((iteration + 1) * 100 / self.max_iterations, 100)
+                    update_progress(progress)
+
+                    # Process iteration with status tracking
+                    with live_update_handler.update_activity_status(
+                        f"{self.name}_iteration_{iteration}",
+                        f"Iteration {iteration + 1}",
+                        "Running",
+                        {"max_iterations": self.max_iterations},
+                    ):
+                        result = await self._process_iteration(context, iteration)
+                        results.append(result)
+
+                    # Check termination conditions
+                    if not self._should_continue(result, iteration):
+                        break
+
+                    iteration += 1
+
+            flock_logger.success(
+                f"Loop completed successfully",
+                total_iterations=iteration + 1,
+                agent=self.name,
+            )
+            return {
+                "iterations": iteration + 1,
+                "results": results,
+                "final_result": results[-1] if results else None,
+            }
+
+        except Exception as e:
+            flock_logger.error(
+                "Loop execution failed",
+                error=str(e),
+                agent=self.name,
+                iteration=iteration,
+            )
+            raise
+
+    async def run_temporal(self, context: FlockContext) -> dict[str, Any]:
+        """Run the loop agent via Temporal."""
+        try:
+            from temporalio.client import Client
+
+            from flock.workflow.agent_activities import run_agent_activity
+            from flock.workflow.temporal_setup import run_activity
+
+            with performance_handler.track_time("temporal_setup"):
+                flock_logger.info(f"Starting temporal loop agent: {self.name}")
+                client = await Client.connect("localhost:7233", namespace="default")
+
+            results = []
+            iteration = 0
+
+            with live_update_handler.update_workflow_status(
+                self.name, "Running", {"phase": "loop_execution", "max_iterations": self.max_iterations}
+            ):
+                while True:
+                    # Process iteration as a temporal activity
+                    with performance_handler.track_time(f"temporal_iteration_{iteration}"):
+                        context_data = {
+                            "state": context.state,
+                            "history": [record.__dict__ for record in context.history],
+                            "agent_definitions": [definition.__dict__ for definition in context.agent_definitions],
+                        }
+                        agent_data = self.dict()
+
+                        flock_logger.info(f"Starting temporal iteration {iteration + 1}")
+                        result = await run_activity(
+                            client,
+                            f"{self.name}_iteration_{iteration}",
+                            run_agent_activity,
+                            {"agent_data": agent_data, "context_data": context_data},
+                        )
+                        results.append(result)
+
+                    # Check termination conditions
+                    if not self._should_continue(result, iteration):
+                        break
+
+                    iteration += 1
+
+            flock_logger.success(
+                "Temporal loop completed successfully",
+                total_iterations=iteration + 1,
+                agent=self.name,
+            )
+            return {
+                "iterations": iteration + 1,
+                "results": results,
+                "final_result": results[-1] if results else None,
+            }
+
+        except Exception as e:
+            flock_logger.error(
+                "Temporal loop execution failed",
+                error=str(e),
+                agent=self.name,
+                iteration=iteration,
+            )
+            raise