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

flock/core/context/context.py

@@ -1,31 +1,31 @@
-"""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 opentelemetry import trace
+
 from flock.core.context.context_vars import FLOCK_LAST_AGENT, FLOCK_LAST_RESULT
+from flock.core.logging.logging import get_logger
 from flock.core.util.serializable import Serializable
 
+logger = get_logger("context")
+tracer = trace.get_tracer(__name__)
+
 
 @dataclass
 class AgentRunRecord:
-    """A record of an agent run, including the agent name, data, and timestamp."""
-
-    agent: str = field(default="")  # Agent name
+    agent: str = field(default="")
     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
+    hand_off: dict = field(default_factory=dict)
+    called_from: str = field(default="")
 
 
 @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)
+    agent_data: dict = field(default_factory=dict)
     serializer: Literal["json", "cloudpickle", "msgpack"] = field(
         default="cloudpickle"
     )
@@ -33,8 +33,6 @@ class AgentDefinition:
 
 @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)
@@ -50,162 +48,135 @@ class FlockContext(Serializable):
         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,
+        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)
+        logger.info(
+            "Agent run recorded",
+            agent=agent_name,
+            timestamp=timestamp,
+            data=data,
+        )
+        current_span = trace.get_current_span()
+        if current_span.get_span_context().is_valid:
+            current_span.add_event(
+                "record",
+                attributes={"agent": agent_name, "timestamp": timestamp},
             )
-            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
+        return self.state.get(key)
 
     def set_variable(self, key: str, value: Any) -> None:
-        """Set the value of a state variable."""
-        try:
-            self.state[key] = value
-        except Exception:
-            raise
+        old_value = self.state.get(key)
+        self.state[key] = value
+        if old_value != value:
+            logger.info(
+                "Context variable updated",
+                variable=key,
+                old=old_value,
+                new=value,
+            )
+            current_span = trace.get_current_span()
+            if current_span.get_span_context().is_valid:
+                current_span.add_event(
+                    "set_variable",
+                    attributes={
+                        "key": key,
+                        "old": str(old_value),
+                        "new": str(value),
+                    },
+                )
 
     def deepcopy(self) -> "FlockContext":
-        """Create a deep copy of the context."""
-        try:
-            return FlockContext.from_dict(self.to_dict())
-        except Exception:
-            raise
+        return FlockContext.from_dict(self.to_dict())
 
     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
+        return [record for record in self.history if record.agent == agent_name]
 
     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:
+        except Exception as e:
+            logger.error(
+                "Error getting next input for agent",
+                agent=agent.name,
+                error=str(e),
+            )
             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
+        for history_record in reversed(self.history):
+            if variable_name in history_record.data:
+                return history_record.data[variable_name]
 
     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
+        return self.agent_definitions.get(agent_name)
 
     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.
+        definition = AgentDefinition(
+            agent_type=agent_type.__name__,
+            agent_name=agent_name,
+            agent_data=agent_data,
+        )
+        self.agent_definitions[agent_name] = definition
+
+    # Use the reactive setter for dict-like access.
     def __getitem__(self, key: str) -> Any:
-        """Get the current value of a state variable."""
-        value = self.state[key]
-        return value
+        return self.get_variable(key)
 
     def __setitem__(self, key: str, value: Any) -> None:
-        """Set the value of a state variable."""
-        self.state[key] = value
+        self.set_variable(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
+        def convert(obj):
+            if isinstance(obj, datetime):
+                return obj.isoformat()
+            if hasattr(obj, "__dataclass_fields__"):
+                return asdict(
+                    obj, dict_factory=lambda x: {k: convert(v) for k, v in x}
+                )
+            return obj
 
-            return convert(asdict(self))
-        except Exception:
-            raise
+        return convert(asdict(self))
 
     @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
+        def convert(obj):
+            if isinstance(obj, dict):
+                if "timestamp" in obj:
+                    return AgentRunRecord(
+                        **{
+                            **obj,
+                            "timestamp": datetime.fromisoformat(
+                                obj["timestamp"]
+                            ),
+                        }
+                    )
+                if "agent_type" in obj:
+                    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)

flock/core/execution/temporal_executor.py

@@ -1,14 +1,24 @@
 # src/your_package/core/execution/temporal_executor.py
+from devtools import pprint
+
 from flock.core.context.context import FlockContext
+from flock.core.context.context_vars import FLOCK_RUN_ID
+from flock.core.logging.formatters.formatter_factory import FormatterFactory
 from flock.core.logging.logging import get_logger
-from flock.workflow.activities import run_agent  # Activity function used in Temporal
+from flock.workflow.activities import (
+    run_agent,  # Activity function used in Temporal
+)
 from flock.workflow.temporal_setup import create_temporal_client, setup_worker
 from flock.workflow.workflow import FlockWorkflow  # Your workflow class
 
 logger = get_logger("flock")
 
 
-async def run_temporal_workflow(context: FlockContext, box_result: bool = True) -> dict:
+async def run_temporal_workflow(
+    context: FlockContext,
+    output_formatter,
+    box_result: bool = True,
+) -> dict:
     """Execute the agent workflow via Temporal for robust, distributed processing.
 
     Args:
