relationalai 0.12.13__py3-none-any.whl → 0.13.0.dev0__py3-none-any.whl

relationalai/util/source.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+import ast
+from dataclasses import dataclass, field
+from functools import lru_cache
+import os
+import sys
+import textwrap
+import dis
+from bisect import bisect_right
+from types import CodeType
+
+cwd = os.getcwd()
+
+#------------------------------------------------------
+# Path utilities
+#------------------------------------------------------
+
+# relpath is a shockingly slow operation, so caching it
+@lru_cache(maxsize=128)
+def relative_path(path: str) -> str:
+    try:
+        return os.path.relpath(path, cwd)
+    except ValueError:
+        return path
+
+#------------------------------------------------------
+# Frame utilities
+#------------------------------------------------------
+
+_line_map_cache = {}
+
+def instruction_line(code, lasti):
+    pairs = _line_map_cache.get(code)
+    if pairs is None:
+        pairs = list(dis.findlinestarts(code))  # (addr, line_number)
+        _line_map_cache[code] = pairs
+
+    idx = bisect_right(pairs, (lasti, float("inf"))) - 1
+    if idx < 0:
+        return code.co_firstlineno
+    return pairs[idx][1]
+
+def first_non_relationalai_frame():
+    frame = sys._getframe(2)  # skip this function and the previous one
+    while frame and frame.f_back:
+        mod_name = frame.f_globals.get("__name__", "")
+        if mod_name.startswith("relationalai") or mod_name.startswith("v0.relationalai"):
+            frame = frame.f_back
+            continue
+        return frame
+    return frame
+
+#------------------------------------------------------
+# Source utilities
+#------------------------------------------------------
+
+@lru_cache(maxsize=128)
+def get_source_code(filename: str) -> str|None:
+    """Get the source code of a file."""
+    try:
+        with open(filename, "r") as f:
+            return f.read()
+    except Exception:
+        return None
+
+parses = {}
+def find_root_expression(source_code: str, target_line: int, relative_filename: str):
+    # Parse the source code into an AST
+    if relative_filename not in parses:
+        parses[relative_filename] = ast.parse(source_code)
+    tree = parses[relative_filename]
+
+    class ExpressionChainFinder(ast.NodeVisitor):
+        def __init__(self, target_line):
+            self.target_line = target_line
+            self.candidates = []
+
+        def visit(self, node):
+            # Consider all expression nodes, including Call nodes
+            is_expr_node = isinstance(node, (ast.Expr, ast.Call, ast.Attribute, ast.BinOp,
+                                          ast.UnaryOp, ast.IfExp, ast.Compare, ast.Subscript,
+                                          ast.Dict, ast.List, ast.Tuple, ast.Set))
+
+            if hasattr(node, "lineno") and is_expr_node:
+                end_lineno = getattr(node, "end_lineno", node.lineno)
+
+                # Check if this node contains our target line
+                if node.lineno <= self.target_line <= end_lineno:
+                    # Store the node with its line span information
+                    self.candidates.append((node, node.lineno, end_lineno))
+
+            # Continue visiting children
+            for child in ast.iter_child_nodes(node):
+                self.visit(child)
+
+    finder = ExpressionChainFinder(target_line)
+    finder.visit(tree)
+
+    if finder.candidates:
+        # Sort candidates by their starting line (ascending) and then by span size (descending)
+        # This prioritizes expressions that start earlier and cover more lines
+        sorted_candidates = sorted(finder.candidates,
+                                  key=lambda x: (x[1], -(x[2] - x[1])))
+
+        # Find the outermost expression that contains our target line
+        # This will capture the entire multi-line function call
+        outermost_candidates = []
+        for node, start_line, end_line in sorted_candidates:
+            # Check if this is a new candidate or extends an existing one
+            if not outermost_candidates or start_line < outermost_candidates[0][1]:
+                outermost_candidates = [(node, start_line, end_line)]
+
+            # If we have a node that starts at the same line but ends later, prefer it
+            elif start_line == outermost_candidates[0][1] and end_line > outermost_candidates[0][2]:
+                outermost_candidates = [(node, start_line, end_line)]
+
+        if outermost_candidates:
+            node, start_line, end_line = outermost_candidates[0]
+
+            # Extract the lines from the source code
+            block_lines = source_code.splitlines()[start_line - 1:end_line]
+            block_code = "\n".join(block_lines)
+            return SourcePos(relative_filename, _line=start_line, _source=textwrap.dedent(block_code))
+
+    # If no nodes were found, return the single line
+    lines = source_code.splitlines()
+    if target_line > len(lines):
+        return SourcePos(relative_filename, _line=target_line)
+    return SourcePos(relative_filename, _line=target_line, _source=lines[target_line - 1])
+
+#------------------------------------------------------
+# SourcePos
+#------------------------------------------------------
+
+counts = 0
+
+@dataclass(slots=True)
+class SourcePos:
+    file: str = field(default="Unknown")
+    _block:SourcePos|None = None
+    _lasti: int|None = field(default=0)
+    _pycode: CodeType|None = field(default=None)
+    _source: str|None = field(default=None)
+    _line: int|None = field(default=None)
+
+    @classmethod
+    def new(cls):
+        caller_frame = first_non_relationalai_frame()
+        if not caller_frame:
+            return cls()
+
+        filename = caller_frame.f_code.co_filename
+        lasti = caller_frame.f_lasti
+        pycode = caller_frame.f_code
+
+        return cls(filename, _lasti=lasti, _pycode=pycode)
+
+    @property
+    def line(self):
+        if self._line is None and self._pycode:
+            self._line = instruction_line(self._pycode, self._lasti)
+        return self._line or 0
+
+    @property
+    def source(self):
+        if self._source is None and self.file:
+            self._source = get_source_code(self.file)
+        return self._source or ""
+
+    @property
+    def block(self):
+        if not self.source or not self.line:
+            return SourcePos()
+        if not self._block:
+            self._block = find_root_expression(self.source, self.line, self.file)
+        return self._block
+
+    def transform(self, transformer:ast.NodeTransformer) -> SourcePos|None:
+        if not self.block or not self.block.source:
+            return
+
+        new_ast = transformer.visit(ast.parse(self.block.source))
+        new_source = ast.unparse(new_ast)
+        new = SourcePos(self.file, _line=self.line, _source=new_source)
+        return new

