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

flock/__init__.py

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

flock/agents/__init__.py

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

flock/agents/batch_agent.py

@@ -1,175 +1,140 @@
-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
+import asyncio
+import uuid
+from typing import Any
+
+from pydantic import Field
+
+from flock.core.context.context import FlockContext
+from flock.core.flock_agent import FlockAgent
+
+
+class BatchAgent(FlockAgent):
+    """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:
+            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."
+                )
+                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)
+            ]
+
+            # Process batches
+            tasks = []
+            for batch in batches:
+                tasks.append(
+                    self.evaluate(
+                        context, input_overrides={self.iter_input: batch}
+                    )
+                )
+
+            batch_results = await asyncio.gather(*tasks)
+
+            # Aggregate the outputs
+            output_keys = self._parse_keys(self.output)
+            aggregated = {key: [] for key in output_keys}
+            for res in batch_results:
+                for key in output_keys:
+                    aggregated[key].append(res.get(key))
+            aggregated["batch_results"] = batch_results
+            return aggregated
+
+        except Exception:
+            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:
+            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
+            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."
+                )
+                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)
+            ]
+
+            # Process batches
+            tasks = []
+            for batch in batches:
+                # 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,
+                        },
+                    )
+                )
+
+            batch_results = await asyncio.gather(*tasks)
+
+            # Aggregate the outputs
+            output_keys = self._parse_keys(self.output)
+            aggregated = {key: [] for key in output_keys}
+            for res in batch_results:
+                for key in output_keys:
+                    aggregated[key].append(res.get(key))
+            aggregated["batch_results"] = batch_results
+            return aggregated
+
+        except Exception:
+            raise

flock/agents/loop_agent.py

@@ -1,178 +1,117 @@
-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
+from collections.abc import Callable
+from typing import Any
+
+from pydantic import Field
+
+from flock.core.context.context import FlockContext
+from flock.core.flock_agent import FlockAgent
+
+
+class LoopAgent(FlockAgent):
+    """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:
+            # Here you would implement the actual iteration logic
+            # For now, we'll just return a simple result
+            return {"iteration": iteration, "status": "completed"}
+        except Exception:
+            raise
+
+    def _should_continue(self, result: dict[str, Any], iteration: int) -> bool:
+        """Determine if the loop should continue."""
+        if iteration >= self.max_iterations:
+            return False
+
+        if self.termination_condition:
+            should_terminate = self.termination_condition(result)
+            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:
+            results = []
+            iteration = 0
+
+            while True:
+                result = await self._process_iteration(context, iteration)
+                results.append(result)
+
+                # Check termination conditions
+                if not self._should_continue(result, iteration):
+                    break
+
+                iteration += 1
+
+            return {
+                "iterations": iteration + 1,
+                "results": results,
+                "final_result": results[-1] if results else None,
+            }
+
+        except Exception:
+            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
+
+            client = await Client.connect("localhost:7233", namespace="default")
+
+            results = []
+            iteration = 0
+
+            while True:
+                # Process iteration as a temporal activity
+                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()
+
+                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
+
+            return {
+                "iterations": iteration + 1,
+                "results": results,
+                "final_result": results[-1] if results else None,
+            }
+
+        except Exception:
+            raise