flock-core 0.2.13__py3-none-any.whl → 0.2.15__py3-none-any.whl

flock/core/flock_agent.py

@@ -1,14 +1,21 @@
 """FlockAgent is the core, declarative base class for all agents in the Flock framework."""
 
+import asyncio
+import json
+import os
 from abc import ABC
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
-from typing import Any, Literal, Union
+from typing import Any, TypeVar, Union
 
 import cloudpickle
 from pydantic import BaseModel, Field
 
 from flock.core.context.context import FlockContext
+from flock.core.logging.formatters.themed_formatter import (
+    ThemedAgentResultFormatter,
+)
+from flock.core.logging.formatters.themes import OutputTheme
 from flock.core.logging.logging import get_logger
 from flock.core.mixin.dspy_integration import AgentType, DSPyIntegrationMixin
 from flock.core.mixin.prompt_parser import PromptParserMixin
@@ -21,6 +28,9 @@ from opentelemetry import trace
 tracer = trace.get_tracer(__name__)
 
 
+T = TypeVar("T", bound="FlockAgent")
+
+
 @dataclass
 class FlockAgentConfig:
     """Configuration options for a FlockAgent."""
@@ -34,13 +44,31 @@ class FlockAgentConfig:
     disable_output: bool = field(
         default=False, metadata={"description": "Disables the agent's output."}
     )
-    save_to_file: bool = field(
-        default=False,
-        metadata={
-            "description": "Saves the serialized agent to a file every time it gets serialized."
-        },
+    temperature: float = field(
+        default=0.0, metadata={"description": "Temperature for the LLM"}
+    )
+    max_tokens: int = field(
+        default=2000, metadata={"description": "Max tokens for the LLM"}
+    )
+
+
+@dataclass
+class FlockAgentOutputConfig:
+    """Configuration options for a FlockAgent."""
+
+    render_table: bool = field(
+        default=False, metadata={"description": "Renders a table."}
+    )
+    theme: OutputTheme = field(  # type: ignore
+        default=OutputTheme.afterglow,
+        metadata={"description": "Disables the agent's output."},
+    )
+    max_length: int = field(
+        default=1000, metadata={"description": "Disables the agent's output."}
+    )
+    wait_for_input: bool = field(
+        default=False, metadata={"description": "Wait for input."}
     )
-    data_type: Literal["json", "cloudpickle", "msgpack"] = "cloudpickle"
 
 
 @dataclass
@@ -181,6 +209,11 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
         description="Configuration options for the agent, such as serialization settings.",
     )
 
+    output_config: FlockAgentOutputConfig = Field(
+        default_factory=FlockAgentOutputConfig,
+        description="Configuration options for the agent's output.",
+    )
+
     # Lifecycle callback fields: if provided, these callbacks are used instead of overriding the methods.
     initialize_callback: Callable[[dict[str, Any]], Awaitable[None]] | None = (
         Field(
@@ -188,6 +221,12 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
             description="Optional callback function for initialization. If provided, this async function is called with the inputs.",
         )
     )
