flow-forge-ai-sdk 0.1.0__py3-none-any.whl

flow_forge_ai/config/config_handler.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from functools import lru_cache
+import tomllib
+import os
+from pathlib import Path
+from typing import Any, Optional
+
+from flow_forge_ai.config.models import Config, InstrumentorConfig, RuntimeListenerConfig, SinkConfig
+from flow_forge_ai.internal_logging.logger import get_logger
+from flow_forge_ai.utils.toml import remove_none_values
+
+logger = get_logger(__name__)
+
+
+class ConfigHandler:
+    """Manages configuration for sinks and instrumentation settings."""
+
+    def __init__(self) -> None:
+        self.__data: Config = Config()
+
+    def _normalize(self, config: Config) -> None:
+        for sink in config.sinks:
+            if sink.options is not None:
+                for key, value in sink.options.items():
+                    if isinstance(value, str) and value.startswith("env:"):
+                        env_key = value[4:]
+                        sink.options[key] = os.getenv(env_key)
+
+    def load_from_file(self, file_path: str) -> None:
+        """Load configuration from a TOML file."""
+        try:
+            with open(file_path, "rb") as f:
+                data = tomllib.load(f)
+            self.__data = Config(**data)
+            self._normalize(self.__data)
+            logger.info(f"Configuration loaded from {file_path}")
+        except FileNotFoundError:
+            logger.warning(f"Configuration file {file_path} not found. Using default settings.")
+        except Exception as e:
+            logger.error(f"Failed to load configuration from {file_path}: {e}")
+            raise
+
+    def list_instrumentors(self) -> list[InstrumentorConfig]:
+        """Get configuration for all instrumentors."""
+        return self.__data.instrumentors
+
+    def get_sink(self, sink_name: str) -> SinkConfig:
+        """Get configuration for a specific sink."""
+        sink = next((sink for sink in self.__data.sinks if sink.name == sink_name), None)
+        if not sink:
+            raise KeyError(f"Sink '{sink_name}' not found in configuration.")
+        return sink
+
+    def get_runtime_sink(self) -> SinkConfig:
+        """Get configuration for the runtime sink."""
+        if not self.__data.runtime or not self.__data.runtime.source_sink:
+            raise KeyError("Runtime source_sink is not configured.")
+        return self.get_sink(self.__data.runtime.source_sink)
+
+    def get_runtime_config(self) -> RuntimeListenerConfig:
+        """Return the :class:`~flow_forge_ai.config.models.RuntimeConfig`."""
+        return self.__data.runtime
+
+    def list_sinks(self) -> list[SinkConfig]:
+        """Get configuration for all sinks."""
+        return self.__data.sinks
+
+    def to_dict(self) -> dict[str, Any]:
+        """Return configuration as dictionary."""
+        res: dict[str, Any] = remove_none_values(self.__data.to_dict())
+        return res
+
+
+@lru_cache
+def get_config_handler(path: Optional[str] = None) -> ConfigHandler:
+    """
+    Load configuration from TOML file.
+
+    Parameters
+    ----------
+    path : Optional[str]
+        Path to the configuration file. If None, defaults to 'config.toml' in the current working directory.
+
+    Returns
+    -------
+    ConfigHandler
+        Configuration object with loaded values.
+    """
+    config_handler = ConfigHandler()
+    cfg_path = Path.cwd() / "config.toml" if path is None else Path(path)
+    if not cfg_path.exists():
+        # No default config found, return empty config
+        return config_handler
+    config_path = str(cfg_path)
+    config_handler.load_from_file(config_path)
+    return config_handler

flow_forge_ai/config/models.py