relationalai/util/structures.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+from collections.abc import Iterable, Iterator, MutableSet
+from typing import Callable, Generic, Hashable, TypeVar
+
+T = TypeVar("T")
+
+#------------------------------------------------------
+# OrderedSet
+#------------------------------------------------------
+
+class OrderedSet(Generic[T], MutableSet[T]):
+    __slots__ = ("_dict",)
+    _dict: dict[T, None]
+
+    def __init__(self, iterable: Iterable[T] | None = None) -> None:
+        self._dict = {}
+        if iterable is not None:
+            self.update(iterable)
+
+    def __contains__(self, item: object) -> bool:  # type: ignore[override]
+        return item in self._dict
+
+    def __iter__(self) -> Iterator[T]:
+        return iter(self._dict.keys())
+
+    def __len__(self) -> int:
+        return len(self._dict)
+
+    def __getitem__(self, index: int) -> T:
+        if index < 0:
+            index += len(self._dict)
+        if index < 0 or index >= len(self._dict):
+            raise IndexError("OrderedSet index out of range")
+        for i, key in enumerate(self._dict):
+            if i == index:
+                return key
+        # This point should never be reached
+        raise IndexError("OrderedSet index out of range")
+
+    def add(self, value: T) -> None:
+        if value not in self._dict:
+            self._dict[value] = None
+
+    def discard(self, value: T) -> None:
+        self._dict.pop(value, None)
+
+    def update(self, iterable: Iterable[T]) -> None:
+        for x in iterable:
+            self.add(x)
+
+    def first(self) -> T | None:
+        for x in self._dict:
+            return x
+        return None
+
+    def try_replace(self, old: T, new: T) -> None:
+        """Replace a potentially existing item with a new one, keeping its position if possible."""
+        if new in self._dict or old not in self._dict:
+            return
+        self._dict = { (new if k == old else k): None for k in self._dict }
+
+    def __repr__(self) -> str:
+        items = ", ".join(repr(x) for x in self._dict)
+        return f"OrderedSet([{items}])"
+
+#------------------------------------------------------
+# KeyedDict
+#------------------------------------------------------
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+class KeyedDict(Generic[K, V]):
+    __slots__ = ("_dict", "_key")
+    _dict: dict[Hashable, tuple[K, V]]
+    _key: Callable[[K], Hashable]
+
+    def __init__(self, key: Callable[[K], Hashable], iterable: Iterable[tuple[K, V]] | None = None) -> None:
+        self._dict = {}
+        self._key = key
+        for k, v in iterable or ():
+            self._dict[self._key(k)] = (k, v)
+
+    def __setitem__(self, k: K, v: V) -> None:
+        self._dict[self._key(k)] = (k, v)
+
+    def __getitem__(self, k: K) -> V:
+        return self._dict[self._key(k)][1]
+
+    def __delitem__(self, k: K) -> None:
+        del self._dict[self._key(k)]
+
+    def __contains__(self, k: object) -> bool:
+        try: return self._key(k) in self._dict  # type: ignore[arg-type]
+        except Exception: return False
+
+    def __len__(self) -> int:
+        return len(self._dict)
+
+    def get(self, k: K, default: V | None = None) -> V | None:
+        entry = self._dict.get(self._key(k))
+        return default if entry is None else entry[1]
+
+    def items(self) -> Iterator[tuple[K, V]]:
+        yield from self._dict.values()
+
+    def values(self) -> Iterator[V]:
+        for _, v in self._dict.values():
+            yield v
+
+    def update(self, iterable: KeyedDict[K, V]|Iterable[tuple[K, V]]) -> None:
+        if isinstance(iterable, KeyedDict):
+            self._dict.update(iterable._dict)
+        else:
+            for k, v in iterable:
+                self._dict[self._key(k)] = (k, v)
+
+
+#------------------------------------------------------
+# KeyedSet
+#------------------------------------------------------
+
+class KeyedSet(Generic[K], MutableSet[K]):
+    __slots__ = ("_dict", "_key")
+    _dict: dict[Hashable, K]
+    _key: Callable[[K], Hashable]
+
+    def __init__(self, key: Callable[[K], Hashable], iterable: Iterable[K] | None = None) -> None:
+        self._dict = {}
+        self._key = key
+        if iterable is not None:
+            self.update(iterable)
+
+    def __contains__(self, item: K) -> bool:  # type: ignore[override]
+        return self._key(item) in self._dict
+
+    def __iter__(self) -> Iterator[K]:
+        return iter(self._dict.values())
+
+    def __len__(self) -> int:
+        return len(self._dict)
+
+    def add(self, value: K) -> None:
+        if self._key(value) not in self._dict:
+            self._dict[self._key(value)] = value
+
+    def discard(self, value: K) -> None:
+        self._dict.pop(self._key(value), None)
+
+    def update(self, iterable: Iterable[K]) -> None:
+        for x in iterable:
+            self.add(x)
+
+    def first(self) -> K | None:
+        for _, v in self._dict.items():
+            return v
+        return None
+
+    def __repr__(self) -> str:
+        items = ", ".join(repr(x) for x in self._dict)
+        return f"KeyedSet([{items}])"
+
+

