relationalai 0.13.5__py3-none-any.whl → 1.0.0a2__py3-none-any.whl

relationalai/util/snowflake_handler.py

@@ -1,88 +0,0 @@
-import queue
-import threading
-import time
-
-from logging import Handler, LogRecord
-import traceback
-from uuid import UUID
-
-from relationalai.util.span_tracker import SpanTracker, record_is_span
-
-
-class SnowflakeHandler(Handler):
-    """
-    This is a logging handler that inserts spans directly into a Snowflake table.
-
-    It uses a queue and a worker thread to buffer spans and then insert them in batches.
-    """
-
-    def __init__(self, trace_id: UUID, snowflake_conn, sleep_interval_s=2):
-        super().__init__()
-        self.span_tracker = SpanTracker(trace_id, span_types_to_skip=None, log_span_attrs_as_str=True)
-        Handler.__init__(self)
-        self.snowflake_conn = snowflake_conn
-        self.queue = queue.Queue()
-        self.is_shut_down = False
-        self.sleep_interval_s = sleep_interval_s
-        self.worker_thread = threading.Thread(target=self._consume_loop)
-        self.worker_thread.start()
-        print(f'snowflake logger started with trace id {trace_id}')
-
-    def emit(self, record: LogRecord):
-        if record_is_span(record):
-            self.queue.put(record.msg)
-
-    def _consume_loop(self):
-        while True:
-            while not self.queue.empty():
-                try:
-                    batch = self._get_batch()
-                    self._send_batch(batch)
-                except Exception as e:
-                    print('snowflake logger error:', e)
-                    traceback.print_exc()
-            if self.is_shut_down:
-                return
-            time.sleep(self.sleep_interval_s)
-
-    def _get_batch(self):
-        batch = []
-        try:
-            while True:
-                item = self.queue.get_nowait()
-                batch.append(item)
-        except queue.Empty:
-            pass
-        return batch
-
-    def _send_batch(self, batch):
-        if not batch:
-            return
-
-        inserts = []
-
-        for event in batch:
-            if event["event"] == "span_start":
-                self.span_tracker.handle_span_start(event)
-            elif event["event"] == "span_end":
-                span_end_result = self.span_tracker.handle_span_end(event)
-                if span_end_result:
-                    formatted_span_msg = span_end_result[0]
-                    if formatted_span_msg:
-                        inserts.append(formatted_span_msg)
-        # execute
-        if inserts:
-            self.snowflake_conn.cursor().executemany(
-                """
-                INSERT INTO spans_raw (id, parent_id, trace_id, type, key, start_ts, finish_ts, attrs)
-                VALUES (%(id)s, %(parent_id)s, %(trace_id)s, %(type)s, %(key)s, %(start_timestamp)s, %(end_timestamp)s, %(attrs)s)
-                """,
-                inserts,
-            )
-
-    
-    def shut_down(self):
-        self.is_shut_down = True
-        print('snowflake logger: waiting to shut down...')
-        self.worker_thread.join()
-        print('snowflake logger: shut down')

relationalai/util/span_format_test.py

@@ -1,43 +0,0 @@
-import logging
-from typing import Any, Dict, List
-
-from relationalai.debugging import Span
-
-class SpanCollectorHandler(logging.Handler):
-    def __init__(self):
-        super().__init__()
-        self.root = None
-        self.nodes = {}
-        self.events = []
-
-    def emit(self, record):
-        if isinstance(record.msg, dict):
-            if record.msg['event'] == 'span_start':
-                span = record.msg['span']
-                self.nodes[str(span.id)] = span
-            else:
-                self.events.append(record.msg)
-
-def assert_valid_span_structure(spans: Dict[str, 'Span'], events: List[Dict[str, Any]]):
-    """
-    Assert that the tree of spans & events has a form that will
-    work for the debuggers.
-    """
-
-    transaction_created_events = [
-        event for event in events if event['event'] == 'transaction_created'
-    ]
-    transaction_spans = [
-        span for span in spans.values() if span.type == 'transaction'
-    ]
-
-    assert len(transaction_created_events) > 0
-    assert len(transaction_spans) > 0
-
-    # only one transaction_created event per transaction span
-    # TODO: Should this be equal?
-    assert len(transaction_created_events) <= len(transaction_spans)
-
-    # transaction_created spans have a txn_id attribute
-    for event in transaction_created_events:
-        assert 'txn_id' in event