@@ -22,16 +32,22 @@ async def run_temporal_workflow(context: FlockContext, box_result: bool = True)
     await setup_worker(workflow=FlockWorkflow, activity=run_agent)
     logger.debug("Creating Temporal client")
     flock_client = await create_temporal_client()
-    workflow_id = context.get_variable("FLOCK_RUN_ID")
+    workflow_id = context.get_variable(FLOCK_RUN_ID)
     logger.info("Executing Temporal workflow", workflow_id=workflow_id)
     result = await flock_client.execute_workflow(
-        FlockWorkflow.run, context.to_dict(), id=workflow_id, task_queue="flock-queue"
+        FlockWorkflow.run,
+        context.to_dict(),
+        id=workflow_id,
+        task_queue="flock-queue",
     )
-    from flock.core.logging.formatters.themed_formatter import ThemedAgentResultFormatter
 
     agent_name = context.get_variable("FLOCK_CURRENT_AGENT")
     logger.debug("Formatting Temporal result", agent=agent_name)
-    ThemedAgentResultFormatter.print_result(result, agent_name)
+    if output_formatter:
+        formatter = FormatterFactory.create_formatter(output_formatter)
+        formatter.display(result, agent_name, output_formatter.wait_for_input)
+    else:
+        pprint(result)
     if box_result:
         from box import Box
 

flock/core/flock.py

@@ -4,6 +4,7 @@ import os
 import uuid
 from typing import TypeVar
 
+from opentelemetry import trace
 from rich.prompt import Prompt
 
 from flock.core.context.context import FlockContext
@@ -14,12 +15,15 @@ from flock.core.flock_agent import FlockAgent
 from flock.core.logging.formatters.base_formatter import FormatterOptions
 from flock.core.logging.formatters.pprint_formatter import PrettyPrintFormatter
 from flock.core.logging.logging import get_logger
+from flock.core.logging.telemetry import setup_tracing
 from flock.core.registry.agent_registry import Registry
 from flock.core.util.cli_helper import display_banner
 from flock.core.util.input_resolver import top_level_to_keys
 
 T = TypeVar("T", bound=FlockAgent)
 logger = get_logger("flock")
+tracer = setup_tracing()
+tracer = trace.get_tracer(__name__)
 
 
 class Flock:
@@ -46,29 +50,36 @@ class Flock:
             enable_logging (bool): If True, enable verbose logging. Defaults to False.
             output_formatter (FormatterOptions): Options for formatting output results.
         """
-        logger.info(
-            "Initializing Flock",
-            model=model,
-            local_debug=local_debug,
-            enable_logging=enable_logging,
-        )
-        logger.enable_logging = enable_logging
+        with tracer.start_as_current_span("flock_init") as span:
+            span.set_attribute("model", model)
+            span.set_attribute("local_debug", local_debug)
+            span.set_attribute("enable_logging", enable_logging)
+            span.set_attribute(
+                "output_formatter", output_formatter.formatter.__name__
+            )
+            logger.info(
+                "Initializing Flock",
+                model=model,
+                local_debug=local_debug,
+                enable_logging=enable_logging,
+            )
+            logger.enable_logging = enable_logging
 
-        display_banner()
+            display_banner()
 
-        self.agents: dict[str, FlockAgent] = {}
-        self.registry = Registry()
-        self.context = FlockContext()
-        self.model = model
-        self.local_debug = local_debug
-        self.output_formatter = output_formatter
+            self.agents: dict[str, FlockAgent] = {}
+            self.registry = Registry()
+            self.context = FlockContext()
+            self.model = model
+            self.local_debug = local_debug
+            self.output_formatter = output_formatter
 
-        if local_debug:
-            os.environ["LOCAL_DEBUG"] = "1"
-            logger.debug("Set LOCAL_DEBUG environment variable")
-        elif "LOCAL_DEBUG" in os.environ:
-            del os.environ["LOCAL_DEBUG"]
-            logger.debug("Removed LOCAL_DEBUG environment variable")
+            if local_debug:
+                os.environ["LOCAL_DEBUG"] = "1"
+                logger.debug("Set LOCAL_DEBUG environment variable")
+            elif "LOCAL_DEBUG" in os.environ:
+                del os.environ["LOCAL_DEBUG"]
+                logger.debug("Removed LOCAL_DEBUG environment variable")
 
     def add_agent(self, agent: T) -> T:
         """Add a new agent to the Flock system.