@@ -0,0 +1,66 @@
+from dataclasses import dataclass, field
+from typing import Any, List, Optional
+
+
+@dataclass
+class InstrumentorConfig:
+    """Configuration for a single instrumentor."""
+    class_path: str
+    options: Optional[dict[str, Any]] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "class_path": self.class_path,
+            "options": self.options.copy() if self.options else None
+        }
+
+@dataclass
+class SinkConfig:
+    """Configuration for a single sink."""
+    name: str
+    class_path: str
+    options: Optional[dict[str, Any]] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "class_path": self.class_path,
+            "options": self.options.copy() if self.options else None
+        }
+
+@dataclass
+class RuntimeListenerConfig:
+    """Configuration for runtime listener."""
+    enabled: bool = False
+    source_sink: Optional[str] = None
+    listener_host: Optional[str] = None
+    listener_port: Optional[int] = None
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "enabled": self.enabled,
+            "source_sink": self.source_sink,
+            "listener_host": self.listener_host,
+            "listener_port": self.listener_port,
+        }
+
+@dataclass
+class Config:
+    """Configuration for AI execution infrastructure."""
+    instrumentors: List[InstrumentorConfig] = field(default_factory=list)
+    sinks: List[SinkConfig] = field(default_factory=list)
+    runtime: RuntimeListenerConfig = field(default_factory=RuntimeListenerConfig)
+
+    def __post_init__(self) -> None:
+        self.instrumentors = [InstrumentorConfig(**instr) if isinstance(instr, dict) else instr for instr in self.instrumentors]
+        self.sinks = [SinkConfig(**sink) if isinstance(sink, dict) else sink for sink in self.sinks]
+        if isinstance(self.runtime, dict):
+            self.runtime = RuntimeListenerConfig(**self.runtime) # pylint: disable=not-a-mapping
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert Config dataclass to dictionary."""
+        return {
+            "instrumentors": [instr.to_dict() for instr in self.instrumentors],
+            "sinks": [sink.to_dict() for sink in self.sinks],
+            "runtime": self.runtime.to_dict() if self.runtime else None,
+        }

flow_forge_ai/context.py

@@ -0,0 +1,121 @@
+import uuid
+import contextvars
+from dataclasses import dataclass
+from typing import Any, Optional
+
+
+_workflow_id: contextvars.ContextVar[str] = contextvars.ContextVar("workflow_id", default="default_workflow")
+_run_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar("run_id", default=None)
+_trace_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar("trace_id", default=None)
+_step_id: contextvars.ContextVar[int] = contextvars.ContextVar("step_id", default=0)
+_step_id_alias: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar("step_id_alias", default=None)
+_span_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar("span_id", default=None)
+
+def get_workflow_id() -> str:
+    return _workflow_id.get()
+
+def get_run_id() -> str:
+    return _run_id.get() or f"run_{uuid.uuid4()}"
+
+def get_trace_id() -> str:
+    return _trace_id.get() or f"trace_{uuid.uuid4()}"
+
+def get_step_id() -> str:
+    return _step_id_alias.get() or f"step_{_step_id.get()}"
+
+def get_span_id() -> str:
+    return _span_id.get() or f"span_{uuid.uuid4()}"
+
+def set_workflow_id(value: str) -> contextvars.Token:
+    return _workflow_id.set(value)
+
+def set_run_id(value: str) -> contextvars.Token:
+    return _run_id.set(value)
+
+def set_trace_id(value: str) -> contextvars.Token:
+    return _trace_id.set(value)
+
+def set_step_id_alias(value: str) -> contextvars.Token:
+    return _step_id_alias.set(value)
+
+def increment_step() -> int:
+    next_step = _step_id.get() + 1
+    _step_id.set(next_step)
+    return next_step
+
+def set_span_id(value: str) -> contextvars.Token:
+    return _span_id.set(value)
+
+def reset_workflow_id(token: contextvars.Token) -> None:
+    _workflow_id.reset(token)
+
+def reset_run_id(token: contextvars.Token) -> None:
+    _run_id.reset(token)
+
+def reset_trace_id(token: contextvars.Token) -> None:
+    _trace_id.reset(token)
+
+def reset_step_id_alias(token: contextvars.Token) -> None:
+    _step_id_alias.reset(token)
+
+def reset_step_id() -> None:
+    _step_id.set(0)
+
+def reset_span_id(token: contextvars.Token) -> None:
+    _span_id.reset(token)
+
+
+
+
+@dataclass(frozen=True)
+class ContextSnapshot:
+    run_id: Optional[str]
+    trace_id: Optional[str]
+    span_id: Optional[str]
+
+    @staticmethod
+    def capture() -> "ContextSnapshot":
+        return ContextSnapshot(
+            run_id=_run_id.get(),
+            trace_id=_trace_id.get(),
+            span_id=_span_id.get(),
+        )
+
+    def restore(self) -> None:
+        """Restore this snapshot into the current context (e.g. in a new thread)."""
+        if self.run_id is not None:
+            _run_id.set(self.run_id)
+        if self.trace_id is not None:
+            _trace_id.set(self.trace_id)
+        if self.span_id is not None:
+            _span_id.set(self.span_id)
+
+
+class Span:
+    """
+    Context manager that sets a fresh span_id for its block and restores
+    the previous one on exit.  Optionally starts a new trace as well.
+
+    Usage::
+
+        with Span(name="my-operation") as span:
+            do_work()
+    """
+
+    def __init__(self, name: str = "", new_trace: bool = False):
+        self.name = name
+        self.span_id = str(uuid.uuid4())
+        self.trace_id = str(uuid.uuid4()) if new_trace else (_trace_id.get() or str(uuid.uuid4()))
+        self._tok_span: Optional[contextvars.Token] = None
+        self._tok_trace: Optional[contextvars.Token] = None
+
+    def __enter__(self) -> "Span":
+        self._tok_trace = _trace_id.set(self.trace_id)
+        self._tok_span = _span_id.set(self.span_id)
+        return self
+
+    def __exit__(self, *_: Any) -> None:
+        if self._tok_span is not None:
+            _span_id.reset(self._tok_span)
+        if self._tok_trace is not None:
+            _trace_id.reset(self._tok_trace)

flow_forge_ai/emitter.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from typing import Any
+
+from flow_forge_ai.context import get_run_id, get_step_id, get_trace_id, get_span_id, get_workflow_id
+from flow_forge_ai.sinks.models.event import EventType
+from flow_forge_ai.sinks.sink_router import Event, SinkRouter, default_router
+
+
+def emit_event(
+    event_type: EventType,
+    payload: dict[str, Any],
+    *,
+    router: SinkRouter = default_router,
+    **kwargs: Any
+) -> Event:
+    """
+    Build an :class:`Event` from the current context and *payload*,
+    then forward it to *router*.
+
+    Returns the emitted event (handy for testing assertions).
+    """
+    event = Event(
+        type = event_type,
+        payload = payload,
+        workflow_id = get_workflow_id(),
+        run_id = get_run_id(),
+        trace_id = get_trace_id(),
+        step_id=kwargs["step_id"] if "step_id" in kwargs else get_step_id(),
+        span_id = get_span_id(),
+    )
+    router.emit_event(event)
+    return event

flow_forge_ai/instrumentation/__init__.py

@@ -0,0 +1,24 @@
+import importlib
+
+from flow_forge_ai.config.models import InstrumentorConfig
+from flow_forge_ai.instrumentation.base import BaseInstrumentor
+from flow_forge_ai.internal_logging.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def create_instrumentor(instr_config: InstrumentorConfig) -> BaseInstrumentor:
+    """Instantiates the instrumentor class specified in config."""
+    class_path = instr_config.class_path
+    if not class_path:
+        raise ValueError("No instrumentor class specified in config")
+    module_path, class_name = class_path.rsplit(".", 1)
+    try:
+        module = importlib.import_module(module_path)
+        instrumentor_class = getattr(module, class_name)
+    except (ImportError, AttributeError) as e:
+        raise ImportError(f"Could not load instrumentor '{class_path}': {e}") from e
+    if not issubclass(instrumentor_class, BaseInstrumentor):
+        raise TypeError(f"'{class_name}' must subclass BaseInstrumentor")
+    instance: BaseInstrumentor = instrumentor_class(**(instr_config.options or {}))
+    return instance

flow_forge_ai/instrumentation/base.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+import functools
+import inspect
+from typing import Any, Callable
+
+from flow_forge_ai.instrumentation.utils import step_guard
+from flow_forge_ai.sinks.models.step import Step
+
+
+class BaseInstrumentor(ABC):
+    """
+    Subclasses patch a target library on :meth:`install` and
+    restore original callables on :meth:`uninstall`.
+    """
+
+    def __init__(self) -> None:
+        self._patched = False
+        self._uninstall_hooks: list[Callable[[], None]] = []
+
+    def install(self) -> None:
+        """Apply monkey-patches.  Safe to call multiple times (idempotent)."""
+        if self._patched:
+            return
+        if not self._is_available():
+            return
+        self._install()
+        self._patched = True
+
+    def uninstall(self) -> None:
+        """Uninstalls hooks and restores originals."""
+        if not self._patched:
+            return
+        for hook in self._uninstall_hooks:
+            hook()
+        self._uninstall_hooks.clear()
+        self._patched = False
+
+    def _wrap(self, fn: Callable) -> Callable:
+        if inspect.iscoroutinefunction(fn):
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                step = step_guard()
+                if step is not None:
+                    return self._build_cached_response(step)
+                return await fn(*args, **kwargs)
+            return async_wrapper
+
+        @functools.wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            step = step_guard()
+            if step is not None:
+                return self._build_cached_response(step)
+            return fn(*args, **kwargs)
+        return wrapper
+
+    def _patch(self, target: Any, attr: str) -> Callable:
+        def decorator(fn: Callable) -> Callable:
+            orig = getattr(target, attr)
+            setattr(target, attr, self._wrap(fn))
+            # register the revert alongside the patch
+            self._uninstall_hooks.append(lambda: setattr(target, attr, orig))
+
+            return fn
+        return decorator
+
+    @abstractmethod
+    def _is_available(self) -> bool:
+        """Return True if the target library is importable."""
+
+    @abstractmethod
+    def _install(self) -> None:
+        """Perform the actual patching."""
+
+    @abstractmethod
+    def _build_cached_response(self, step: Step) -> Any:
+        """Reconstruct a native library response from cached *step* data.
+
+        Called by wrappers when :func:`~flow_forge_ai.instrumentation.utils.step_guard`
+        indicates this step should be replayed rather than making a live call.
+        The implementation should locate the ``LLM_RESPONSE`` (or equivalent)
+        event in ``step.events`` and return the appropriate native object.
+        """

flow_forge_ai/instrumentation/httpx_instr.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+import time
+from typing import Any, Optional
+from dataclasses import dataclass
+
+from flow_forge_ai.emitter import emit_event
+from flow_forge_ai.instrumentation.base import BaseInstrumentor
+from flow_forge_ai.instrumentation.models.llm_payloads import LLMErrorPayload, LLMRequestPayload, LLMResponsePayload
+from flow_forge_ai.instrumentation.utils import safe_headers
+from flow_forge_ai.sinks.models.event import EventType
+from flow_forge_ai.sinks.models.step import Step
+from flow_forge_ai.internal_logging.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+@dataclass(init=False)
+class HttpxRequestPayload(LLMRequestPayload):
+    method: str
+
+    def __init__(self,
+                 messages: list[str],
+                 url: str,
+                 method: str,
+                 model: Optional[str] = None,
+                 instructions: Optional[str] = None,
+                 headers: Optional[dict[str, Any]] = None):
+        super().__init__(provider="request",
+                         messages=messages,
+                         url=url,
+                         model=model,
+                         instructions=instructions,
+                         headers=headers)
+        self.method = method
+
+@dataclass(init=False)
+class HttpxResponsePayload(LLMResponsePayload):
+    status_code: int
+    headers: dict[str, Any]
+
+    def __init__(self,
+                 response: list[str],
+                 latency: int,
+                 status_code: int,
+                 headers: dict[str, Any]):
+        super().__init__(response=response,
+                         latency=latency)
+        self.status_code = status_code
+        self.headers = headers
+
+
+HttpxErrorPayload = LLMErrorPayload
+
+class HttpxInstrumentor(BaseInstrumentor):
+
+    def __init__(self, max_body_bytes: Optional[int] = None):
+        super().__init__()
+        self._max_body_bytes = max_body_bytes
+
+    def _is_available(self) -> bool:
+        try:
+            import httpx
+            httpx_version = getattr(httpx, "__version__", None)
+            logger.info(f"httpx library detected, version: {httpx_version}")
+            return True
+        except ImportError:
+            return False
+
+    def _install(self) -> None:
+        import httpx
+
+        orig_send = httpx.Client.send
+
+        @self._patch(httpx.Client, "send")
+        def patched_sync(client_self: Any, request: Any, *args: Any, **kwargs: Any) -> Any:
+            _emit_request(request, self._max_body_bytes)
+            start = time.perf_counter()
+            try:
+                response = orig_send(client_self, request, *args, **kwargs)
+            except Exception as exc:
+                _emit_error(type(exc).__name__, str(exc), int((time.perf_counter() - start) * 1000))
+                raise
+            _emit_response(response, int((time.perf_counter() - start) * 1000), self._max_body_bytes)
+            return response
+
+        orig_async_send = httpx.AsyncClient.send
+
+        @self._patch(httpx.AsyncClient, "send")
+        async def patched_async(client_self: Any, request: Any, *args: Any, **kwargs: Any) -> Any:
+            _emit_request(request, self._max_body_bytes)
+            start = time.perf_counter()
+            try:
+                response = await orig_async_send(client_self, request, *args, **kwargs)
+            except Exception as exc:
+                _emit_error(type(exc).__name__, str(exc), int((time.perf_counter() - start) * 1000))
+                raise
+            _emit_response(response, int((time.perf_counter() - start) * 1000), self._max_body_bytes)
+            return response
+
+    def _build_cached_response(self, step: Step) -> Any:
+        import httpx
+
+        response_event = next((e for e in step.events if e.type == EventType.LLM_RESPONSE), None)
+        if response_event is None:
+            raise ValueError(f"No LLM_RESPONSE event found in step {step.id!r}")
+
+        payload = response_event.payload
+        content = "".join(payload.get("response") or []).encode()
+        return httpx.Response(
+            status_code=payload["status_code"],
+            headers=payload.get("headers") or {},
+            content=content,
+        )
+
+
+def _emit_request(request: Any, max_body_bytes: Optional[int]) -> None:
+    payload = HttpxRequestPayload(
+        url=str(request.url),
+        headers=safe_headers(request.headers),
+        method=request.method,
+        messages=[request.content[:max_body_bytes].decode("utf-8", errors="replace")]
+    )
+    emit_event(EventType.LLM_REQUEST, payload.to_dict())
+
+def _emit_response(response: Any, latency: int, max_body_bytes: Optional[int]) -> None:
+    payload = HttpxResponsePayload(
+        status_code=response.status_code,
+        latency=latency,
+        headers=safe_headers(response.headers),
+        response=[response.content[:max_body_bytes].decode("utf-8", errors="replace")]
+    )
+    emit_event(EventType.LLM_RESPONSE, payload.to_dict())
+
+def _emit_error(error: str, detail: str, latency: int) -> None:
+    emit_event(EventType.LLM_ERROR, HttpxErrorPayload(
+        error=error,
+        detail=detail,
+        latency=latency,
+    ).to_dict())

flow_forge_ai/instrumentation/models/llm_payloads.py

@@ -0,0 +1,76 @@
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+
+@dataclass(init=False)
+class LLMRequestPayload:
+    _provider: str = field(repr=False)
+    model: Optional[str]
+    instructions: Optional[str]
+    messages: list[str]
+    url: str
+    headers: Optional[dict[str, Any]]
+
+    def __init__(self,
+                 provider: str,
+                 messages: list[str],
+                 url: str,
+                 model: Optional[str] = None,
+                 instructions: Optional[str] = None,
+                 headers: Optional[dict[str, Any]] = None):
+        self._provider = provider
+        self.messages = messages
+        self.url = url
+        self.model = model
+        self.instructions = instructions
+        self.headers = headers
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "provider": self._provider,
+            "model": self.model,
+            "instructions": self.instructions,
+            "messages": self.messages,
+            "url": self.url,
+            "headers": self.headers.copy() if self.headers else None
+        }
+
+
+@dataclass(init=False)
+class LLMResponsePayload:
+    response: Any
+    latency: int
+
+    def __init__(self,
+                 response: Any,
+                 latency: int):
+        self.response = response
+        self.latency = latency
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "response": self.response,
+            "latency": self.latency
+        }
+
+
+@dataclass(init=False)
+class LLMErrorPayload:
+    error: str
+    detail: str
+    latency: int
+
+    def __init__(self,
+                 error: str,
+                 detail: str,
+                 latency: int):
+        self.error = error
+        self.detail = detail
+        self.latency = latency
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "error": self.error,
+            "detail": self.detail,
+            "latency": self.latency
+        }