relationalai/util/span_tracker.py

@@ -1,207 +0,0 @@
-import json
-from logging import LogRecord
-import os
-import subprocess
-from typing import List, Tuple, Union, Dict, Any
-from uuid import UUID, uuid4
-
-from relationalai.debugging import Span, encode_log_message, filter_span_attrs, logger
-from relationalai.util.constants import SPAN_ATTR_ALLOW_LIST, SPAN_TYPES_KEYS
-from relationalai.util.format import default_serialize
-from relationalai.clients.config import Config
-
-TRACE_ID = uuid4()
-
-class SpanTracker:
-    """
-    Abstract base class for tracking spans produced as log messages. 
-    Provides methods to handle the start and end of spans, format span messages, and manage span attributes.
-    """
-
-    def __init__(self, trace_id: UUID, span_types_to_skip=None, log_span_attrs_as_str=True):
-        self.trace_id = trace_id
-        self._set_span_types_to_skip(span_types_to_skip)
-        self._log_span_attrs_as_str = log_span_attrs_as_str
-        self.open_spans = {}
-        self.skipped_spans = set()
-
-    def handle_span_start(self, msg: dict) -> Union[None,Span]:
-        span: Span = msg["span"]
-        assert_key_exists_for_span_type(span)
-        if span.type in self.span_types_to_skip:
-            self.skipped_spans.add(str(span.id))
-            return None
-        self.open_spans[str(span.id)] = span
-        return span
-
-    def handle_span_end(self, msg: dict) -> Union[Tuple[dict, str, dict], None]:
-        if str(msg["id"]) in self.skipped_spans:
-            return None
-        span = self.open_spans.pop(msg["id"])
-        joined_key_path = get_span_path_as_str(span)
-        end_attrs = get_span_allowed_attributes_values(span, msg["end_attrs"])
-        combined_attributes = {
-            **span.attrs, # span_start attrs
-            **end_attrs, # span_end attrs
-            "key_path": joined_key_path,
-            "elapsed_s": (span.end_timestamp - span.start_timestamp).total_seconds(),
-        }
-        filtered_attrs = filter_span_attrs(combined_attributes)
-        formatted_span_msg = format_span_msg(self.trace_id, span, filtered_attrs, attrs_as_str=self._log_span_attrs_as_str)
-        return formatted_span_msg, joined_key_path, filtered_attrs
-
-    def _set_span_types_to_skip(self, span_types_to_skip):
-        if span_types_to_skip is None:
-            self.span_types_to_skip = set()
-        elif isinstance(span_types_to_skip, str):
-            self.span_types_to_skip = {span_types_to_skip}
-        elif isinstance(span_types_to_skip, (list, tuple)):
-            self.span_types_to_skip = set(span_types_to_skip)
-        elif isinstance(span_types_to_skip, set):
-            self.span_types_to_skip = span_types_to_skip
-        else:
-            raise ValueError(f"Invalid span_types_to_skip: {span_types_to_skip}, type: {type(span_types_to_skip)}")
-
-def format_span_json(trace_id: UUID, span: Span, attrs: dict, attrs_as_str:bool=True) -> dict:
-    span_key = get_span_value_from_key(span)
-    if attrs_as_str:
-        a = json.dumps(attrs, default=default_serialize)
-    else:
-        a = encode_log_message(attrs)
-
-    span_json = span.to_json()
-    extra_attrs = {
-        "trace_id": trace_id,
-        "key": str(span_key),
-        "attrs": a
-    }
-    formatted_span_msg = {**span_json, **extra_attrs}
-    return formatted_span_msg
-
-#region Span formatting and filtering functions
-def format_span_msg(trace_id: UUID, span: Span, attrs: dict, attrs_as_str:bool=True) -> dict:
-    span_key = get_span_value_from_key(span)
-    if attrs_as_str:
-        a = json.dumps(attrs, default=default_serialize)
-    else:
-        a = encode_log_message(attrs)
-
-    span_json = span.to_json()
-    extra_attrs = {
-        "trace_id": str(trace_id) if attrs_as_str else trace_id,
-        "key": str(span_key) if attrs_as_str else span_key,
-        "attrs": a
-    }
-    formatted_span_msg = {**span_json, **extra_attrs}
-    return formatted_span_msg
-
-def get_span_allowed_attributes_values(span: Span, attrs=None) -> dict:
-    # For an input span with attrs (or instead for an `attrs` dict), return a dict with only the allowed attributes for the span type
-    attrs = attrs or span.attrs
-    out = {}
-    allowed_attributes = SPAN_ATTR_ALLOW_LIST.get(span.type, set())
-    for attr in allowed_attributes:
-        if attr in attrs:
-            out[attr] = attrs[attr]
-    return out
-
-def assert_key_exists_for_span_type(span: Span):
-    if span.type in SPAN_TYPES_KEYS:
-        key = SPAN_TYPES_KEYS[span.type]
-        if key not in span.attrs:
-            raise ValueError(f"Attribute {key} is required for spans of type {span.type}, but wasn't provided. Provided attributes:\n{span.attrs}")
-
-def get_span_key_path(span: Span) -> str:
-    # Return the key path of the current span
-    value = get_span_value_from_key(span)
-    if value is None:
-        return span.type
-    return f"{span.type}({value})"
-
-def get_span_path_as_list(span: Span) -> List[Span]:
-    # Return the path from the root to the current span
-    path = []
-    cur = span
-    while cur is not None:
-        path.append(cur)
-        cur = cur.parent
-    path.reverse()
-    return path
-
-def get_span_path_as_str(span: Span) -> str:
-    # Return a string representation of the path from the root to the current span
-    # formerly format_key_path function
-    return '.'.join([get_span_key_path(span) for span in get_span_path_as_list(span)])
-
-def get_span_value_from_key(span: Span):
-    if span.type in SPAN_TYPES_KEYS:
-        key = SPAN_TYPES_KEYS[span.type]
-        if key in span.attrs:
-            return span.attrs[key]
-        raise ValueError(f"Key {key} not found in span of type {span.type}\n{span.attrs}")
-    return None
-
-def record_is_span(record: LogRecord) -> bool:
-    return (
-        isinstance(record.msg, dict) and
-        "event" in record.msg and
-        record.msg["event"] in ("span_start", "span_end")
-    )
-
-def _get_local_git_info() -> Dict[str, str]:
-    try:
-        branch = subprocess.check_output(['git', 'rev-parse', '--abbrev-ref', 'HEAD']).strip().decode('utf-8')
-        sha = subprocess.check_output(['git', 'rev-parse', 'HEAD']).strip().decode('utf-8')
-        repo = subprocess.check_output(['git', 'config', '--get', 'remote.origin.url']).strip().decode('utf-8').rsplit('/', maxsplit=1)[-1].replace('.git', '')
-        return {
-            'commit': sha,
-            'branch': branch,
-            'repo': repo,
-        }
-    except subprocess.CalledProcessError as e:
-        logger.warning(f"Could not get git info: {e}")
-        return {
-            'commit': 'local',
-            'branch': 'local',
-            'repo': 'local',
-        }
-
-def _get_git_info() -> Dict[str, str]:
-    if os.getenv("GITHUB_REF_NAME") is not None:
-        return {
-            'commit': os.getenv("GITHUB_SHA", ""),
-            'branch': os.getenv("GITHUB_REF_NAME", ""),
-            'repo': os.getenv("GITHUB_REPOSITORY", ""),
-        }
-    return _get_local_git_info()
-
-def _get_masked_config() -> Dict[str, Any]:
-    """Return a config with all sensitive keys (secrets and passwords) removed, such that it can be safely logged"""
-    config = Config()
-    masked_props = {
-        k: v for k, v in config.props.items()
-        if not isinstance(v, str)
-        or not any(prop_key_pattern in k for prop_key_pattern in ["secret", "password"])
-    }
-    return masked_props
-
-def _get_cloud_provider() -> Dict[str, str]:
-    provider = os.getenv("RAI_CLOUD_PROVIDER")
-    if not provider:
-        config = Config()
-        provider = config.get("platform")
-    assert provider, "Could not retrieve cloud provider from environment or config"
-    return {"platform": provider}
-
-def get_root_span_attrs(get_full_config: bool = False) -> Dict[str, Any]:
-    git_info = _get_git_info()
-    cloud_provider = _get_cloud_provider()
-    masked_config = _get_masked_config() if get_full_config else {}
-    attrs = {
-        **git_info,
-        **cloud_provider,
-        **masked_config
-    }
-    return attrs
-
-#endregion

