writer 0.8.3rc21__py3-none-any.whl → 1.25.1rc1__py3-none-any.whl

writer/journal.py

@@ -0,0 +1,227 @@
+import contextlib
+import json
+import logging
+from contextvars import ContextVar
+from copy import deepcopy
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, Dict, Literal, Optional
+
+import writer.abstract
+from writer.core import Config
+from writer.keyvalue_storage import writer_kv_storage
+
+if TYPE_CHECKING:
+    from writer.blueprints import Graph, GraphNode
+    from writer.core import Component
+
+
+logger = logging.getLogger("journal")
+
+JOURNAL_KEY_PREFIX = "wf-journal-"
+INIT_LOGS_KEY_PREFIX = "wf-init-logs-"
+
+class JournalRecord:
+    def __init__(
+        self,
+        execution_environment: Dict,
+        title: str,
+        graph: "Graph"
+    ):
+        from writer import core_ui
+
+        self.started_at = datetime.now(timezone.utc)
+        self.instance_type = "editor" if Config.mode == "edit" else "agent"
+
+        # Get blueprint_id from the parent of any node in the graph
+        # All nodes in a blueprint share the same parent blueprint component
+        self.blueprint_id = graph.nodes[0].component.parentId if graph.nodes else None
+
+        self.execution_environment = execution_environment
+        self.trigger = {
+            "event": execution_environment.get("context", {}).get("event"),
+            "payload": execution_environment.get("payload"),
+            "component": {}
+        }
+
+        if self.trigger["event"] == "wf-run-blueprint":
+            self.trigger["component"]["type"] = "blueprint"
+            self.trigger["component"]["id"] = self.blueprint_id
+            blueprint_component = core_ui.current_component_tree().get_component(self.trigger["component"]["id"])
+            if blueprint_component is not None:
+                self.trigger["component"]["title"] = blueprint_component.content.get("key")
+        else:
+            self.trigger["component"]["type"] = "block"
+            component = graph.get_start_nodes()[0].component
+            self.trigger["component"]["id"] = component.id
+            self.trigger["component"]["title"] = self._get_block_info(component)["title"]
+
+        if "API" in title:
+            self.trigger["type"] = "API"
+        elif "Cron" in title:
+            self.trigger["type"] = "Cron"
+        elif "UI" in title:
+            self.trigger["type"] = "UI"
+        else:
+            self.trigger["type"] = "On demand"
+
+        self.graph = graph
+        self.block_outputs: Dict[str, Any] = {}
+        for graph_node in self.graph.nodes:
+            block_info = self._get_block_info(graph_node.component)
+            self.block_outputs[graph_node.id] = {
+                "component": {
+                    "type": graph_node.component.type,
+                    "id": graph_node.component.id,
+                    "title": block_info["title"],
+                    "category": block_info["category"]
+                },
+                "executions": []
+            }
+
+        self.is_runable = True
+        self.result: Optional[Literal["success", "error", "stopped"]] = None
+
+    def _get_block_info(self, component: "Component") -> Dict[str, str]:
+        block_title = component.content.get("alias")
+        component_definition = writer.abstract.templates.get(component.type)
+
+        # If component has an alias, use it as title
+        if block_title is not None:
+            category = "Unknown category"
+            if component_definition is not None:
+                category = component_definition.writer.get("category", "Unknown category")
+            return {
+                "title": block_title,
+                "category": category
+            }
+
+        # If no component definition found, return defaults
+        if component_definition is None:
+            return {
+                "title": "Unknown block",
+                "category": "Unknown category"
+            }
+
+        # Use component definition for both title and category
+        return {
+            "title": component_definition.writer.get("name", "Unknown block"),
+            "category": component_definition.writer.get("category", "Unknown category")
+        }
+
+    def to_dict(self) -> Dict[str, Any]:
+        block_outputs = deepcopy(self.block_outputs)
+        for graph_node in self.graph.nodes:
+            block_outputs[graph_node.id]["executions"].append(self.get_execution_data(graph_node))
+
+        data = {
+            "timestamp": self.started_at.isoformat(),
+            "instanceType": self.instance_type,
+            "blueprintId": self.blueprint_id,
+            "trigger": self.trigger,
+            "blockOutputs": block_outputs,
+            "result": self.result,
+        }
+        sanitized_data = self._sanitize_data(data)
+        return {
+            **sanitized_data,
+            "isRunable": self.is_runable,
+        }
+
+    def get_execution_data(self, graph_node: "GraphNode") -> Dict[str, Any]:
+        execution_data: Dict[str, Any] = {
+            "result": graph_node.result,
+            "outcome": graph_node.outcome,
+        }
+
+        # Add timing information if available
+        if graph_node.tool:
+            if hasattr(graph_node.tool, 'started_at') and graph_node.tool.started_at >= 0:
+                execution_data["startedAt"] = graph_node.tool.started_at
+            if hasattr(graph_node.tool, 'execution_time_in_seconds') and graph_node.tool.execution_time_in_seconds >= 0:
+                execution_data["executionTimeInSeconds"] = graph_node.tool.execution_time_in_seconds
+
+            # Add captured logs if available
+            if hasattr(graph_node.tool, 'captured_stdout') and graph_node.tool.captured_stdout:
+                execution_data["stdout"] = graph_node.tool.captured_stdout
+            if hasattr(graph_node.tool, 'captured_logs') and graph_node.tool.captured_logs:
+                execution_data["logs"] = graph_node.tool.captured_logs
+
+            # Add error message if available (contains the traceback for errors)
+            if hasattr(graph_node.tool, 'message') and graph_node.tool.message:
+                execution_data["message"] = graph_node.tool.message
+
+        return execution_data
+
+    def _sanitize_data(self, data):
+        if data is None:
+            return None
+
+        if isinstance(data, list):
+            return [self._sanitize_data(item) for item in data]
+        if isinstance(data, dict):
+            return {
+                k: self._sanitize_data(v)
+                for k, v in data.items()
+            }
+        if isinstance(data, (str, int, float, bool, type(None))):
+            return data
+
+        try:
+            return json.loads(json.dumps(data))
+        except (TypeError, OverflowError):
+            self.is_runable = False
+            return f"Can't be displayed in the Journal. Value of type: {str(type(data))}."
+
+    def construct_key(self) -> str:
+        return f"{JOURNAL_KEY_PREFIX}{self.instance_type[0]}-{int(self.started_at.timestamp() * 1000)}"
+    
+    def set_result(self, result: Literal["success", "error", "stopped"]) -> None:
+        self.result = result
+
+    def add_nested_execution(self, nested_record: "JournalRecord") -> None:
+        for graph_node in nested_record.graph.nodes:
+            if graph_node.id not in self.block_outputs:
+                self.block_outputs[graph_node.id] = nested_record.block_outputs[graph_node.id]
+            self.block_outputs[graph_node.id]["executions"].append(nested_record.get_execution_data(graph_node))
+    
+    def save(self) -> None:
+        if "journal" not in Config.feature_flags or not writer_kv_storage.is_accessible():
+            return
+        data = self.to_dict()
+        writer_kv_storage.save(self.construct_key(), data)
+
+
+_parent_journal_record: ContextVar[Optional[JournalRecord]] = ContextVar("parent_journal_record", default=None)
+_current_journal_record: ContextVar[Optional[JournalRecord]] = ContextVar("current_journal_record", default=None)
+
+@contextlib.contextmanager
+def use_journal_record_context(
+    execution_environment: Dict,
+    title: str,
+    graph: "Graph"
+):
+    parent_record = _parent_journal_record.get()
+    current_record = JournalRecord(execution_environment, title, graph)
+    _current_journal_record.set(current_record)
+    if parent_record is None:
+        _parent_journal_record.set(current_record)
+
+    try:
+        yield current_record
+    except BaseException as e:
+        current_record.set_result("error")
+        raise e
+    finally:
+        _current_journal_record.set(None)
+        if parent_record is not None:
+            parent_record.add_nested_execution(current_record)
+        else:
+            try:
+                current_record.save()
+            except Exception:
+                logger.exception("Failed to save a Journal entry")
+            _parent_journal_record.set(None)
+
+
+def get_current_journal_record() -> Optional[JournalRecord]:
+    return _current_journal_record.get()