+    evaluate_callback: (
+        Callable[[dict[str, Any]], Awaitable[dict[str, Any]]] | None
+    ) = Field(
+        default=None,
+        description="Optional callback function for evaluate. If provided, this async function is called with the inputs instead of the internal evaluate",
+    )
     terminate_callback: (
         Callable[[dict[str, Any], dict[str, Any]], Awaitable[None]] | None
     ) = Field(
@@ -296,6 +335,8 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
         with tracer.start_as_current_span("agent.evaluate") as span:
             span.set_attribute("agent.name", self.name)
             span.set_attribute("inputs", str(inputs))
+            if self.evaluate_callback is not None:
+                return await self.evaluate_callback(self, inputs)
             try:
                 # Create and configure the signature and language model.
                 self.__dspy_signature = self.create_dspy_signature_class(
@@ -323,7 +364,33 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
                 span.record_exception(eval_error)
                 raise
 
-    async def run(self, inputs: dict[str, Any]) -> dict[str, Any]:
+    def save_to_file(self, file_path: str | None = None) -> None:
+        """Save the serialized agent to a file."""
+        if file_path is None:
+            file_path = f"{self.name}.json"
+        dict_data = self.to_dict()
+
+        # create all needed directories
+        path = os.path.dirname(file_path)
+        if path:
+            os.makedirs(os.path.dirname(file_path), exist_ok=True)
+
+        with open(file_path, "w") as file:
+            file.write(json.dumps(dict_data))
+
+    @classmethod
+    def load_from_file(cls: type[T], file_path: str) -> T:
+        """Load a serialized agent from a file."""
+        with open(file_path) as file:
+            data = json.load(file)
+        # Fallback: use the current class.
+        return cls.from_dict(data)
+
+    def run(self, inputs: dict[str, Any]) -> dict[str, Any]:
+        """Run the agent with the given inputs and return its generated output."""
+        return asyncio.run(self.run_async(inputs))
+
+    async def run_async(self, inputs: dict[str, Any]) -> dict[str, Any]:
         """Run the agent with the given inputs and return its generated output.
 
         This method represents the primary execution flow for a FlockAgent and performs the following
@@ -375,6 +442,7 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
             try:
                 await self.initialize(inputs)
                 result = await self.evaluate(inputs)
+                self.display_output(result)
                 await self.terminate(inputs, result)
                 span.set_attribute("result", str(result))
                 logger.info("Agent run completed", agent=self.name)
@@ -387,6 +455,18 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
                 span.record_exception(run_error)
                 raise
 
+    def display_info(self) -> None:
+        pass
+
+    def display_output(self, result: dict[str, Any]) -> None:
+        """Display the agent's output using the configured output formatter."""
+        ThemedAgentResultFormatter(
+            self.output_config.theme,
+            self.output_config.max_length,
+            self.output_config.render_table,
+            self.output_config.wait_for_input,
+        ).display_result(result, self.name)
+
     async def run_temporal(self, inputs: dict[str, Any]) -> dict[str, Any]:
         """Execute this agent via a Temporal workflow for enhanced fault tolerance and asynchronous processing.
 
@@ -526,7 +606,7 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
         return convert_callable(data)
 
     @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> "FlockAgent":
+    def from_dict(cls: type[T], data: dict[str, Any]) -> T:
         """Deserialize a FlockAgent instance from a dictionary.
 
         This class method reconstructs a FlockAgent from its serialized dictionary representation, as produced

flock/core/logging/formatters/enum_builder.py

@@ -0,0 +1,38 @@
+"""Enum Builder."""
+
+import os
+import pathlib
+import re
+
+theme_folder = pathlib.Path(__file__).parent.parent.parent.parent / "themes"
+
+if not theme_folder.exists():
+    raise FileNotFoundError(f"Theme folder not found: {theme_folder}")
+
+theme_files = [
+    pathlib.Path(f.path).stem for f in os.scandir(theme_folder) if f.is_file()
+]
+
+theme_enum_entries = {}
+for theme in theme_files:
+    safe_name = (
+        theme.replace("-", "_")
+        .replace(" ", "_")
+        .replace("(", "_")
+        .replace(")", "_")
+        .replace("+", "_")
+        .replace(".", "_")
+    )
+
+    if re.match(r"^\d", safe_name):
+        safe_name = f"_{safe_name}"
+
+    theme_enum_entries[safe_name] = theme
+
+with open("theme_enum.py", "w") as f:
+    f.write("from enum import Enum\n\n")
+    f.write("class OutputOptionsTheme(Enum):\n")
+    for safe_name, original_name in theme_enum_entries.items():
+        f.write(f'    {safe_name} = "{original_name}"\n')
+
+print("Generated theme_enum.py ✅")

flock/core/logging/formatters/theme_builder.py

@@ -371,7 +371,7 @@ def save_theme(theme: dict, filename: pathlib.Path) -> None:
 # --- Main Interactive Loop --- #
 
 
-def main():
+def theme_builder():
     console = Console(force_terminal=True, color_system="truecolor")
     themes_dir = pathlib.Path(__file__).parent.parent.parent.parent / "themes"
     theme_files = load_theme_files(themes_dir)
@@ -458,7 +458,7 @@ def main():
     )
     if sel2.lower() == "r":
         console.print("Regenerating samples...")
-        main()  # restart the builder
+        theme_builder()  # restart the builder
         return
     try:
         sel2 = int(sel2)
@@ -474,7 +474,3 @@ def main():
     save_path = themes_dir / filename
     save_theme(chosen_sample_theme, save_path)
     console.print(f"\n[green]Theme saved as {save_path}.[/green]")
-
-
-if __name__ == "__main__":
-    main()

flock/core/logging/formatters/themed_formatter.py

@@ -5,16 +5,19 @@ import random
 import re
 from typing import Any
 
-from devtools import pprint
 from temporalio import workflow
 
-from flock.core.logging.formatters.base_formatter import BaseFormatter
+from flock.core.logging.formatters.themes import OutputTheme
 
 with workflow.unsafe.imports_passed_through():
+    from pygments.style import Style
+    from pygments.token import Token
     from rich import box
     from rich.console import Console, Group
     from rich.panel import Panel
+    from rich.syntax import PygmentsSyntaxTheme, Syntax
     from rich.table import Table
+    from rich.theme import Theme
 
 import toml  # install with: pip install toml
 
@@ -335,18 +338,101 @@ def create_rich_renderable(
         return s
 
 
-class ThemedAgentResultFormatter(BaseFormatter):
+def load_syntax_theme_from_file(filepath: str) -> dict:
+    """Load a syntax highlighting theme from a TOML file and map it to Rich styles."""
+    with open(filepath) as f:
+        theme = toml.load(f)
+
+    if "colors" not in theme:
+        raise ValueError(
+            f"Theme file {filepath} does not contain a 'colors' section."
+        )
+
+    # Map theme colors to syntax categories
+    syntax_theme = {
+        "background": theme["colors"]["primary"].get("background", "#161719"),
+        "text": theme["colors"]["primary"].get("foreground", "#c5c8c6"),
+        "comment": theme["colors"]["normal"].get("black", "#666666"),
+        "keyword": theme["colors"]["bright"].get("magenta", "#ff79c6"),
+        "builtin": theme["colors"]["bright"].get("cyan", "#8be9fd"),
+        "string": theme["colors"]["bright"].get("green", "#50fa7b"),
+        "name": theme["colors"]["bright"].get("blue", "#6272a4"),
+        "number": theme["colors"]["bright"].get("yellow", "#f1fa8c"),
+        "operator": theme["colors"]["bright"].get("red", "#ff5555"),
+        "punctuation": theme["colors"]["normal"].get("white", "#bbbbbb"),
+        "error": theme["colors"]["bright"].get("red", "#ff5555"),
+    }
+
+    return syntax_theme
+
+
+def create_rich_syntax_theme(syntax_theme: dict) -> Theme:
+    """Convert a syntax theme dict to a Rich-compatible Theme."""
+    return Theme(
+        {
+            "background": f"on {syntax_theme['background']}",
+            "text": syntax_theme["text"],
+            "keyword": f"bold {syntax_theme['keyword']}",
+            "builtin": f"bold {syntax_theme['builtin']}",
+            "string": syntax_theme["string"],
+            "name": syntax_theme["name"],
+            "number": syntax_theme["number"],
+            "operator": syntax_theme["operator"],
+            "punctuation": syntax_theme["punctuation"],
+            "error": f"bold {syntax_theme['error']}",
+        }
+    )
+
+
+def create_pygments_syntax_theme(syntax_theme: dict) -> PygmentsSyntaxTheme:
+    """Convert a syntax theme dict to a Pygments-compatible Rich syntax theme."""
+
+    class CustomSyntaxStyle(Style):
+        """Dynamically generated Pygments style based on the loaded theme."""
+
+        background_color = syntax_theme["background"]
+        styles = {
+            Token.Text: syntax_theme["text"],
+            Token.Comment: f"italic {syntax_theme['comment']}",
+            Token.Keyword: f"bold {syntax_theme['keyword']}",
+            Token.Name.Builtin: f"bold {syntax_theme['builtin']}",
+            Token.String: syntax_theme["string"],
+            Token.Name: syntax_theme["name"],
+            Token.Number: syntax_theme["number"],
+            Token.Operator: syntax_theme["operator"],
+            Token.Punctuation: syntax_theme["punctuation"],
+            Token.Error: f"bold {syntax_theme['error']}",
+        }
+
+    return PygmentsSyntaxTheme(CustomSyntaxStyle)
+
+
+class ThemedAgentResultFormatter:
     """Formats agent results in a Rich table with nested subtables and theme support."""
 
-    def __init__(self, theme: str = "atom", max_length: int = -1):
+    def __init__(
+        self,
+        theme: OutputTheme = OutputTheme.afterglow,
+        max_length: int = -1,
+        render_table: bool = True,
+        wait_for_input: bool = False,
+    ):
         """Initialize the formatter with a theme and optional max length."""
         self.theme = theme
         self.styles = None
         self.max_length = max_length
+        self.render_table = render_table
+        self.wait_for_input = wait_for_input
 
     def format_result(
-        self, result: dict[str, Any], agent_name: str, theme, styles
+        self,
+        result: dict[str, Any],
+        agent_name: str,
+        theme,
+        styles,
     ) -> Panel:
+        from devtools import pformat
+
         """Format an agent's result as a Rich Panel containing a table."""
         box_style = (
             getattr(box, styles["table_box"])
@@ -394,14 +480,32 @@ class ThemedAgentResultFormatter(BaseFormatter):
             )
             table.add_row(key, rich_renderable)
 
-        return Panel(
-            table,
-            title="🐤🐧🐓🦆",
-            title_align=styles["panel_title_align"],
-            border_style=styles["panel_border_style"],
-            padding=styles["panel_padding"],
-            style=styles["panel_style"],
-        )
+        s = pformat(result, highlight=False)
+
+        if self.render_table:
+            return Panel(
+                table,
+                title="🐤🐧🐓🦆",
+                title_align=styles["panel_title_align"],
+                border_style=styles["panel_border_style"],
+                padding=styles["panel_padding"],
+                style=styles["panel_style"],
+            )
+        else:
+            syntax = Syntax(
+                s,  # The formatted string
+                "python",  # Highlight as Python (change this for other formats)
+                theme=self.syntax_style,  # Choose a Rich theme (matches your color setup)
+                line_numbers=False,
+            )
+            return Panel(
+                syntax,
+                title=agent_name,
+                title_align=styles["panel_title_align"],
+                border_style=styles["panel_border_style"],
+                padding=styles["panel_padding"],
+                style=styles["panel_style"],
+            )
 
     def display_result(self, result: dict[str, Any], agent_name: str) -> None:
         """Print an agent's result using Rich formatting."""
@@ -410,7 +514,11 @@ class ThemedAgentResultFormatter(BaseFormatter):
             pathlib.Path(__file__).parent.parent.parent.parent / "themes"
         )
         all_themes = list(themes_dir.glob("*.toml"))
-        theme = theme + ".toml" if not theme.endswith(".toml") else theme
+        theme = (
+            theme.value + ".toml"
+            if not theme.value.endswith(".toml")
+            else theme.value
+        )
         theme = (
             pathlib.Path(__file__).parent.parent.parent.parent
             / "themes"
@@ -426,6 +534,9 @@ class ThemedAgentResultFormatter(BaseFormatter):
 
         styles = get_default_styles(theme_dict)
         self.styles = styles
+        self.syntax_style = create_pygments_syntax_theme(
+            load_syntax_theme_from_file(theme)
+        )
 
         console = Console()
         panel = self.format_result(
@@ -435,8 +546,5 @@ class ThemedAgentResultFormatter(BaseFormatter):
             styles=styles,
         )
         console.print(panel)
-
-    @staticmethod
-    def display_data(data: dict[str, Any]) -> None:
-        """Print agent data using Rich formatting."""
-        pprint(data)
+        if self.wait_for_input:
+            console.input(prompt="Press Enter to continue...")