flock-core 0.2.4__tar.gz → 0.2.6__tar.gz

flock_core-0.2.6/.env_template

@@ -0,0 +1,28 @@
+# Main Config
+
+# Your LLM Service Keys
+OPENAI_API_KEY=sk-proj-...
+
+# Logging/Telemetry
+
+LOCAL_DEBUG = true
+LOG_LEVEL = Debug
+
+OTEL_SERVICE_NAME = "flock"
+OTEL_JAEGER_ENDPOINT = "http://localhost:14268/api/traces"
+OTEL_JAEGER_TRANSPORT = "http"
+OTEL_SQL_PATH = "sqlite:///otel.db"
+OTEL_FILE_PATH = "telemetry.json"
+OTEL_ENABLE_SQL = false
+OTEL_ENABLE_FILE = true
+OTEL_ENABLE_JAEGER = false
+
+
+# Tools
+
+GITHUB_PAT=github_pat...
+GITHUB_REPO=
+GITHUB_USERNAME=
+
+TAVILY_API_KEY=tvly-...
+SERPER_API_KEY=

{flock_core-0.2.4 → flock_core-0.2.6}/.gitignore

@@ -25,6 +25,7 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
+logs/
 
 
 # PyInstaller

{flock_core-0.2.4 → flock_core-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flock-core
-Version: 0.2.4
+Version: 0.2.6
 Summary: Declarative LLM Orchestration at Scale
 Author-email: Andre Ratzenberger <andre.ratzenberger@whiteduck.de>
 License-File: LICENSE
@@ -11,6 +11,7 @@ Requires-Python: >=3.10
 Requires-Dist: cloudpickle>=3.1.1
 Requires-Dist: devtools>=0.12.2
 Requires-Dist: dspy==2.5.42
+Requires-Dist: duckduckgo-search>=7.3.2
 Requires-Dist: httpx>=0.28.1
 Requires-Dist: loguru>=0.7.3
 Requires-Dist: msgpack>=1.1.0
@@ -22,6 +23,7 @@ Requires-Dist: opentelemetry-instrumentation-logging>=0.51b0
 Requires-Dist: opentelemetry-sdk>=1.30.0
 Requires-Dist: pydantic>=2.10.5
 Requires-Dist: python-box>=7.3.2
+Requires-Dist: python-decouple>=3.8
 Requires-Dist: rich>=13.9.4
 Requires-Dist: temporalio>=1.9.0
 Requires-Dist: toml>=0.10.2
@@ -36,7 +38,9 @@ Description-Content-Type: text/markdown
 
 <p align="center">
 <img src="docs/img/flock.png" width="600"><br>
-<img alt="Dynamic TOML Badge" src="https://img.shields.io/badge/dynamic/toml?url=https%3A%2F%2Fraw.githubusercontent.com%2Fwhiteducksoftware%2Fflock%2Frefs%2Fheads%2Fbadges%2Fpyproject.toml%3Ftoken%3DGHSAT0AAAAAACVFDVNBU3S6HLJSC36P3YNQZ5LNPQQ&query=%24.project.version&style=for-the-badge&logo=pypi&logoSize=large&label=pip%20version&link=https%3A%2F%2Fpypi.org%2Fproject%2Fflock-core%2F">
+<img alt="Dynamic TOML Badge" src="https://img.shields.io/badge/dynamic/toml?url=https%3A%2F%2Fraw.githubusercontent.com%2Fwhiteducksoftware%2Fflock%2Frefs%2Fheads%2Fmaster%2Fpyproject.toml&query=%24.project.version&style=for-the-badge&logo=pypi&label=pip%20version">
+<img alt="X (formerly Twitter) Follow" src="https://img.shields.io/twitter/follow/whiteduck_gmbh?style=for-the-badge&logo=X">
+
 
 
 

{flock_core-0.2.4 → flock_core-0.2.6}/README.md

@@ -1,6 +1,8 @@
 <p align="center">
 <img src="docs/img/flock.png" width="600"><br>
-<img alt="Dynamic TOML Badge" src="https://img.shields.io/badge/dynamic/toml?url=https%3A%2F%2Fraw.githubusercontent.com%2Fwhiteducksoftware%2Fflock%2Frefs%2Fheads%2Fbadges%2Fpyproject.toml%3Ftoken%3DGHSAT0AAAAAACVFDVNBU3S6HLJSC36P3YNQZ5LNPQQ&query=%24.project.version&style=for-the-badge&logo=pypi&logoSize=large&label=pip%20version&link=https%3A%2F%2Fpypi.org%2Fproject%2Fflock-core%2F">
+<img alt="Dynamic TOML Badge" src="https://img.shields.io/badge/dynamic/toml?url=https%3A%2F%2Fraw.githubusercontent.com%2Fwhiteducksoftware%2Fflock%2Frefs%2Fheads%2Fmaster%2Fpyproject.toml&query=%24.project.version&style=for-the-badge&logo=pypi&label=pip%20version">
+<img alt="X (formerly Twitter) Follow" src="https://img.shields.io/twitter/follow/whiteduck_gmbh?style=for-the-badge&logo=X">
+
 
 
 

flock_core-0.2.6/docs/architecture.drawio

@@ -0,0 +1,24 @@
+<mxfile host="65bd71144e">
+    <diagram id="iQ7YuHOuaVDECIdskFDT" name="Page-1">
+        <mxGraphModel grid="0" page="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" pageScale="1" pageWidth="850" pageHeight="1100" background="light-dark(#FFFFFF, #d8cfb1)" math="0" shadow="1">
+            <root>
+                <mxCell id="0"/>
+                <mxCell id="1" parent="0"/>
+                <mxCell id="2" value="sequenceDiagram&#xa;    participant Client&#xa;    participant Flock&#xa;    participant FlockAgent&#xa;    participant Registry&#xa;    participant Context&#xa;    participant LLM&#xa;    &#xa;    Client-&gt;&gt;Flock: run_async(agent, input)&#xa;    Flock-&gt;&gt;Registry: get_agent(name)&#xa;    Registry--&gt;&gt;Flock: agent instance&#xa;    &#xa;    Flock-&gt;&gt;Context: initialize_context(agent, input)&#xa;    &#xa;    Flock-&gt;&gt;FlockAgent: run(inputs)&#xa;    activate FlockAgent&#xa;    &#xa;    FlockAgent-&gt;&gt;FlockAgent: initialize(inputs)&#xa;    FlockAgent-&gt;&gt;FlockAgent: _evaluate(inputs)&#xa;    &#xa;    FlockAgent-&gt;&gt;LLM: process with tools/prompts&#xa;    LLM--&gt;&gt;FlockAgent: response&#xa;    &#xa;    FlockAgent-&gt;&gt;FlockAgent: terminate(inputs, result)&#xa;    FlockAgent--&gt;&gt;Flock: result&#xa;    deactivate FlockAgent&#xa;    &#xa;    Flock--&gt;&gt;Client: final result" style="shadow=0;dashed=0;align=left;strokeWidth=1;shape=mxgraph.mermaid.abstract.mermaid;labelBackgroundColor=#ffffff;noLabel=1;theme=default;" vertex="1" parent="1">
+                    <mxGeometry x="271" y="360" width="300" height="300" as="geometry"/>
+                </mxCell>
+            </root>
+        </mxGraphModel>
+    </diagram>
+    <diagram id="mQwFuZfSuTjpqv4SXbYC" name="Page-2">
+        <mxGraphModel dx="354" dy="259" grid="0" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="850" pageHeight="1100" background="light-dark(#FFFFFF, #ffffde)" math="0" shadow="1">
+            <root>
+                <mxCell id="0"/>
+                <mxCell id="1" parent="0"/>
+                <mxCell id="gh-ZxXSGDDD0mwRwWCxR-1" value="graph TD&#xa;    subgraph Client Application&#xa;        client[Client Code]&#xa;    end&#xa;&#xa;    subgraph Flock Framework&#xa;        orchestrator[Flock Orchestrator]&#xa;        registry[Agent Registry]&#xa;        context[Context Manager]&#xa;        &#xa;        subgraph Agents&#xa;            base[FlockAgent Base]&#xa;&#xa;        end&#xa;        &#xa;        subgraph Execution&#xa;            local[Local Executor]&#xa;            temporal[Temporal Executor]&#xa;        end&#xa;        &#xa;        subgraph Core Services&#xa;            logging[Logging System]&#xa;            tools[Tool Registry]&#xa;            formatter[Output Formatters]&#xa;        end&#xa;    end&#xa;&#xa;    subgraph External Services&#xa;        llm[Language Models]&#xa;        temporal_server[Temporal Server]&#xa;    end&#xa;&#xa;    %% Client interactions&#xa;    client --&gt; orchestrator&#xa;    &#xa;    %% Orchestrator relations&#xa;    orchestrator --&gt; registry&#xa;    orchestrator --&gt; context&#xa;    orchestrator --&gt; local&#xa;    orchestrator --&gt; temporal&#xa;    &#xa;    %% Agent relations&#xa;    base --&gt; tools&#xa;    base --&gt; logging&#xa;    base --&gt; formatter&#xa;&#xa;    &#xa;    %% Execution relations&#xa;    local --&gt; llm&#xa;    temporal --&gt; temporal_server&#xa;    temporal_server --&gt; llm&#xa;    &#xa;    %% Registry relations&#xa;    registry --&gt; base&#xa;    registry --&gt; tools&#xa;&#xa;    %% Style&#xa;    classDef default fill:#f9f9f9,stroke:#333,stroke-width:2px;&#xa;    classDef external fill:#e1f5fe,stroke:#0288d1,stroke-width:2px;&#xa;    classDef core fill:#f3e5f5,stroke:#7b1fa2,stroke-width:2px;&#xa;    &#xa;    class llm,temporal_server external;&#xa;    class orchestrator,registry,context,base,local,temporal core;" style="shadow=0;dashed=0;align=left;strokeWidth=1;shape=mxgraph.mermaid.abstract.mermaid;labelBackgroundColor=#ffffff;noLabel=1;theme=default;" vertex="1" parent="1">
+                    <mxGeometry x="274" y="325" width="300" height="300" as="geometry"/>
+                </mxCell>
+            </root>
+        </mxGraphModel>
+    </diagram>
+</mxfile>

{flock_core-0.2.4 → flock_core-0.2.6}/examples/01_introduction/03_tool_and_code_agent.py

@@ -52,7 +52,7 @@ async def main():
         name="my_celebrity_age_agent",
         input="a_person",
         output="persons_age_in_days",
-        tools=[basic_tools.web_search_tavily, basic_tools.code_eval],
+        tools=[basic_tools.web_search_duckduckgo, basic_tools.code_eval],
         use_cache=True,
     )
     flock.add_agent(agent)