relationalai/util/spans_file_handler.py

@@ -1,72 +0,0 @@
-from copy import deepcopy
-from logging import FileHandler
-import os
-import sys
-from uuid import UUID
-from typing import Dict, Union
-
-from relationalai.debugging import JsonFormatter
-from relationalai.util import get_timestamp
-from relationalai.util.format import str_to_bool
-from relationalai.util.span_tracker import SpanTracker, record_is_span
-
-# Span types that are automatically generated by PyRel but are not relevant for performance analysis
-SKIP_SPAN_TYPES = ['check_status']
-
-
-class SpansFileHandler(FileHandler):
-    """
-    This logger writes spans to a local .jsonl file inside the `spans_output` folder as they finish.
-    The primary goal is to ingest these spans into Snowflake for analytics.
-    """
-
-    def __init__(self, trace_id: UUID, filename: Union[str,None] = None):
-        self.span_tracker = SpanTracker(trace_id, span_types_to_skip=SKIP_SPAN_TYPES, log_span_attrs_as_str=False)
-        if filename is None:
-            filename = self._get_default_log_file_path()
-            if not os.path.exists(os.path.dirname(filename)):
-                os.makedirs(os.path.dirname(filename))
-        FileHandler.__init__(self, filename=filename)
-        self.setFormatter(JsonFormatter())
-
-    def emit(self, record):
-        if record_is_span(record):
-            msg: Dict = record.msg # type: ignore
-            if msg["event"] == "span_start":
-                self.span_tracker.handle_span_start(msg)
-            if msg["event"] == "span_end":
-                span_end_result = self.span_tracker.handle_span_end(msg)
-                if span_end_result:
-                    formatted_span_msg = span_end_result[0]
-                    if formatted_span_msg is None: # This will happen for span_types_to_skip
-                        return
-                    record_copy = deepcopy(record)  # Create a deep copy so as to not interfere with other handlers
-                    record_copy.msg = formatted_span_msg
-                    FileHandler.emit(self, record_copy)
-
-    def _get_default_log_file_path(self) -> str:
-        now_str = get_timestamp().strftime("%Y%m%d_%H%M%S")
-        trace_id_trunc = str(self.span_tracker.trace_id)[:8]
-        spans_subfolder = os.getenv("SPANS_SUBFOLDER", "")
-        return os.path.join(
-            os.getcwd(), "output_spans", spans_subfolder, f"{now_str}__{trace_id_trunc}.jsonl"
-        )
-
-
-def is_collect_only():
-    return "pytest" in sys.modules and "--collect-only" in sys.argv
-
-
-def maybe_add_spans_file_handler(logger, trace_id:UUID, if_not_set: bool = False):
-    """
-    Add a SpansFileHandler to the logger if the 'ENABLE_SPANS_FILE_HANDLER' environment variable is set to True,
-    or if the 'if_not_set' parameter is True.
-    """
-    if str_to_bool(os.getenv("ENABLE_SPANS_FILE_HANDLER")) or if_not_set:
-        if is_collect_only():
-            print("Running in collect-only mode, skipping SpansFileHandler")
-        elif any(isinstance(handler, SpansFileHandler) for handler in logger.handlers):
-            print("logger already has a SpansFileHandler instance, won't add another one")
-        else:
-            logger.debug("Adding SpansFileHandler to logger")
-            logger.addHandler(SpansFileHandler(trace_id))

