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

flock/core/flock_agent.py

@@ -1,7 +1,7 @@
 """FlockAgent is the core, declarative base class for all agents in the Flock framework."""
 
 from abc import ABC
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
 from typing import Any, Literal, Union
 
@@ -10,16 +10,28 @@ from pydantic import BaseModel, Field
 
 from flock.core.context.context import FlockContext
 from flock.core.logging.logging import get_logger
-from flock.core.mixin.dspy_integration import DSPyIntegrationMixin
+from flock.core.mixin.dspy_integration import AgentType, DSPyIntegrationMixin
 from flock.core.mixin.prompt_parser import PromptParserMixin
 
 logger = get_logger("flock")
+from opentelemetry import trace
+
+tracer = trace.get_tracer(__name__)
 
 
 @dataclass
 class FlockAgentConfig:
     """Configuration options for a FlockAgent."""
 
+    agent_type_override: AgentType = field(
+        default=None,
+        metadata={
+            "description": "Overrides the agent type. TOOL USE ONLY WORKS WITH REACT"
+        },
+    )
+    disable_output: bool = field(
+        default=False, metadata={"description": "Disables the agent's output."}
+    )
     save_to_file: bool = field(
         default=False,
         metadata={
@@ -30,7 +42,7 @@ class FlockAgentConfig:
 
 
 @dataclass
-class HandoffBase:
+class HandOff:
     """Base class for handoff returns."""
 
     next_agent: Union[str, "FlockAgent"] = field(
@@ -124,14 +136,14 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
         "", description="A human-readable description of the agent."
     )
 
-    input: str | None = Field(
+    input: str | Callable[..., str] | None = Field(
         None,
         description=(
             "A comma-separated list of input keys. Optionally supports type hints (:) and descriptions (|). "
             "For example: 'query: str | The search query, chapter_list: list[str] | The chapter list of the document'."
         ),
     )
-    output: str | None = Field(
+    output: str | Callable[..., str] | None = Field(
         None,
         description=(
             "A comma-separated list of output keys.  Optionally supports type hints (:) and descriptions (|). "
@@ -157,7 +169,7 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
         ),
     )
 
-    termination: str | None = Field(
+    termination: str | Callable[..., str] | None = Field(
         None,
         description="An optional termination condition or phrase used to indicate when the agent should stop processing.",
     )
@@ -167,32 +179,58 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
         description="Configuration options for the agent, such as serialization settings.",
     )
 
+    # Lifecycle callback fields: if provided, these callbacks are used instead of overriding the methods.
+    initialize_callback: Callable[[dict[str, Any]], Awaitable[None]] | None = (
+        Field(
+            default=None,
+            description="Optional callback function for initialization. If provided, this async function is called with the inputs.",
+        )
+    )
+    terminate_callback: (
+        Callable[[dict[str, Any], dict[str, Any]], Awaitable[None]] | None
+    ) = Field(
+        default=None,
+        description="Optional callback function for termination. If provided, this async function is called with the inputs and result.",
+    )
+    on_error_callback: (
+        Callable[[Exception, dict[str, Any]], Awaitable[None]] | None
+    ) = Field(
+        default=None,
+        description="Optional callback function for error handling. If provided, this async function is called with the error and inputs.",
+    )
+
     # Lifecycle hooks
     async def initialize(self, inputs: dict[str, Any]) -> None:
-        """The very first thing to get called.
+        """Called at the very start of the agent's execution.
 
-        Override this method to perform any setup or configuration tasks,
-        such as loading resources or validating inputs.
+        Override this method or provide an `initialize_callback` to perform setup tasks such as input validation or resource loading.
         """
-        pass
+        if self.initialize_callback is not None:
+            await self.initialize_callback(self, inputs)
+        else:
+            pass
 
     async def terminate(
         self, inputs: dict[str, Any], result: dict[str, Any]
     ) -> None:
-        """The very last thing to get called.
+        """Called at the very end of the agent's execution.
 
-        Override this method to perform any cleanup tasks,
-        such as releasing resources or logging results.
+        Override this method or provide a `terminate_callback` to perform cleanup tasks such as releasing resources or logging results.
         """
-        pass
+        if self.terminate_callback is not None:
+            await self.terminate_callback(self, inputs, result)
+        else:
+            pass
 
     async def on_error(self, error: Exception, inputs: dict[str, Any]) -> None:
         """Called if the agent encounters an error during execution.
 
-        Override this method to implement
-        custom error handling or recovery strategies.
+        Override this method or provide an `on_error_callback` to implement custom error handling or recovery strategies.
         """
-        pass
+        if self.on_error_callback is not None:
+            await self.on_error_callback(self, error, inputs)
+        else:
+            pass
 
     async def evaluate(self, inputs: dict[str, Any]) -> dict[str, Any]:
         """Process the agent's task using the provided inputs and return the result.
@@ -253,24 +291,35 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
                 - Return a dictionary similar to:
                     {"idea": "A fun app idea based on ...", "query": "build an app", "context": {"previous_idea": "messaging app"}}
         """
-        try:
-            self.__dspy_signature = self.create_dspy_signature_class(
-                self.name, self.description, f"{self.input} -> {self.output}"
-            )
-            # Initialize the language model.
-            self._configure_language_model()
-            # Select the appropriate DSPy task based on tool availability.
-            agent_task = self._select_task(self.__dspy_signature)
-            # Execute the task with the provided inputs.
-            result = agent_task(**inputs)
-            # Process the result and ensure fallback values for missing keys.
-            result = self._process_result(result, inputs)
-            return result
-        except Exception as eval_error:
-            logger.error(
-                f"Error during evaluation in agent '{self.name}': {eval_error}"
-            )
-            raise
+        with tracer.start_as_current_span("agent.evaluate") as span:
+            span.set_attribute("agent.name", self.name)
+            span.set_attribute("inputs", str(inputs))
+            try:
+                # Create and configure the signature and language model.
+                self.__dspy_signature = self.create_dspy_signature_class(
+                    self.name,
+                    self.description,
+                    f"{self.input} -> {self.output}",
+                )
+                self._configure_language_model()
+                agent_task = self._select_task(
+                    self.__dspy_signature,
+                    agent_type_override=self.config.agent_type_override,
+                )
+                # Execute the task.
+                result = agent_task(**inputs)
+                result = self._process_result(result, inputs)
+                span.set_attribute("result", str(result))
+                logger.info("Evaluation successful", agent=self.name)
+                return result
+            except Exception as eval_error:
+                logger.error(
+                    "Error during evaluation",
+                    agent=self.name,
+                    error=str(eval_error),
+                )
+                span.record_exception(eval_error)
+                raise
 
     async def run(self, inputs: dict[str, Any]) -> dict[str, Any]:
         """Run the agent with the given inputs and return its generated output.
@@ -318,15 +367,23 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
             The method might return:
                 {"result": "A conversational chatbot that uses AI to...", "query": "build a chatbot", "context": {"user": "Alice"}}
         """
-        try:
-            await self.initialize(inputs)
-            result = await self.evaluate(inputs)
-            await self.terminate(inputs, result)
-            return result
-        except Exception as run_error:
-            await self.on_error(run_error, inputs)
-            logger.error(f"Error running agent '{self.name}': {run_error}")
-            raise
+        with tracer.start_as_current_span("agent.run") as span:
+            span.set_attribute("agent.name", self.name)
+            span.set_attribute("inputs", str(inputs))
+            try:
+                await self.initialize(inputs)
+                result = await self.evaluate(inputs)
+                await self.terminate(inputs, result)
+                span.set_attribute("result", str(result))
+                logger.info("Agent run completed", agent=self.name)
+                return result
+            except Exception as run_error:
+                logger.error(
+                    "Error running agent", agent=self.name, error=str(run_error)
+                )
+                await self.on_error(run_error, inputs)
+                span.record_exception(run_error)
+                raise
 
     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.
@@ -369,29 +426,54 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
                 result = await agent.run_temporal({"query": "analyze data", "context": {"source": "sales"}})
             will execute the agent on a Temporal worker and return the output in a structured dictionary format.
         """
-        try:
-            from temporalio.client import Client
-
-            from flock.workflow.agent_activities import run_flock_agent_activity
-            from flock.workflow.temporal_setup import run_activity
-
-            client = await Client.connect("localhost:7233", namespace="default")
-            agent_data = self.to_dict()
-            inputs_data = inputs
-
-            result = await run_activity(
-                client,
-                self.name,
-                run_flock_agent_activity,
-                {"agent_data": agent_data, "inputs": inputs_data},
-            )
-            return result
-
-        except Exception as temporal_error:
-            logger.error(
-                f"Error running Temporal workflow for agent '{self.name}': {temporal_error}"
-            )
-            raise
+        with tracer.start_as_current_span("agent.run_temporal") as span:
+            span.set_attribute("agent.name", self.name)
+            span.set_attribute("inputs", str(inputs))
+            try:
+                from temporalio.client import Client
+
+                from flock.workflow.agent_activities import (
+                    run_flock_agent_activity,
+                )
+                from flock.workflow.temporal_setup import run_activity
+
+                client = await Client.connect(
+                    "localhost:7233", namespace="default"
+                )
+                agent_data = self.to_dict()
+                inputs_data = inputs
+
+                result = await run_activity(
+                    client,
+                    self.name,
+                    run_flock_agent_activity,
+                    {"agent_data": agent_data, "inputs": inputs_data},
+                )
+                span.set_attribute("result", str(result))
+                logger.info("Temporal run successful", agent=self.name)
+                return result
+            except Exception as temporal_error:
+                logger.error(
+                    "Error in Temporal workflow",
+                    agent=self.name,
+                    error=str(temporal_error),
+                )
+                span.record_exception(temporal_error)
+                raise
+
+    def resolve_callables(self, context) -> None:
+        """Resolve any callable fields in the agent instance using the provided context.
+
+        This method resolves any callable fields in the agent instance using the provided context. It iterates over
+        the agent's fields and replaces any callable objects (such as lifecycle hooks or tools) with their corresponding
+        resolved values from the context. This ensures that the agent is fully configured and ready
+        """
+        if isinstance(self.input, Callable):
+            self.input = self.input(context)
+        if isinstance(self.output, Callable):
+            self.output = self.output(context)
+        if isinstance(self.description, Callable):
+            self.description = self.description(context)
 
     def to_dict(self) -> dict[str, Any]:
         """Serialize the FlockAgent instance to a dictionary.
@@ -438,7 +520,7 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
                 return {k: convert_callable(v) for k, v in obj.items()}
             return obj
 
-        data = self.dict()
+        data = self.model_dump()
         return convert_callable(data)
 
     @classmethod

flock/core/logging/logging.py

@@ -13,6 +13,8 @@ Key points:
 
 import sys
 
+from opentelemetry import trace
+
 # Always import Temporal workflow (since it's part of the project)
 from temporalio import workflow
 
@@ -37,6 +39,16 @@ def in_workflow_context() -> bool:
         return False
 
 
+def get_current_trace_id() -> str:
+    """Fetch the current trace ID from OpenTelemetry, if available."""
+    current_span = trace.get_current_span()
+    span_context = current_span.get_span_context()
+    # Format the trace_id as hex (if valid)
+    if span_context.is_valid:
+        return format(span_context.trace_id, "032x")
+    return "no-trace"
+
+
 # Configure Loguru for non-workflow (local/worker) contexts.
 # Note that in workflow code, we will use Temporal's workflow.logger instead.
 loguru_logger.remove()
@@ -44,7 +56,10 @@ loguru_logger.add(
     sys.stderr,
     level="DEBUG",
     colorize=True,
-    format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{message}</cyan>",
+    format=(
+        "<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | "
+        "<cyan>[trace_id: {extra[trace_id]}]</cyan> | <magenta>[{extra[category]}]</magenta> | {message}"
+    ),
 )
 # Optionally add a file handler, e.g.:
 # loguru_logger.add("logs/flock.log", rotation="100 MB", retention="30 days", level="DEBUG")
@@ -92,9 +107,12 @@ class FlockLogger:
         if in_workflow_context():
             # Use Temporal's workflow.logger inside a workflow context.
             return workflow.logger
-        else:
-            # Bind a new Loguru logger with the given name as context.
-            return loguru_logger.bind(name=self.name)
+        # Bind our logger with category and trace_id
+        return loguru_logger.bind(
+            name=self.name,
+            category=self.name,  # Customize this per module (e.g., "flock", "agent", "context")
+            trace_id=get_current_trace_id(),
+        )
 
     def debug(self, message: str, *args, **kwargs):
         self._get_logger().debug(message, *args, **kwargs)

flock/core/logging/telemetry.py

@@ -0,0 +1,109 @@
+"""This module sets up OpenTelemetry tracing for a service."""
+
+from opentelemetry import trace
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import (
+    BatchSpanProcessor,
+)
+
+from flock.core.logging.telemetry_exporter.file_span import FileSpanExporter
+from flock.core.logging.telemetry_exporter.sqllite_span import (
+    SQLiteSpanExporter,
+)
+
+
+class TelemetryConfig:
+    """This configuration class sets up OpenTelemetry tracing.
+
+      - Export spans to a Jaeger collector using gRPC.
+      - Write spans to a file.
+      - Save spans in a SQLite database.
+
+    Only exporters with a non-None configuration will be activated.
+    """
+
+    def __init__(
+        self,
+        service_name: str,
+        jaeger_endpoint: str = None,
+        jaeger_transport: str = "grpc",
+        file_export_path: str = None,
+        sqlite_db_path: str = None,
+        batch_processor_options: dict = None,
+    ):
+        """:param service_name: Name of your service.
+
+        :param jaeger_endpoint: The Jaeger collector gRPC endpoint (e.g., "localhost:14250").
+        :param file_export_path: If provided, spans will be written to this file.
+        :param sqlite_db_path: If provided, spans will be stored in this SQLite DB.
+        :param batch_processor_options: Dict of options for BatchSpanProcessor (e.g., {"max_export_batch_size": 10}).
+        """
+        self.service_name = service_name
+        self.jaeger_endpoint = jaeger_endpoint
+        self.jaeger_transport = jaeger_transport
+        self.file_export_path = file_export_path
+        self.sqlite_db_path = sqlite_db_path
+        self.batch_processor_options = batch_processor_options or {}
+
+    def setup_tracing(self):
+        # Create a Resource with the service name.
+        resource = Resource(attributes={"service.name": self.service_name})
+        provider = TracerProvider(resource=resource)
+        trace.set_tracer_provider(provider)
+
+        # List to collect our span processors.
+        span_processors = []
+
+        # If a Jaeger endpoint is specified, add the Jaeger gRPC exporter.
+        if self.jaeger_endpoint:
+            if self.jaeger_transport == "grpc":
+                from opentelemetry.exporter.jaeger.proto.grpc import (
+                    JaegerExporter,
+                )
+
+                jaeger_exporter = JaegerExporter(
+                    endpoint=self.jaeger_endpoint,
+                    insecure=True,
+                )
+            elif self.jaeger_transport == "http":
+                from opentelemetry.exporter.jaeger.thrift import JaegerExporter
+
+                jaeger_exporter = JaegerExporter(
+                    collector_endpoint=self.jaeger_endpoint,
+                )
+            else:
+                raise ValueError(
+                    "Invalid JAEGER_TRANSPORT specified. Use 'grpc' or 'http'."
+                )
+
+            span_processors.append(
+                BatchSpanProcessor(
+                    jaeger_exporter, **self.batch_processor_options
+                )
+            )
+
+        # If a file path is provided, add the custom file exporter.
+        if self.file_export_path:
+            file_exporter = FileSpanExporter(self.file_export_path)
+            span_processors.append(
+                BatchSpanProcessor(
+                    file_exporter, **self.batch_processor_options
+                )
+            )
+
+        # If a SQLite database path is provided, add the custom SQLite exporter.
+        if self.sqlite_db_path:
+            sqlite_exporter = SQLiteSpanExporter(self.sqlite_db_path)
+            span_processors.append(
+                BatchSpanProcessor(
+                    sqlite_exporter, **self.batch_processor_options
+                )
+            )
+
+        # Register all span processors with the provider.
+        for processor in span_processors:
+            provider.add_span_processor(processor)
+
+        # Return a tracer instance.
+        return trace.get_tracer(__name__)

flock/core/logging/telemetry_exporter/file_span.py

@@ -0,0 +1,37 @@
+import json
+
+from opentelemetry.sdk.trace.export import (
+    SpanExporter,
+    SpanExportResult,
+)
+
+
+class FileSpanExporter(SpanExporter):
+    """A simple exporter that writes span data as JSON lines into a file."""
+
+    def __init__(self, file_path: str):
+        self.file_path = file_path
+
+    def export(self, spans) -> SpanExportResult:
+        try:
+            with open(self.file_path, "a") as f:
+                for span in spans:
+                    # Create a dictionary representation of the span.
+                    span_dict = {
+                        "name": span.name,
+                        "trace_id": format(span.context.trace_id, "032x"),
+                        "span_id": format(span.context.span_id, "016x"),
+                        "start_time": span.start_time,
+                        "end_time": span.end_time,
+                        "attributes": span.attributes,
+                        "status": str(span.status),
+                    }
+                    f.write(json.dumps(span_dict) + "\n")
+            return SpanExportResult.SUCCESS
+        except Exception as e:
+            print("Error exporting spans to file:", e)
+            return SpanExportResult.FAILURE
+
+    def shutdown(self) -> None:
+        # Nothing special needed on shutdown.
+        pass

flock/core/logging/telemetry_exporter/sqllite_span.py

@@ -0,0 +1,68 @@
+import json
+import sqlite3
+
+from opentelemetry.sdk.trace.export import (
+    SpanExporter,
+    SpanExportResult,
+)
+
+
+class SQLiteSpanExporter(SpanExporter):
+    """A custom exporter that writes span data into a SQLite database."""
+
+    def __init__(self, sqlite_db_path: str):
+        self.sqlite_db_path = sqlite_db_path
+        self.conn = sqlite3.connect(
+            self.sqlite_db_path, check_same_thread=False
+        )
+        self._create_table()
+
+    def _create_table(self):
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """
+            CREATE TABLE IF NOT EXISTS spans (
+                id TEXT PRIMARY KEY,
+                name TEXT,
+                trace_id TEXT,
+                span_id TEXT,
+                start_time INTEGER,
+                end_time INTEGER,
+                attributes TEXT,
+                status TEXT
+            )
+            """
+        )
+        self.conn.commit()
+
+    def export(self, spans) -> SpanExportResult:
+        try:
+            cursor = self.conn.cursor()
+            for span in spans:
+                span_id = format(span.context.span_id, "016x")
+                trace_id = format(span.context.trace_id, "032x")
+                cursor.execute(
+                    """
+                    INSERT OR REPLACE INTO spans 
+                    (id, name, trace_id, span_id, start_time, end_time, attributes, status)
+                    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+                    """,
+                    (
+                        span_id,
+                        span.name,
+                        trace_id,
+                        span_id,
+                        span.start_time,
+                        span.end_time,
+                        json.dumps(span.attributes),
+                        str(span.status),
+                    ),
+                )
+            self.conn.commit()
+            return SpanExportResult.SUCCESS
+        except Exception as e:
+            print("Error exporting spans to SQLite:", e)
+            return SpanExportResult.FAILURE
+
+    def shutdown(self) -> None:
+        self.conn.close()

flock/core/logging/trace_and_logged.py

@@ -0,0 +1,55 @@
+import functools
+import inspect
+
+from opentelemetry import trace
+
+from flock.core.logging.logging import get_logger
+
+logger = get_logger("tools")
+tracer = trace.get_tracer(__name__)
+
+
+def traced_and_logged(func):
+    """A decorator that wraps a function in an OpenTelemetry span and logs its inputs,
+    outputs, and exceptions. Supports both synchronous and asynchronous functions.
+    """
+    if inspect.iscoroutinefunction(func):
+
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            with tracer.start_as_current_span(func.__name__) as span:
+                span.set_attribute("args", str(args))
+                span.set_attribute("kwargs", str(kwargs))
+                try:
+                    result = await func(*args, **kwargs)
+                    span.set_attribute("result", str(result))
+                    logger.debug(
+                        f"{func.__name__} executed successfully", result=result
+                    )
+                    return result
+                except Exception as e:
+                    logger.error(f"Error in {func.__name__}", error=str(e))
+                    span.record_exception(e)
+                    raise
+
+        return async_wrapper
+    else:
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            with tracer.start_as_current_span(func.__name__) as span:
+                span.set_attribute("args", str(args))
+                span.set_attribute("kwargs", str(kwargs))
+                try:
+                    result = func(*args, **kwargs)
+                    span.set_attribute("result", str(result))
+                    logger.debug(
+                        f"{func.__name__} executed successfully", result=result
+                    )
+                    return result
+                except Exception as e:
+                    logger.error(f"Error in {func.__name__}", error=str(e))
+                    span.record_exception(e)
+                    raise
+
+        return wrapper