{flock_core-0.2.4 → flock_core-0.2.6}/examples/01_introduction/05_typed_output2.py

@@ -39,6 +39,8 @@ class RandomPerson:
     favorite_movie: str  
     short_bio: str
 
+RandomUserList = list[RandomPerson]
+
 async def main():
    
     flock = Flock(local_debug=True)
@@ -53,7 +55,7 @@ async def main():
     people_agent = FlockAgent(
         name="people_agent",
         input="amount_of_people",
-        output="random_user_list: list[RandomPerson]",
+        output="random_user_list: RandomUserList",
         use_cache=True,
     )
     flock.add_agent(people_agent)

{flock_core-0.2.4 → flock_core-0.2.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "flock-core"
-version = "0.2.4"
+version = "0.2.6"
 description = "Declarative LLM Orchestration at Scale"
 readme = "README.md"
 authors = [
@@ -16,6 +16,7 @@ dependencies = [
     "cloudpickle>=3.1.1",
     "devtools>=0.12.2",
     "dspy==2.5.42",
+    "duckduckgo-search>=7.3.2",
     "httpx>=0.28.1",
     "loguru>=0.7.3",
     "msgpack>=1.1.0",
@@ -27,6 +28,7 @@ dependencies = [
     "opentelemetry-sdk>=1.30.0",
     "pydantic>=2.10.5",
     "python-box>=7.3.2",
+    "python-decouple>=3.8",
     "rich>=13.9.4",
     "temporalio>=1.9.0",
     "toml>=0.10.2",
@@ -65,6 +67,9 @@ build-backend = "hatchling.build"
 [tool.hatch.build.targets.wheel]
 packages = ["src/flock"]
 
+[project.scripts]
+flock = "flock:main"
+
 
 [tool.ruff]
 cache-dir = ".cache/ruff"

{flock_core-0.2.4 → flock_core-0.2.6}/requirements.txt

@@ -6,16 +6,32 @@ devtools==0.12.2
     # via flock-core (./pyproject.toml)
 dspy==2.5.42
     # via flock-core (./pyproject.toml)
+duckduckgo-search==7.3.2
+    # via flock-core (./pyproject.toml)
 httpx==0.28.1
     # via flock-core (./pyproject.toml)
 loguru==0.7.3
     # via flock-core (./pyproject.toml)
 msgpack==1.1.0
     # via flock-core (./pyproject.toml)
+opentelemetry-api==1.30.0
+    # via flock-core (./pyproject.toml)
+opentelemetry-exporter-jaeger==1.21.0
+    # via flock-core (./pyproject.toml)
+opentelemetry-exporter-jaeger-proto-grpc==1.21.0
+    # via flock-core (./pyproject.toml)
+opentelemetry-exporter-otlp==1.30.0
+    # via flock-core (./pyproject.toml)
+opentelemetry-instrumentation-logging==0.51b0
+    # via flock-core (./pyproject.toml)
+opentelemetry-sdk==1.30.0
+    # via flock-core (./pyproject.toml)
 pydantic==2.10.6
     # via flock-core (./pyproject.toml)
 python-box==7.3.2
     # via flock-core (./pyproject.toml)
+python-decouple==3.8
+    # via flock-core (./pyproject.toml)
 rich==13.9.4
     # via flock-core (./pyproject.toml)
 temporalio==1.9.0

{flock_core-0.2.4 → flock_core-0.2.6}/src/flock/__init__.py

@@ -1,5 +1,8 @@
 """Flock package initialization."""
 
+from opentelemetry import trace
+
 from flock.config import TELEMETRY
 
 tracer = TELEMETRY.setup_tracing()
+tracer = trace.get_tracer(__name__)

flock_core-0.2.6/src/flock/config.py

@@ -0,0 +1,47 @@
+# flock/config.py
+from decouple import config
+
+from flock.core.logging.telemetry import TelemetryConfig
+
+# -- Connection and External Service Configurations --
+TEMPORAL_SERVER_URL = config("TEMPORAL_SERVER_URL", "localhost:7233")
+DEFAULT_MODEL = config("DEFAULT_MODEL", "openai/gpt-4o")
+
+
+# API Keys and related settings
+TAVILY_API_KEY = config("TAVILY_API_KEY", "")
+GITHUB_PAT = config("GITHUB_PAT", "")
+GITHUB_REPO = config("GITHUB_REPO", "")
+GITHUB_USERNAME = config("GITHUB_USERNAME", "")
+
+# -- Debugging and Logging Configurations --
+LOCAL_DEBUG = config("LOCAL_DEBUG", True)
+LOG_LEVEL = config("LOG_LEVEL", "DEBUG")
+LOGGING_DIR = config("LOGGING_DIR", "logs")
+
+OTEL_SERVICE_NAME = config("OTL_SERVICE_NAME", "otel-flock")
+JAEGER_ENDPOINT = config(
+    "JAEGER_ENDPOINT", "http://localhost:14268/api/traces"
+)  # Default gRPC endpoint for Jaeger
+JAEGER_TRANSPORT = config(
+    "JAEGER_TRANSPORT", "http"
+).lower()  # Options: "grpc" or "http"
+OTEL_SQL_DATABASE_NAME = config("OTEL_SQL_DATABASE", "flock_events.db")
+OTEL_FILE_NAME = config("OTEL_FILE_NAME", "flock_events.jsonl")
+OTEL_ENABLE_SQL: bool = config("OTEL_ENABLE_SQL", True) == "True"
+OTEL_ENABLE_FILE: bool = config("OTEL_ENABLE_FILE", True) == "True"
+OTEL_ENABLE_JAEGER: bool = config("OTEL_ENABLE_JAEGER", True) == "True"
+
+
+TELEMETRY = TelemetryConfig(
+    OTEL_SERVICE_NAME,
+    JAEGER_ENDPOINT,
+    JAEGER_TRANSPORT,
+    LOGGING_DIR,
+    OTEL_FILE_NAME,
+    OTEL_SQL_DATABASE_NAME,
+    OTEL_ENABLE_JAEGER,
+    OTEL_ENABLE_FILE,
+    OTEL_ENABLE_SQL,
+)
+TELEMETRY.setup_tracing()

{flock_core-0.2.4 → flock_core-0.2.6}/src/flock/core/flock.py

@@ -5,7 +5,7 @@ import uuid
 from typing import TypeVar
 
 from opentelemetry import trace
-from rich.prompt import Prompt
+from opentelemetry.baggage import get_baggage, set_baggage
 
 from flock.core.context.context import FlockContext
 from flock.core.context.context_manager import initialize_context
@@ -63,6 +63,11 @@ class Flock:
                 enable_logging=enable_logging,
             )
             logger.enable_logging = enable_logging
+            session_id = get_baggage("session_id")
+            if not session_id:
+                session_id = str(uuid.uuid4())
+            set_baggage("session_id", session_id)
+            span.set_attribute("session_id", get_baggage("session_id"))
 
             display_banner()
 
@@ -193,12 +198,16 @@ class Flock:
                     run_id = f"{start_agent.name}_{uuid.uuid4().hex[:4]}"
                     logger.debug("Generated run ID", run_id=run_id)
 
+                set_baggage("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:
+                        from rich.prompt import Prompt
+
                         input[key] = Prompt.ask(
                             f"Please enter {key} for {start_agent.name}"
                         )

{flock_core-0.2.4 → flock_core-0.2.6}/src/flock/core/flock_agent.py

@@ -151,7 +151,7 @@ class FlockAgent(BaseModel, ABC, PromptParserMixin, DSPyIntegrationMixin):
         ),
     )
 
-    tools: list[Callable[..., Any]] | None = Field(
+    tools: list[Callable[..., Any] | Any] | None = Field(
         default=None,
         description="An optional list of callable tools that the agent can leverage during execution.",
     )

{flock_core-0.2.4 → flock_core-0.2.6}/src/flock/core/logging/formatters/pprint_formatter.py

@@ -1,4 +1,3 @@
-import json
 from typing import Any
 
 from flock.core.logging.formatters.base_formatter import BaseFormatter
@@ -10,12 +9,14 @@ class PrettyPrintFormatter(BaseFormatter):
     ) -> None:
         """Print an agent's result using Rich formatting."""
         from rich.console import Console
-        from rich.json import JSON
 
         console = Console()
 
         console.print(agent_name)
-        console.print(JSON(json.dumps(result)))
+        console.print(str(result), markup=True)
+        # try:
+        #     console.print(JSON(json.dumps(result)))
+        # except Exception:
 
     def display_data(self, data: dict[str, Any], **kwargs) -> None:
         """Print an agent's result using Rich formatting."""

{flock_core-0.2.4 → flock_core-0.2.6}/src/flock/core/logging/logging.py

@@ -1,6 +1,5 @@
 # File: src/flock/core/logging.py
-"""A unified logging module for Flock that works both in local/worker contexts
-and inside Temporal workflows.
+"""A unified logging module for Flock that works both in local/worker contexts and inside Temporal workflows.
 
 Key points:
   - We always have Temporal imported, so we cannot decide based on import.
@@ -24,6 +23,7 @@ with workflow.unsafe.imports_passed_through():
 
 def in_workflow_context() -> bool:
     """Returns True if this code is running inside a Temporal workflow context.
+
     It does this by attempting to call workflow.info() and returning True
     if successful. Otherwise, it returns False.
     """
@@ -31,10 +31,7 @@ def in_workflow_context() -> bool:
         workflow.logger.debug("Checking if in workflow context...")
         # loguru_logger.debug("Checking if in workflow context...")
         # This call will succeed only if we're in a workflow context.
-        if hasattr(workflow.info(), "is_replaying"):
-            return True
-        else:
-            return False
+        return bool(hasattr(workflow.info(), "is_replaying"))
     except Exception:
         return False
 
@@ -67,22 +64,24 @@ loguru_logger.add(
 
 # Define a dummy logger that does nothing
 class DummyLogger:
-    def debug(self, *args, **kwargs):
+    """A dummy logger that does nothing when called."""
+
+    def debug(self, *args, **kwargs):  # noqa: D102
         pass
 
-    def info(self, *args, **kwargs):
+    def info(self, *args, **kwargs):  # noqa: D102
         pass
 
-    def warning(self, *args, **kwargs):
+    def warning(self, *args, **kwargs):  # noqa: D102
         pass
 
-    def error(self, *args, **kwargs):
+    def error(self, *args, **kwargs):  # noqa: D102
         pass
 
-    def exception(self, *args, **kwargs):
+    def exception(self, *args, **kwargs):  # noqa: D102
         pass
 
-    def success(self, *args, **kwargs):
+    def success(self, *args, **kwargs):  # noqa: D102
         pass
 
 
@@ -92,7 +91,7 @@ dummy_logger = DummyLogger()
 class FlockLogger:
     """A unified logger that selects the appropriate logging mechanism based on context.
 
-    - If running in a workflow context, it uses Temporal's built‑in logger.
+    - If running in a workflow context, it uses Temporal's built-in logger.
       Additionally, if workflow.info().is_replaying is True, it suppresses debug/info/warning logs.
     - Otherwise, it uses Loguru.
     """
@@ -114,22 +113,22 @@ class FlockLogger:
             trace_id=get_current_trace_id(),
         )
 
-    def debug(self, message: str, *args, **kwargs):
+    def debug(self, message: str, *args, **kwargs):  # noqa: D102
         self._get_logger().debug(message, *args, **kwargs)
 
-    def info(self, message: str, *args, **kwargs):
+    def info(self, message: str, *args, **kwargs):  # noqa: D102
         self._get_logger().info(message, *args, **kwargs)
 
-    def warning(self, message: str, *args, **kwargs):
+    def warning(self, message: str, *args, **kwargs):  # noqa: D102
         self._get_logger().warning(message, *args, **kwargs)
 
-    def error(self, message: str, *args, **kwargs):
+    def error(self, message: str, *args, **kwargs):  # noqa: D102
         self._get_logger().error(message, *args, **kwargs)
 
-    def exception(self, message: str, *args, **kwargs):
+    def exception(self, message: str, *args, **kwargs):  # noqa: D102
         self._get_logger().exception(message, *args, **kwargs)
 
-    def success(self, message: str, *args, **kwargs):
+    def success(self, message: str, *args, **kwargs):  # noqa: D102
         self._get_logger().success(message, *args, **kwargs)
 
 

flock_core-0.2.6/src/flock/core/logging/span_middleware/baggage_span_processor.py

@@ -0,0 +1,31 @@
+from opentelemetry.baggage import get_baggage
+from opentelemetry.sdk.trace import SpanProcessor
+
+
+class BaggageAttributeSpanProcessor(SpanProcessor):
+    """A custom span processor that, on span start, inspects the baggage items from the parent context
+    and attaches specified baggage keys as attributes on the span.
+    """
+
+    def __init__(self, baggage_keys=None):
+        # baggage_keys: list of baggage keys to attach to spans (e.g. ["session_id", "run_id"])
+        if baggage_keys is None:
+            baggage_keys = []
+        self.baggage_keys = baggage_keys
+
+    def on_start(self, span, parent_context):
+        # For each desired key, look up its value in the parent context baggage and set it as an attribute.
+        for key in self.baggage_keys:
+            value = get_baggage(key, context=parent_context)
+            if value is not None:
+                span.set_attribute(key, value)
+
+    def on_end(self, span):
+        # No action required on span end for this processor.
+        pass
+
+    def shutdown(self):
+        pass
+
+    def force_flush(self, timeout_millis: int = 30000):
+        pass

flock_core-0.2.6/src/flock/core/logging/telemetry.py

@@ -0,0 +1,137 @@
+"""This module sets up OpenTelemetry tracing for a service."""
+
+import sys
+
+from opentelemetry import trace
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from temporalio import workflow
+
+from flock.core.logging.span_middleware.baggage_span_processor import (
+    BaggageAttributeSpanProcessor,
+)
+
+with workflow.unsafe.imports_passed_through():
+    from flock.core.logging.telemetry_exporter.file_exporter import (
+        FileSpanExporter,
+    )
+    from flock.core.logging.telemetry_exporter.sqlite_exporter import (
+        SqliteTelemetryExporter,
+    )
+
+
+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 = None,
+        jaeger_transport: str = "grpc",
+        local_logging_dir: str | None = None,
+        file_export_name: str | None = None,
+        sqlite_db_name: str | None = None,
+        enable_jaeger: bool = True,
+        enable_file: bool = True,
+        enable_sql: bool = True,
+        batch_processor_options: dict | None = 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_name = file_export_name
+        self.sqlite_db_name = sqlite_db_name
+        self.local_logging_dir = local_logging_dir
+        self.batch_processor_options = batch_processor_options or {}
+        self.enable_jaeger = enable_jaeger
+        self.enable_file = enable_file
+        self.enable_sql = enable_sql
+
+    def setup_tracing(self):
+        """Set up OpenTelemetry tracing with the specified exporters."""
+        # 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 exporter.
+        if self.jaeger_endpoint and self.enable_jaeger:
+            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(SimpleSpanProcessor(jaeger_exporter))
+
+        # If a file path is provided, add the custom file exporter.
+        if self.file_export_name and self.enable_file:
+            file_exporter = FileSpanExporter(
+                self.local_logging_dir, self.file_export_name
+            )
+            span_processors.append(SimpleSpanProcessor(file_exporter))
+
+        # If a SQLite database path is provided, ensure the DB exists and add the SQLite exporter.
+        if self.sqlite_db_name and self.enable_sql:
+            sqlite_exporter = SqliteTelemetryExporter(
+                self.local_logging_dir, self.sqlite_db_name
+            )
+            span_processors.append(SimpleSpanProcessor(sqlite_exporter))
+
+        # Register all span processors with the provider.
+        for processor in span_processors:
+            provider.add_span_processor(processor)
+
+        provider.add_span_processor(
+            BaggageAttributeSpanProcessor(baggage_keys=["session_id", "run_id"])
+        )
+
+        sys.excepthook = self.log_exception_to_otel
+
+    def log_exception_to_otel(self, exc_type, exc_value, exc_traceback):
+        """Log unhandled exceptions to OpenTelemetry."""
+        if issubclass(exc_type, KeyboardInterrupt):
+            # Allow normal handling of KeyboardInterrupt
+            sys.__excepthook__(exc_type, exc_value, exc_traceback)
+            return
+
+        # Use OpenTelemetry to record the exception
+        with self.global_tracer.start_as_current_span(
+            "UnhandledException"
+        ) as span:
+            span.record_exception(exc_value)
+            span.set_status(
+                trace.Status(trace.StatusCode.ERROR, str(exc_value))
+            )

flock_core-0.2.6/src/flock/core/logging/telemetry_exporter/base_exporter.py

@@ -0,0 +1,38 @@
+"""Base class for custom OpenTelemetry exporters."""
+
+from abc import ABC, abstractmethod
+
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+
+
+class TelemetryExporter(SpanExporter, ABC):
+    """Base class for custom OpenTelemetry exporters."""
+
+    def __init__(self):
+        """Base class for custom OpenTelemetry exporters."""
+        super().__init__()
+
+    def _export(self, spans):
+        """Forward spans to the Jaeger exporter."""
+        try:
+            result = self.export(spans)
+            if result is None:
+                return SpanExportResult.SUCCESS
+            return result
+        except Exception:
+            return SpanExportResult.FAILURE
+        finally:
+            self.shutdown()
+
+    @abstractmethod
+    def export(self, spans) -> SpanExportResult | None:
+        """Export spans to the configured backend.
+
+        To be implemented by subclasses.
+        """
+        raise NotImplementedError("Subclasses must implement the export method")
+
+    @abstractmethod
+    def shutdown(self):
+        """Cleanup resources, if any. Optional for subclasses."""
+        pass