relationalai/util/tracing.py

@@ -0,0 +1,261 @@
+# tracer.py — stdlib-only OTEL-shaped tracer with live span_start + 200ms write throttling
+from __future__ import annotations
+import os, json, time, secrets, threading, atexit
+from contextlib import ContextDecorator, asynccontextmanager
+from contextvars import ContextVar
+from typing import Optional, Callable, Any, Dict
+
+# ---------- utils ----------
+def _hex(bits: int) -> str: return secrets.token_hex(bits // 8)
+def now_ns() -> int: return time.time_ns()
+
+# Each logical execution path (thread/task) has its own span stack
+_current_stack: ContextVar[list["Span"]] = ContextVar("_current_stack", default=[])
+
+# ---------- core ----------
+class Span:
+    """
+    Represents a span; use Tracer.span(...) / Tracer.aspan(...) to create via context manager.
+    """
+    __slots__ = ("tracer","name","kind","trace_id","span_id","parent_span_id",
+                 "start","end","attrs","events","status","_finished")
+
+    def __init__(self, tracer:"Tracer", name:str, kind:str, parent:"Span|None", **attrs):
+        self.tracer = tracer
+        self.name = name
+        self.kind = kind
+        self.trace_id = parent.trace_id if parent else _hex(128)
+        self.span_id = _hex(64)
+        self.parent_span_id = parent.span_id if parent else None
+        self.start = now_ns()
+        self.end: Optional[int] = None
+        self.attrs: Dict[str, Any] = dict(attrs)
+        self.events: list[dict] = []
+        self.status = "UNSET"
+        self._finished = False
+
+    # ---- span APIs ----
+    def add_event(self, name:str, **attributes):
+        self.events.append({"name": name, "ts": now_ns(), "attributes": attributes})
+
+    def set_status(self, ok: bool = True, message: Optional[str] = None):
+        self.status = "OK" if ok else "ERROR"
+        if message:
+            self.attrs["error.message"] = message
+
+    def record_exception(self, exc: BaseException):
+        self.add_event("exception",
+                       type=exc.__class__.__name__,
+                       message=str(exc),
+                       module=getattr(exc.__class__, "__module__", ""))
+        self.set_status(False, str(exc))
+
+    def finish(self):
+        if self._finished:
+            return
+        self._finished = True
+        self.end = now_ns()
+        self.tracer._emit_full_span(self)
+
+    def __setitem__(self, key, value):
+        self.attrs[key] = value
+
+class Tracer(ContextDecorator):
+    """
+    Tracer with:
+      - context-managed spans (sync & async)
+      - automatic parent linkage via ContextVar stack
+      - live 'span_start' records + batched/throttled writes
+      - OTEL-shaped final span records
+    """
+    def __init__(
+        self,
+        out_path: str = "spans.jsonl",
+        service: str = "rel-modeler",
+        flush_interval_ms: int = 200,
+        emit_span_start: bool = True,
+    ):
+        self.out_path = out_path
+        self.service = service
+        self.emit_span_start = emit_span_start
+        os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
+
+        self._program_span: Span | None = None
+
+        # Buffered/throttled writer
+        self._lock = threading.Lock()
+        self._buffer: list[str] = []
+        self._last_flush_ns: int = 0
+        self._flush_interval_ns: int = flush_interval_ms * 1_000_000
+        atexit.register(self.flush)
+
+    # ---- context manager for sync code ----
+    def span(self, name:str, *, kind:str="INTERNAL", **attrs):
+        tracer = self
+
+        class SpanCtx:
+            def __enter__(self):
+                stack = _current_stack.get()
+                parent = stack[-1] if stack else (tracer._program_span if tracer._program_span and not tracer._program_span._finished else None)
+                self.span = Span(tracer, name, kind, parent, **attrs)
+                _current_stack.set(stack + [self.span])
+                if tracer.emit_span_start:
+                    tracer._emit_span_start(self.span)
+                return self.span
+
+            def __exit__(self, exc_type, exc, tb):
+                try:
+                    if exc_type is not None:
+                        self.span.record_exception(exc)  # marks ERROR
+                    elif self.span.status == "UNSET":
+                        self.span.set_status(True)
+                finally:
+                    self.span.finish()
+                    # pop this span
+                    cur = _current_stack.get()
+                    if cur and cur[-1] is self.span:
+                        _current_stack.set(cur[:-1])
+                # do not suppress exceptions
+                return False
+
+        return SpanCtx()
+
+    # ---- async context manager ----
+    @asynccontextmanager
+    async def aspan(self, name:str, *, kind:str="INTERNAL", **attrs):
+        stack = _current_stack.get()
+        parent = stack[-1] if stack else (self._program_span if self._program_span and not self._program_span._finished else None)
+        span = Span(self, name, kind, parent, **attrs)
+        _current_stack.set(stack + [span])
+        if self.emit_span_start:
+            self._emit_span_start(span)
+        try:
+            yield span
+            if span.status == "UNSET":
+                span.set_status(True)
+        except BaseException as e:
+            span.record_exception(e)
+            raise
+        finally:
+            span.finish()
+            cur = _current_stack.get()
+            if cur and cur[-1] is span:
+                _current_stack.set(cur[:-1])
+
+    # ---- helpers ----
+    def current_span(self) -> Optional[Span]:
+        stack = _current_stack.get()
+        parent = stack[-1] if stack else (self._program_span if self._program_span and not self._program_span._finished else None)
+        return parent
+
+    def add_event(self, name:str, **attrs):
+        s = self.current_span()
+        if s:
+            s.add_event(name, **attrs)
+
+    def set_status(self, ok:bool=True, message:str|None=None):
+        s = self.current_span()
+        if s:
+            s.set_status(ok, message)
+
+    def traceparent(self) -> str:
+        s = self.current_span()
+        if not s: return ""
+        # version(00)-trace_id-span_id-flags(01)
+        return f"00-{s.trace_id}-{s.span_id}-01"
+
+    def start_program(self, name: str = "program", *, kind: str = "INTERNAL", **attrs) -> Span:
+        """Open one root span for the whole process; auto-finished via atexit."""
+        if self._program_span and not self._program_span._finished:
+            return self._program_span  # already started
+        root = Span(self, name, kind, parent=None, **attrs)
+        self._program_span = root
+        if getattr(self, "emit_span_start", False):  # works if you added span_start earlier
+            self._emit_span_start(root)
+        # close on interpreter shutdown
+        def _close():
+            s = self._program_span
+            if s and not s._finished:
+                if s.status == "UNSET":
+                    s.set_status(True)
+                s.finish()
+        atexit.register(_close)
+        return root
+
+    # ---- emitters (buffered) ----
+    def _emit_span_start(self, span: Span):
+        doc = {
+            "type": "span_start",
+            "trace_id": span.trace_id,
+            "span_id": span.span_id,
+            "parent_span_id": span.parent_span_id,
+            "name": span.name,
+            "kind": span.kind,
+            "start_time_unix_nano": span.start,
+            "resource": {"service.name": self.service},
+            "attributes": span.attrs,  # snapshot of initial attrs
+        }
+        self._write(json.dumps(doc, separators=(",", ":")) + "\n")
+
+    def _emit_full_span(self, span: Span):
+        doc = {
+            "trace_id": span.trace_id,
+            "span_id": span.span_id,
+            "parent_span_id": span.parent_span_id,
+            "name": span.name,
+            "kind": span.kind,
+            "start_time_unix_nano": span.start,
+            "end_time_unix_nano": span.end,
+            "status": span.status,
+            "attributes": span.attrs,
+            "events": span.events,
+            "resource": {"service.name": self.service},
+        }
+        self._write(json.dumps(doc, separators=(",", ":")) + "\n")
+
+    def _write(self, line: str):
+        now = now_ns()
+        with self._lock:
+            self._buffer.append(line)
+            if (now - self._last_flush_ns) >= self._flush_interval_ns:
+                self._flush_locked(now)
+
+    def flush(self):
+        with self._lock:
+            self._flush_locked(now_ns())
+
+    def _flush_locked(self, now: int):
+        if not self._buffer:
+            self._last_flush_ns = now
+            return
+        # Single append write for all buffered lines
+        with open(self.out_path, "a", encoding="utf-8") as f:
+            f.writelines(self._buffer)
+        self._buffer.clear()
+        self._last_flush_ns = now
+
+# ---------- decorator ----------
+def traced(name: str | None = None, *, kind: str = "INTERNAL"):
+    def _wrap(fn: Callable):
+        nm = name or fn.__name__
+        def inner(*a, **kw):
+            tracer = get_tracer()
+            with tracer.span(nm, kind=kind):
+                return fn(*a, **kw)
+        return inner
+    return _wrap
+
+# ---------- global default (optional) ----------
+_DEFAULT_TRACER: Tracer | None = None
+def set_tracer(t: Tracer):  # call once at app start
+    global _DEFAULT_TRACER
+    _DEFAULT_TRACER = t
+def get_tracer() -> Tracer:
+    if _DEFAULT_TRACER is None:
+        set_tracer(Tracer())  # lazy init to spans.jsonl
+    assert _DEFAULT_TRACER is not None
+    return _DEFAULT_TRACER
+
+def span(name:str, *, kind:str="INTERNAL", **attrs):
+    tracer = get_tracer()
+    return tracer.span(name, kind=kind, **attrs)

relationalai-0.13.0.dev0.dist-info/METADATA

@@ -0,0 +1,46 @@
+Metadata-Version: 2.4
+Name: relationalai
+Version: 0.13.0.dev0
+Summary: RelationalAI Library and CLI
+Author-email: RelationalAI <support@relational.ai>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: lqp==0.1.19
+Requires-Dist: snowflake-connector-python[secure-local-storage]
+Requires-Dist: snowflake-snowpark-python>=1.38.0
+Requires-Dist: typing-extensions
+Requires-Dist: rich
+Requires-Dist: nicegui>3.0.0
+Requires-Dist: numpy<2
+Requires-Dist: pandas
+Requires-Dist: colorama
+Requires-Dist: duckdb
+Requires-Dist: inquirerpy
+Requires-Dist: click==8.2.1
+Requires-Dist: gravis
+Requires-Dist: scipy
+Requires-Dist: toml
+Requires-Dist: tomlkit
+Requires-Dist: requests
+Requires-Dist: pyarrow
+Requires-Dist: websockets
+Requires-Dist: aiohttp
+Requires-Dist: packaging
+Requires-Dist: opentelemetry-sdk
+Requires-Dist: opentelemetry-api
+Requires-Dist: opentelemetry-exporter-otlp-proto-http
+Requires-Dist: wcwidth
+Requires-Dist: cryptography
+Requires-Dist: pyjwt
+Requires-Dist: more-itertools
+Requires-Dist: sqlglot
+Requires-Dist: pyyaml
+Requires-Dist: pydantic<2.12.0,>=2.0.0
+Requires-Dist: pydantic-settings<2.11.0,>=2.0.0
+Provides-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"