writer/keyvalue_storage.py

@@ -0,0 +1,93 @@
+import logging
+import os
+from functools import partial
+from typing import Any, Dict, List, Literal, Optional, Protocol
+
+import httpx
+
+logger = logging.getLogger("kv_storage")
+
+
+class _WrappedRequestFunc(Protocol):
+    def __call__(
+        self,
+        headers: Dict[str, str],
+        timeout: int,
+    ) -> httpx.Response: ...
+
+
+class KeyValueStorage:
+    def __init__(self, client: Optional[httpx.Client] = None) -> None:
+        base_url = os.getenv("WRITER_BASE_URL")
+        self.api_key = os.getenv("WRITER_API_KEY")
+        if None in (base_url, self.api_key):
+            logger.warning("Missing required environment variables for KV storage access")
+        self.api_url = f"{base_url}/v1" if base_url else None
+
+        self._client = client if client is not None else httpx
+
+    def _get_agent_ids(self):
+        from writer.core import get_session
+        current_session = get_session()
+
+        if current_session:
+            headers = current_session.headers or {}
+            agent_id = headers.get("x-agent-id") or os.getenv("WRITER_APP_ID")
+            org_id = headers.get("x-organization-id") or os.getenv("WRITER_ORG_ID")
+            return (agent_id, org_id)
+
+        agent_id = os.getenv("WRITER_APP_ID")
+        org_id = os.getenv("WRITER_ORG_ID")
+        return (agent_id, org_id)
+
+    def get(self, key: str, type_: Literal["data", "secret"]) -> Dict[str, Any]:
+        return self._request(partial(self._client.get, url=f"{self.api_url}/agent_{type_}/{key}")).json()
+
+    def get_data_keys(self) -> List[str]:
+        return self._request(partial(self._client.get, url=f"{self.api_url}/agent_data")).json()["keys"]
+    
+    def save(self, key: str, data: Any) -> Dict[str, Any]:
+        try:
+            return self._create(key, data).json()
+        except httpx.HTTPStatusError as e:
+            if "already exists" in e.response.text:
+                return self._update(key, data).json()
+            raise e
+
+    def _create(self, key: str, data: Any) -> httpx.Response:
+        return self._request(partial(self._client.post, url=f"{self.api_url}/agent_data", json={"key": key, "data": data}))
+
+    def _update(self, key: str, data: Any) -> httpx.Response:
+        return self._request(partial(self._client.put, url=f"{self.api_url}/agent_data/{key}", json={"data": data}))
+
+    def delete(self, key: str) -> Dict[str, str]:
+        self._request(partial(self._client.delete, url=f"{self.api_url}/agent_data/{key}"))
+        return {"key": key}
+
+    def _request(self, request_func: _WrappedRequestFunc) -> httpx.Response:
+
+        agent_id, org_id = self._get_agent_ids()
+        if None in (agent_id, org_id):
+            raise ValueError("Can't access KV storage. Missing agent id or org id")
+
+        if None in (self.api_key, self.api_url):
+            raise ValueError("Can't access KV storage. Missing required env vars")
+
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "X-Organization-Id": org_id,
+            "X-Agent-Id": agent_id,
+        }
+
+        response = request_func(headers=headers, timeout=3)
+        response.raise_for_status()
+        return response
+
+    def is_accessible(self) -> bool:
+        if None in self._get_agent_ids():
+            return False
+        if None in (self.api_key, self.api_url):
+            return False
+        return True
+
+writer_kv_storage = KeyValueStorage()

