flock-core 0.3.15__tar.gz → 0.3.16__tar.gz

{flock_core-0.3.15 → flock_core-0.3.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flock-core
-Version: 0.3.15
+Version: 0.3.16
 Summary: Declarative LLM Orchestration at Scale
 Author-email: Andre Ratzenberger <andre.ratzenberger@whiteduck.de>
 License-File: LICENSE

flock_core-0.3.16/examples/02_concepts/context/context.py

@@ -0,0 +1,57 @@
+
+from flock.core import Flock, FlockFactory
+from flock.core.flock_module import FlockModule, FlockModuleConfig
+from flock.routers.default.default_router import DefaultRouter, DefaultRouterConfig
+
+class ContextModule(FlockModule):
+    def terminate(self, agent, inputs, result, context=None):
+        context.set_variable("flock_agent1.a_random_name", "John Doe")
+     
+
+flock = Flock()
+
+
+
+# Flock has an advanced context system that allows you to store and retrieve data
+# across agents. The context is a dictionary that can be accessed and modified by 
+# all agents in the flock via modules
+# Agents will write their inputs and outputs to the context prefixed with "agentname."
+
+
+
+flock_agent_1 = FlockFactory.create_default_agent(
+    name="flock_agent1",
+    input="",
+    output="a_random_name: str",
+)
+
+
+flock_agent_2 = FlockFactory.create_default_agent(
+    name="flock_agent2",
+    input="a_random_name",
+    output="name_in_caps: str",
+)
+flock_agent_2.add_module(ContextModule(name="context_module", config= FlockModuleConfig()))
+
+# Agent3 will reverse John Doe to eoD nhoJ
+flock_agent_3 = FlockFactory.create_default_agent(
+    name="flock_agent3",
+    input="flock_agent1.a_random_name",
+    output="name_reversed: str",
+    wait_for_input=True,
+    print_context=True,
+)
+
+
+flock_agent_1.handoff_router = DefaultRouter(config=DefaultRouterConfig(hand_off=flock_agent_2.name))
+flock_agent_2.handoff_router = DefaultRouter(config=DefaultRouterConfig(hand_off=flock_agent_3.name))
+flock_agent_3.handoff_router = DefaultRouter(config=DefaultRouterConfig(hand_off=flock_agent_1.name))
+flock.add_agent(flock_agent_1)
+flock.add_agent(flock_agent_2)
+flock.add_agent(flock_agent_3)
+
+flock.run(
+    input={},
+    start_agent=flock_agent_1,
+    agents=[flock_agent_1, flock_agent_2, flock_agent_3]
+)

{flock_core-0.3.15 → flock_core-0.3.16}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "flock-core"
-version = "0.3.15"
+version = "0.3.16"
 description = "Declarative LLM Orchestration at Scale"
 readme = "README.md"
 authors = [

{flock_core-0.3.15 → flock_core-0.3.16}/src/flock/core/context/context.py

@@ -1,9 +1,10 @@
 import uuid
-from dataclasses import asdict, dataclass, field
+from dataclasses import asdict
 from datetime import datetime
 from typing import Any, Literal
 
 from opentelemetry import trace
+from pydantic import BaseModel, Field
 
 from flock.core.context.context_vars import FLOCK_LAST_AGENT, FLOCK_LAST_RESULT
 from flock.core.logging.logging import get_logger
@@ -13,34 +14,31 @@ logger = get_logger("context")
 tracer = trace.get_tracer(__name__)
 
 
-@dataclass
-class AgentRunRecord:
-    id: str = field(default="")
-    agent: str = field(default="")
-    data: dict[str, Any] = field(default_factory=dict)
-    timestamp: str = field(default="")
-    hand_off: dict = field(default_factory=dict)
-    called_from: str = field(default="")
+class AgentRunRecord(BaseModel):
+    id: str = Field(default="")
+    agent: str = Field(default="")
+    data: dict[str, Any] = Field(default_factory=dict)
+    timestamp: str = Field(default="")
+    hand_off: dict | None = Field(default_factory=dict)
+    called_from: str = Field(default="")
 
 
-@dataclass
-class AgentDefinition:
-    agent_type: str = field(default="")
-    agent_name: str = field(default="")
-    agent_data: dict = field(default_factory=dict)
-    serializer: Literal["json", "cloudpickle", "msgpack"] = field(
+class AgentDefinition(BaseModel):
+    agent_type: str = Field(default="")
+    agent_name: str = Field(default="")
+    agent_data: dict = Field(default_factory=dict)
+    serializer: Literal["json", "cloudpickle", "msgpack"] = Field(
         default="cloudpickle"
     )
 
 
-@dataclass
-class FlockContext(Serializable):
-    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="")
+class FlockContext(Serializable, BaseModel):
+    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,

{flock_core-0.3.15 → flock_core-0.3.16}/src/flock/core/flock_agent.py

@@ -11,6 +11,7 @@ import cloudpickle
 from opentelemetry import trace
 from pydantic import BaseModel, Field
 
+from flock.core.context.context import FlockContext
 from flock.core.flock_evaluator import FlockEvaluator
 from flock.core.flock_module import FlockModule
 from flock.core.flock_router import FlockRouter
@@ -73,6 +74,11 @@ class FlockAgent(BaseModel, ABC, DSPyIntegrationMixin):
         description="FlockModules attached to this agent",
     )
 
+    context: FlockContext | None = Field(
+        default=None,
+        description="Context associated with flock",
+    )
+
     def add_module(self, module: FlockModule) -> None:
         """Add a module to this agent."""
         self.modules[module.name] = module
@@ -95,6 +101,13 @@ class FlockAgent(BaseModel, ABC, DSPyIntegrationMixin):
         with tracer.start_as_current_span("agent.initialize") as span:
             span.set_attribute("agent.name", self.name)
             span.set_attribute("inputs", str(inputs))
+            if not self.context:
+                self.context = FlockContext()
+
+            if self.name not in self.context.agent_definitions:
+                self.context.add_agent_definition(
+                    type(self), self.name, self.to_dict()
+                )
 
             try:
                 for module in self.get_enabled_modules():
@@ -102,7 +115,7 @@ class FlockAgent(BaseModel, ABC, DSPyIntegrationMixin):
                         f"agent.initialize - module {module.name}",
                         agent=self.name,
                     )