@@ -151,49 +162,71 @@ class Flock:
             ValueError: If the specified agent is not found in the registry.
             Exception: For any other errors encountered during execution.
         """
-        try:
-            if isinstance(start_agent, str):
-                logger.debug("Looking up agent by name", agent_name=start_agent)
-                start_agent = self.registry.get_agent(start_agent)
-                if not start_agent:
-                    logger.error("Agent not found", agent_name=start_agent)
-                    raise ValueError(
-                        f"Agent '{start_agent}' not found in registry"
-                    )
-            if context:
-                logger.debug("Using provided context")
-                self.context = context
-            if not run_id:
-                run_id = f"{start_agent.name}_{uuid.uuid4().hex[:4]}"
-                logger.debug("Generated run ID", run_id=run_id)
-
-            # TODO - Add a check for required input keys
-            input_keys = top_level_to_keys(start_agent.input)
-            for key in input_keys:
-                if key.startswith("flock."):
-                    key = key[6:]  # Remove the "flock." prefix
-                if key not in input:
-                    input[key] = Prompt.ask(
-                        f"Please enter {key} for {start_agent.name}"
-                    )
-
-            # Initialize the context with standardized variables
-            initialize_context(
-                self.context, start_agent.name, input, run_id, self.local_debug
+        with tracer.start_as_current_span("run_async") as span:
+            span.set_attribute(
+                "start_agent",
+                start_agent.name
+                if hasattr(start_agent, "name")
+                else start_agent,
             )
 
-            logger.info(
-                "Starting agent execution",
-                agent=start_agent.name,
-                local_debug=self.local_debug,
-            )
+            span.set_attribute("input", str(input))
+            span.set_attribute("context", str(context))
+            span.set_attribute("run_id", run_id)
+            span.set_attribute("box_result", box_result)
+
+            try:
+                if isinstance(start_agent, str):
+                    logger.debug(
+                        "Looking up agent by name", agent_name=start_agent
+                    )
+                    start_agent = self.registry.get_agent(start_agent)
+                    if not start_agent:
+                        logger.error("Agent not found", agent_name=start_agent)
+                        raise ValueError(
+                            f"Agent '{start_agent}' not found in registry"
+                        )
+                    start_agent.resolve_callables(context=self.context)
+                if context:
+                    logger.debug("Using provided context")
+                    self.context = context
+                if not run_id:
+                    run_id = f"{start_agent.name}_{uuid.uuid4().hex[:4]}"
+                    logger.debug("Generated run ID", run_id=run_id)
+
+                # TODO - Add a check for required input keys
+                input_keys = top_level_to_keys(start_agent.input)
+                for key in input_keys:
+                    if key.startswith("flock."):
+                        key = key[6:]  # Remove the "flock." prefix
+                    if key not in input:
+                        input[key] = Prompt.ask(
+                            f"Please enter {key} for {start_agent.name}"
+                        )
+
+                # Initialize the context with standardized variables
+                initialize_context(
+                    self.context,
+                    start_agent.name,
+                    input,
+                    run_id,
+                    self.local_debug,
+                )
 
-            if self.local_debug:
-                return await run_local_workflow(
-                    self.context, self.output_formatter, box_result
+                logger.info(
+                    "Starting agent execution",
+                    agent=start_agent.name,
+                    local_debug=self.local_debug,
                 )
-            else:
-                return await run_temporal_workflow(self.context, box_result)
-        except Exception as e:
-            logger.exception("Execution failed", error=str(e))
-            raise
+
+                if self.local_debug:
+                    return await run_local_workflow(
+                        self.context, self.output_formatter, box_result
+                    )
+                else:
+                    return await run_temporal_workflow(
+                        self.context, self.output_formatter, box_result
+                    )
+            except Exception as e:
+                logger.exception("Execution failed", error=str(e))
+                raise