writer/logs.py

@@ -0,0 +1,277 @@
+import io
+import json
+import logging
+import logging.config
+import os
+from contextlib import contextmanager, redirect_stdout
+from time import time
+from typing import Any, Callable, Dict, List, Optional, Union
+
+WRITER_LOG_LEVEL = os.getenv("WRITER_LOG_LEVEL", "INFO")
+WRITER_LOG_FORMAT = os.getenv("WRITER_LOG_FORMAT", "text")  # 'text' or 'json'
+
+FAILOVER_ROUTING_KEY = "unset"
+FAILOVER_BUFFER = "failover"
+
+
+def get_routing_key(prefix: Optional[str] = None) -> str:
+    try:
+        from writer.blueprints import get_current_block
+        current_block = get_current_block()
+    except RuntimeError:
+        current_block = None
+
+    key = FAILOVER_ROUTING_KEY
+    if current_block is not None:
+        key = current_block.component.id
+    return f"{prefix}-{key}"
+
+
+class RoutingMap():
+    """
+    Maintains a map of routing keys to in-memory output buffers (io.StringIO).
+    Used for capturing logs or stdout output in different contexts.
+    """
+
+    def __init__(self) -> None:
+        # It's not expected that this will be used without context.
+        # But just in case a fail-over buffer is provided
+        self._buffer_map: Dict[str, io.StringIO] = {
+            FAILOVER_BUFFER: io.StringIO(),
+        }
+    
+    def get_buffer(self, key: str) -> io.StringIO:
+        """
+        Retrieve the buffer associated with a routing key.
+
+        If buffer for the key is not present uses a fail-over buffer
+        """
+        return self._buffer_map.get(key, self._buffer_map[FAILOVER_BUFFER])
+
+    def add_buffer(self, key: str) -> io.StringIO:
+        """
+        Add a new buffer with the given key.
+        """
+        buffer = io.StringIO()
+        self._buffer_map[key] = buffer
+        return buffer
+
+    def remove_buffer(self, key: str) -> None:
+        """
+        Remove the buffer associated with the specified key.
+        """
+        self._buffer_map.pop(key, None)
+
+
+routing_map = RoutingMap()
+
+
+class RoutingStream(io.StringIO):
+    """
+    Custom stream that re-routes stdout to the correct io.StringIO buffer
+    based on the current context routing key.
+    """
+
+    def write(self, s: str) -> int:
+        if s.strip():
+            logging.getLogger("stdout").info(s)
+        routing_key = get_routing_key(prefix="stdout")
+        return routing_map.get_buffer(routing_key).write(s)
+    
+    def getvalue(self) -> str:
+        routing_key = get_routing_key(prefix="stdout")
+        return routing_map.get_buffer(routing_key).getvalue()
+
+
+class RoutingHandler(logging.StreamHandler):
+    """
+    Custom logging handler that re-routes logs to different buffers
+    based on the current context routing key.
+
+    Overwritten methods are mirroring original ones from logging.StreamHandler.
+    The only difference is how 'stream' object is acquired
+    """
+
+    def emit(self, record):
+        try:
+            msg = self.format(record)
+            routing_key = get_routing_key(prefix="logging")
+            stream = routing_map.get_buffer(routing_key)
+            stream.write(msg + self.terminator)
+            self.flush()
+        except RecursionError:
+            raise
+        except Exception:
+            self.handleError(record)
+
+    def flush(self):
+        routing_key = get_routing_key(prefix="logging")
+        stream = routing_map.get_buffer(routing_key)
+        with self.lock:
+            if stream and hasattr(stream, "flush"):
+                stream.flush()
+
+
+@contextmanager
+def use_stdout_redirect(add_log_entry_func: Union[Callable[[str], None], List[Callable[[str], None]]]):
+    """
+    Context manager that redirects stdout to a context-specific buffer.
+    Supports single callback or list of callbacks for multiple output destinations.
+    """
+    # Normalize to list of callbacks
+    callbacks = [add_log_entry_func] if callable(add_log_entry_func) else add_log_entry_func
+
+    key = get_routing_key(prefix="stdout")
+    buffer = routing_map.add_buffer(key)
+    try:
+        with redirect_stdout(RoutingStream()):
+            yield buffer
+    finally:
+        routing_map.remove_buffer(key)
+        stdout = buffer.getvalue()
+        if stdout:
+            for callback in callbacks:
+                callback(stdout)
+
+
+@contextmanager
+def use_logging_redirect(add_log_entry_func: Union[Callable[[str], None], List[Callable[[str], None]]]):
+    """
+    Context manager that redirects logging to a context-specific buffer.
+    Supports single callback or list of callbacks for multiple output destinations.
+    """
+    # Normalize to list of callbacks
+    callbacks = [add_log_entry_func] if callable(add_log_entry_func) else add_log_entry_func
+
+    key = get_routing_key(prefix="logging")
+    buffer = routing_map.add_buffer(key)
+    try:
+        yield buffer
+    finally:
+        routing_map.remove_buffer(key)
+        logs = buffer.getvalue()
+        if logs:
+            for callback in callbacks:
+                callback(logs)
+
+
+class JSONFormatter(logging.Formatter):
+    def format(self, record: logging.LogRecord, **kwargs) -> str:
+        try:
+            from writer.blueprints import get_current_block
+            from writer.core import get_app_process, get_session
+            current_block = get_current_block()
+            app_process = get_app_process()
+            session = get_session()
+        except RuntimeError:
+            current_block = None
+            session = None
+            app_process = None
+
+        data: Dict[str, Any] = {
+            "severity": record.levelname.upper(),
+            "message": super().format(record),
+            "logger_name": record.name,
+            "process": {
+                "name": record.processName
+            },
+            "timestamp": time(),
+        }
+
+        if current_block is not None:
+            data["component"] = {
+                "id": current_block.component.id,
+                "type": current_block.component.type
+            }
+
+        if session is not None:
+            data["session"] = {
+                "id": session.session_id,
+            }
+
+        if app_process is not None:
+            data["process"]["mode"] = app_process.mode
+
+        if isinstance(record.args, dict):
+            data.update(record.args)
+
+        # if record.args[0] is a dict add to the json dict
+        if isinstance(record.args, tuple):
+            if len(record.args) > 0 and isinstance(record.args[0], dict):
+                data.update(record.args[0])
+
+        try:
+            data_as_json = json.dumps(data)
+            return data_as_json 
+        except Exception:
+            return '{"unserializable": true}' 
+
+
+def get_handler(format: str = WRITER_LOG_FORMAT, level: str = WRITER_LOG_LEVEL):
+    return {
+        "level": level,
+        "class": "logging.StreamHandler",
+        "formatter": format,
+    }
+
+
+def get_logger(level: str = WRITER_LOG_LEVEL):
+    return {
+        "handlers": ["basic"],
+        "level": level,
+        "propagate": False,
+    }
+
+
+LOGGING_CONFIG: Dict[str, Any] = {
+    "version": 1,
+    "disable_existing_loggers": False,
+    "formatters": {
+        "text": {
+            "format": "%(levelname)s - %(name)s - %(message)s",
+        },
+        "json": {
+            "()": JSONFormatter,
+        },
+        "user": {
+            "format": "%(levelname)s - %(message)s"
+        }
+    },
+    "handlers": {
+        "basic": get_handler(),
+        "stdout": {
+            "level": "DEBUG",
+            "class": "logging.StreamHandler",
+            "formatter": "json" if WRITER_LOG_FORMAT == "json" else None
+        },
+        "routing": {
+            "()": RoutingHandler,
+            "formatter": "user",
+        }
+    },
+    "loggers": {
+        "root": get_logger(),
+        "writer": get_logger(),
+        "app": get_logger(),
+        "from_app": get_logger(),
+        "journal": get_logger(),
+        "kv_storage": get_logger(),
+        "vault": get_logger(),
+        "exec_logger": {
+            "handlers": ["basic", "routing"],
+            "level": "DEBUG",
+            "propagate": False,
+        },
+        "user_code": {
+            "handlers": ["basic", "routing"],
+            "level": "DEBUG",
+            "propagate": False,
+        },
+        "stdout": {
+            "handlers": ["stdout"],
+            "propagate": False,
+        }
+    }
+}
+
+logging.config.dictConfig(LOGGING_CONFIG)