-                    await module.initialize(self, inputs)
+                    await module.initialize(self, inputs, self.context)
             except Exception as module_error:
                 logger.error(
                     "Error during initialize",
@@ -124,7 +137,7 @@ class FlockAgent(BaseModel, ABC, DSPyIntegrationMixin):
             )
             try:
                 for module in self.get_enabled_modules():
-                    await module.terminate(self, inputs, inputs)
+                    await module.terminate(self, inputs, inputs, self.context)
             except Exception as module_error:
                 logger.error(
                     "Error during terminate",
@@ -139,7 +152,7 @@ class FlockAgent(BaseModel, ABC, DSPyIntegrationMixin):
             span.set_attribute("inputs", str(inputs))
             try:
                 for module in self.get_enabled_modules():
-                    await module.on_error(self, error, inputs)
+                    await module.on_error(self, error, inputs, self.context)
             except Exception as module_error:
                 logger.error(
                     "Error during on_error",
@@ -154,13 +167,15 @@ class FlockAgent(BaseModel, ABC, DSPyIntegrationMixin):
             span.set_attribute("inputs", str(inputs))
 
             for module in self.get_enabled_modules():
-                inputs = await module.pre_evaluate(self, inputs)
+                inputs = await module.pre_evaluate(self, inputs, self.context)
 
             try:
                 result = await self.evaluator.evaluate(self, inputs, self.tools)
 
                 for module in self.get_enabled_modules():
-                    result = await module.post_evaluate(self, inputs, result)
+                    result = await module.post_evaluate(
+                        self, inputs, result, self.context
+                    )
 
                 span.set_attribute("result", str(result))
 

{flock_core-0.3.15 → flock_core-0.3.16}/src/flock/core/flock_factory.py

@@ -35,6 +35,7 @@ class FlockFactory:
         max_tokens: int = 4096,
         alert_latency_threshold_ms: int = 30000,
         no_output: bool = False,
+        print_context: bool = False,
     ) -> FlockAgent:
         """Creates a default FlockAgent.
 
@@ -65,6 +66,7 @@ class FlockFactory:
             theme=output_theme,
             wait_for_input=wait_for_input,
             no_output=no_output,
+            print_context=print_context,
         )
         output_module = OutputModule("output", config=output_config)
 

{flock_core-0.3.15 → flock_core-0.3.16}/src/flock/core/flock_module.py

@@ -5,6 +5,8 @@ from typing import Any, TypeVar
 
 from pydantic import BaseModel, Field, create_model
 
+from flock.core.context.context import FlockContext
+
 T = TypeVar("T", bound="FlockModuleConfig")
 
 
@@ -50,30 +52,50 @@ class FlockModule(BaseModel, ABC):
         default_factory=FlockModuleConfig, description="Module configuration"
     )
 
-    async def initialize(self, agent: Any, inputs: dict[str, Any]) -> None:
+    async def initialize(
+        self,
+        agent: Any,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
+    ) -> None:
         """Called when the agent starts running."""
         pass
 
     async def pre_evaluate(
-        self, agent: Any, inputs: dict[str, Any]
+        self,
+        agent: Any,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> dict[str, Any]:
         """Called before agent evaluation, can modify inputs."""
         return inputs
 
     async def post_evaluate(
-        self, agent: Any, inputs: dict[str, Any], result: dict[str, Any]
+        self,
+        agent: Any,
+        inputs: dict[str, Any],
+        result: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> dict[str, Any]:
         """Called after agent evaluation, can modify results."""
         return result
 
     async def terminate(
-        self, agent: Any, inputs: dict[str, Any], result: dict[str, Any]
+        self,
+        agent: Any,
+        inputs: dict[str, Any],
+        result: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> None:
         """Called when the agent finishes running."""
         pass
 
     async def on_error(
-        self, agent: Any, error: Exception, inputs: dict[str, Any]
+        self,
+        agent: Any,
+        error: Exception,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> None:
         """Called when an error occurs during agent execution."""
         pass

{flock_core-0.3.15 → flock_core-0.3.16}/src/flock/modules/callback/callback_module.py

@@ -6,6 +6,7 @@ from typing import Any
 from pydantic import Field
 
 from flock.core import FlockModule, FlockModuleConfig
+from flock.core.context.context import FlockContext
 
 
 class CallbackModuleConfig(FlockModuleConfig):
@@ -44,13 +45,21 @@ class CallbackModule(FlockModule):
         description="Callback module configuration",
     )
 
-    async def pre_initialize(self, agent: Any, inputs: dict[str, Any]) -> None:
+    async def pre_initialize(
+        self,
+        agent: Any,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
+    ) -> None:
         """Run initialize callback if configured."""
         if self.config.initialize_callback:
             await self.config.initialize_callback(agent, inputs)
 
     async def pre_evaluate(
-        self, agent: Any, inputs: dict[str, Any]
+        self,
+        agent: Any,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> dict[str, Any]:
         """Run evaluate callback if configured."""
         if self.config.evaluate_callback:
@@ -58,14 +67,22 @@ class CallbackModule(FlockModule):
         return inputs
 
     async def pre_terminate(
-        self, agent: Any, inputs: dict[str, Any], result: dict[str, Any]
+        self,
+        agent: Any,
+        inputs: dict[str, Any],
+        result: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> None:
         """Run terminate callback if configured."""
         if self.config.terminate_callback:
             await self.config.terminate_callback(agent, inputs, result)
 
     async def on_error(
-        self, agent: Any, error: Exception, inputs: dict[str, Any]
+        self,
+        agent: Any,
+        error: Exception,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> None:
         """Run error callback if configured."""
         if self.config.on_error_callback:

{flock_core-0.3.15 → flock_core-0.3.16}/src/flock/modules/memory/memory_module.py

@@ -7,6 +7,7 @@ from pydantic import Field
 from tqdm import tqdm
 
 from flock.core import FlockAgent, FlockModule, FlockModuleConfig
+from flock.core.context.context import FlockContext
 from flock.core.logging.logging import get_logger
 from flock.modules.memory.memory_parser import MemoryMappingParser
 from flock.modules.memory.memory_storage import FlockMemoryStore, MemoryEntry
@@ -77,7 +78,10 @@ class MemoryModule(FlockModule):
         )
 
     async def initialize(
-        self, agent: FlockAgent, inputs: dict[str, Any]
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> None:
         """Initialize memory store if needed."""
         if not self.memory_store:
@@ -92,7 +96,10 @@ class MemoryModule(FlockModule):
         logger.debug(f"Initialized memory module for agent {agent.name}")
 
     async def pre_evaluate(
-        self, agent: FlockAgent, inputs: dict[str, Any]
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> dict[str, Any]:
         """Check memory before evaluation."""
         if not self.memory_store:
@@ -214,7 +221,11 @@ class MemoryModule(FlockModule):
             logger.warning(f"Memory storage failed: {e}", agent=agent.name)
 
     async def post_evaluate(
-        self, agent: FlockAgent, inputs: dict[str, Any], result: dict[str, Any]
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        result: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> dict[str, Any]:
         """Store results in memory after evaluation."""
         if not self.memory_store:
@@ -229,7 +240,11 @@ class MemoryModule(FlockModule):
         return result
 
     async def terminate(
-        self, agent: Any, inputs: dict[str, Any], result: dict[str, Any]
+        self,
+        agent: Any,
+        inputs: dict[str, Any],
+        result: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> None:
         """Save memory store if configured."""
         if self.config.save_after_update and self.memory_store:

{flock_core-0.3.15 → flock_core-0.3.16}/src/flock/modules/output/output_module.py

@@ -8,6 +8,7 @@ from typing import Any
 from pydantic import Field
 
 from flock.core import FlockAgent
+from flock.core.context.context import FlockContext
 from flock.core.flock_module import FlockModule, FlockModuleConfig
 from flock.core.logging.formatters.themed_formatter import (
     ThemedAgentResultFormatter,
@@ -58,6 +59,10 @@ class OutputModuleConfig(FlockModuleConfig):
         default=False,
         description="Whether to suppress output",
     )
+    print_context: bool = Field(
+        default=False,
+        description="Whether to print the context",
+    )
 
 
 class OutputModule(FlockModule):
@@ -159,12 +164,18 @@ class OutputModule(FlockModule):
             json.dump(output_data, f, indent=2)
 
     async def post_evaluate(
-        self, agent: FlockAgent, inputs: dict[str, Any], result: dict[str, Any]
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        result: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> dict[str, Any]:
         """Format and display the output."""
         logger.debug("Formatting and displaying output")
         if self.config.no_output:
             return result
+        if self.config.print_context:
+            result["context"] = context
         # Display the result using the formatter
         self._formatter.display_result(result, agent.name)
 

{flock_core-0.3.15 → flock_core-0.3.16}/src/flock/modules/performance/metrics_module.py

@@ -11,6 +11,7 @@ import numpy as np
 import psutil
 from pydantic import BaseModel, Field, validator
 
+from flock.core.context.context import FlockContext
 from flock.core.flock_agent import FlockAgent
 from flock.core.flock_module import FlockModule, FlockModuleConfig
 
@@ -222,7 +223,11 @@ class MetricsModule(FlockModule):
         return stats
 
     async def terminate(
-        self, agent: FlockAgent, inputs: dict[str, Any], result: dict[str, Any]
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        result: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> None:
         """Clean up and final metric recording."""
         if self.config.storage_type == "json":
@@ -331,7 +336,10 @@ class MetricsModule(FlockModule):
         return False
 
     async def initialize(
-        self, agent: FlockAgent, inputs: dict[str, Any]
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> None:
         """Initialize metrics collection."""
         self._start_time = time.time()
@@ -370,7 +378,10 @@ class MetricsModule(FlockModule):
             return token_count, total_cost
 
     async def pre_evaluate(
-        self, agent: FlockAgent, inputs: dict[str, Any]
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> dict[str, Any]:
         """Record pre-evaluation metrics."""
         if self.config.collect_token_usage:
@@ -408,7 +419,11 @@ class MetricsModule(FlockModule):
         return inputs
 
     async def post_evaluate(
-        self, agent: FlockAgent, inputs: dict[str, Any], result: dict[str, Any]
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        result: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> dict[str, Any]:
         """Record post-evaluation metrics."""
         if self.config.collect_timing and self._start_time:
@@ -463,7 +478,11 @@ class MetricsModule(FlockModule):
         return result
 
     async def on_error(
-        self, agent: FlockAgent, error: Exception, inputs: dict[str, Any]
+        self,
+        agent: FlockAgent,
+        error: Exception,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> None:
         """Record error metrics."""
         self._record_metric(

{flock_core-0.3.15 → flock_core-0.3.16}/src/flock/modules/zep/zep_module.py

@@ -5,6 +5,7 @@ from pydantic import Field
 from zep_python.client import Zep
 from zep_python.types import Message as ZepMessage, SessionSearchResult
 
+from flock.core.context.context import FlockContext
 from flock.core.flock_agent import FlockAgent
 from flock.core.flock_module import FlockModule, FlockModuleConfig
 from flock.core.logging.logging import get_logger
@@ -139,7 +140,11 @@ class ZepModule(FlockModule):
         return response.results
 
     async def post_evaluate(
-        self, agent: FlockAgent, inputs: dict[str, Any], result: dict[str, Any]
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        result: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> dict[str, Any]:
         """Format and display the output."""
         if not self.config.enable_write:
@@ -152,7 +157,10 @@ class ZepModule(FlockModule):
         return result
 
     async def pre_evaluate(
-        self, agent: FlockAgent, inputs: dict[str, Any]
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
     ) -> dict[str, Any]:
         """Format and display the output."""
         if not self.config.enable_read:

{flock_core-0.3.15 → flock_core-0.3.16}/src/flock/workflow/activities.py

@@ -50,7 +50,7 @@ async def run_agent(context: FlockContext) -> dict:
             # Create a nested span for this iteration.
             with tracer.start_as_current_span("agent_iteration") as iter_span:
                 iter_span.set_attribute("agent.name", agent.name)
-
+                agent.context = context
                 # Resolve inputs for the agent.
                 agent_inputs = resolve_inputs(
                     agent.input, context, previous_agent_name
@@ -160,7 +160,7 @@ async def run_agent(context: FlockContext) -> dict:
                     agent.name,
                     result,
                     timestamp=datetime.now().isoformat(),
-                    hand_off=handoff_data,
+                    hand_off=handoff_data.model_dump(),
                     called_from=previous_agent_name,
                 )
                 previous_agent_name = agent.name

{flock_core-0.3.15 → flock_core-0.3.16}/uv.lock

@@ -1119,7 +1119,7 @@ wheels = [
 
 [[package]]
 name = "flock-core"
-version = "0.3.15"
+version = "0.3.16"
 source = { editable = "." }
 dependencies = [
     { name = "chromadb" },