relationalai/util/tracing_handler.py

@@ -1,34 +0,0 @@
-from logging import Handler
-from typing import Dict
-from uuid import UUID
-
-from relationalai.util.span_tracker import SpanTracker, get_span_allowed_attributes_values, get_span_path_as_str, record_is_span
-
-
-class TracingHandler(Handler):
-    """
-    This logger prints spans to stdout as they start and finish.
-    """
-
-    def __init__(self, trace_id: UUID):
-        super().__init__()
-        self.span_tracker = SpanTracker(trace_id, span_types_to_skip=None)
-
-    def emit(self, record):
-        if record_is_span(record):
-            msg: Dict = record.msg # type: ignore
-            if msg["event"] == "span_start":
-                span_to_print = self.span_tracker.handle_span_start(msg)
-                if span_to_print is not None:
-                    attrs = get_span_allowed_attributes_values(span_to_print)
-                    print("start", get_span_path_as_str(span_to_print), attrs)
-            elif msg["event"] == "span_end":
-                span_end_result = self.span_tracker.handle_span_end(msg)
-                if span_end_result:
-                    joined_key_path, filtered_attrs = span_end_result[1], span_end_result[2]
-                    if joined_key_path is not None and filtered_attrs is not None:
-                        # key_path is already printed separately, and name is already in the key_path
-                        filtered_attrs = {k: v for k, v in filtered_attrs.items() if k not in ["name", "key_path"]}
-                        # Exclude class instances, which can have a large nested structure to print
-                        filtered_attrs = {k: v for k, v in filtered_attrs.items() if not hasattr(v, "__dict__")}
-                        print("  end", joined_key_path, filtered_attrs)

relationalai-0.13.5.dist-info/METADATA

@@ -1,74 +0,0 @@
-Metadata-Version: 2.4
-Name: relationalai
-Version: 0.13.5
-Summary: RelationalAI Library and CLI
-Author-email: RelationalAI <support@relational.ai>
-License-File: LICENSE
-Requires-Python: >=3.10
-Requires-Dist: aiohttp
-Requires-Dist: click==8.2.1
-Requires-Dist: colorama
-Requires-Dist: cryptography
-Requires-Dist: gravis
-Requires-Dist: inquirerpy
-Requires-Dist: lqp==0.2.3
-Requires-Dist: nicegui==2.16.1
-Requires-Dist: numpy<2
-Requires-Dist: opentelemetry-api
-Requires-Dist: opentelemetry-exporter-otlp-proto-http
-Requires-Dist: opentelemetry-sdk
-Requires-Dist: packaging
-Requires-Dist: pandas
-Requires-Dist: pyarrow
-Requires-Dist: pyjwt
-Requires-Dist: requests
-Requires-Dist: rich
-Requires-Dist: snowflake-connector-python[secure-local-storage]
-Requires-Dist: snowflake-snowpark-python>=1.38.0
-Requires-Dist: toml
-Requires-Dist: tomlkit
-Requires-Dist: typing-extensions
-Requires-Dist: wcwidth
-Requires-Dist: websockets
-Provides-Extra: dev
-Requires-Dist: argcomplete; extra == 'dev'
-Requires-Dist: build; extra == 'dev'
-Requires-Dist: coverage; extra == 'dev'
-Requires-Dist: duckdb; extra == 'dev'
-Requires-Dist: faker; extra == 'dev'
-Requires-Dist: hupper; extra == 'dev'
-Requires-Dist: hypothesis; extra == 'dev'
-Requires-Dist: jupyter; extra == 'dev'
-Requires-Dist: logscope; extra == 'dev'
-Requires-Dist: pandas-stubs; extra == 'dev'
-Requires-Dist: pexpect; extra == 'dev'
-Requires-Dist: pip-tools; extra == 'dev'
-Requires-Dist: polars; extra == 'dev'
-Requires-Dist: prompt-toolkit; extra == 'dev'
-Requires-Dist: pytest; extra == 'dev'
-Requires-Dist: pytest-cov; extra == 'dev'
-Requires-Dist: pytest-mock; extra == 'dev'
-Requires-Dist: pytest-snapshot; extra == 'dev'
-Requires-Dist: pytest-timeout; extra == 'dev'
-Requires-Dist: ruff; extra == 'dev'
-Requires-Dist: scipy; extra == 'dev'
-Requires-Dist: sqlite-utils; extra == 'dev'
-Requires-Dist: sympy; extra == 'dev'
-Requires-Dist: tablescope; extra == 'dev'
-Requires-Dist: watchdog; extra == 'dev'
-Provides-Extra: ea
-Requires-Dist: dataclasses-json; extra == 'ea'
-Requires-Dist: owlready2; extra == 'ea'
-Requires-Dist: parameterized; extra == 'ea'
-Requires-Dist: pysmt; extra == 'ea'
-Requires-Dist: simple-ddl-parser; extra == 'ea'
-Requires-Dist: xmltodict; extra == 'ea'
-Provides-Extra: legacy
-Requires-Dist: rai-sdk>=0.7.5; extra == 'legacy'
-Description-Content-Type: text/markdown
-
-# The RelationalAI Python Library
-
-The Python library for building and querying knowledge graphs with RelationalAI.
-
-For more information, visit our documentation at [https://relational.ai/docs](https://relational.ai/docs).