datajet-compass 0.1.0__py3-none-any.whl

datajet_compass/config.py

@@ -0,0 +1,11 @@
+import os
+from dotenv import load_dotenv, find_dotenv
+
+
+# find_dotenv() walks UP the directory tree from cwd until it finds a .env file.
+# This works regardless of where the annotators library is installed or imported from.
+# usecwd=True ensures the search starts from the caller's working directory, not this file's location.
+load_dotenv(find_dotenv(usecwd=True), override=False)
+
+LOG_PATH = os.getenv("LOG_PATH", None)
+ENABLE_ONLINE_COMPASS_EVAL = True if os.getenv("ENABLE_ONLINE_COMPASS_EVAL", "false").lower() in ["true", "yes", "1"] else False

datajet_compass/context.py

@@ -0,0 +1,3 @@
+import contextvars
+
+_state_ctx = contextvars.ContextVar("_parent_state")

datajet_compass/decorators.py

@@ -0,0 +1,50 @@
+import asyncio
+import functools
+
+from .config import ENABLE_ONLINE_COMPASS_EVAL
+from .trace import Trace
+
+
+def _make_decorator(call_type) :
+    def decorator(func):
+        
+        @functools.wraps(func)
+        def _sync_wrapper(*args, **kwargs):
+            if not ENABLE_ONLINE_COMPASS_EVAL:
+                return func(*args, **kwargs)
+            
+            with Trace(call_type=call_type) as trace:
+                trace.start_trace(function_name=func.__name__, request={"args": args, "kwargs": kwargs})
+                
+                try:
+                    result = func(*args, **kwargs)
+                    return trace.consolidate_trace(output=result, status="SUCCESS")
+                except Exception as e:
+                    return trace.consolidate_trace(status="FAILED", error=e)
+        
+        @functools.wraps(func)
+        async def _async_wrapper(*args, **kwargs):
+            if not ENABLE_ONLINE_COMPASS_EVAL:
+                return await func(*args, **kwargs)
+            
+            with Trace(call_type=call_type) as trace:
+                trace.start_trace(function_name=func.__name__, request={"args": args, "kwargs": kwargs})
+                
+                try:
+                    result = await func(*args, **kwargs)
+                    return trace.consolidate_trace(output=result, status="SUCCESS")
+                except Exception as e:
+                    return trace.consolidate_trace(status="FAILED", error=e)
+    
+        return _async_wrapper if asyncio.iscoroutinefunction(func) else _sync_wrapper
+    
+    return decorator
+
+
+# Syntactic Sugar wrappers matching standard trace requirements
+observe = _make_decorator("function")
+
+observe_agent = _make_decorator("agent")
+
+observe_tool = _make_decorator("tool")
+        

datajet_compass/executors.py

@@ -0,0 +1,76 @@
+import multiprocessing
+from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
+import contextvars
+
+from .context import _state_ctx
+from .log_entry import LogEntry
+
+
+class ContextThreadPoolExecutor(ThreadPoolExecutor):
+    """Thread pool that automatically propagates ContextVar state to worker threads."""
+
+    def submit(self, fn, *args, **kwargs):
+        ctx = contextvars.copy_context()
+        return super().submit(ctx.run, fn, *args, **kwargs)
+
+
+class ContextProcessPoolExecutor(ProcessPoolExecutor):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._manager = multiprocessing.Manager()
+        self._child_states = self._manager.list()   # each slot holds ONE CallEntry dict
+        self._invocation_counter = 0
+
+        try:
+            self._parent_state = _state_ctx.get()
+            self._serialized_parent = self._parent_state.to_object_dict()
+        except LookupError:
+            self._parent_state = None
+            self._serialized_parent = None
+
+    def submit(self, fn, *args, **kwargs):
+        self._invocation_counter += 1
+        invocation_order = self._invocation_counter
+        self._child_states.append({})
+
+        return super().submit(
+            _run_with_context,
+            fn,
+            self._serialized_parent,
+            self._child_states,
+            invocation_order,
+            *args,
+            **kwargs
+        )
+
+    def __exit__(self, *args):
+        result = super().__exit__(*args)
+
+        # All processes finished — each slot in _child_states is ONE LogEntry dict
+        if self._parent_state is not None:
+            for call_dict in self._child_states:
+                if call_dict:
+                    entry = LogEntry.from_dict(call_dict)
+                    self._parent_state.functionCalls.append(entry)
+
+        self._manager.shutdown()
+        return result
+
+
+def _run_with_context(fn, serialized_state, child_states_list, invocation_order, *args, **kwargs):
+    """Runs inside each worker process."""
+    if serialized_state:
+        state = LogEntry.from_dict(serialized_state)
+        _state_ctx.set(state)
+
+    result = fn(*args, **kwargs)
+
+    if serialized_state:
+        # Always exactly 1 top-level child — store it directly, no loop needed
+        child_call = state.functionCalls[0]
+        child_call.order = invocation_order
+        child_states_list[invocation_order - 1] = child_call.to_object_dict()
+
+    return result
+

datajet_compass/file_writter.py

@@ -0,0 +1,37 @@
+import json
+import os
+from typing import Dict, Any
+
+
+class JSONLWriter:
+    """
+    Writes annotation records to a JSONL file.
+    Each execution record is written as ONE JSON object per line.
+    """
+    
+    def __init__(self, file_path: str, indent: int = 2):
+        """
+        :param file_path: Path to the .jsonl file
+        :param indent: JSON indentation (default = 2 for readable output)
+        """
+        self.file_path = file_path
+        self.indent = indent
+    
+    def _serialize(self, obj: Any) -> Any:
+        # Any object with __dict__ (custom classes)
+        if hasattr(obj, "__dict__"):
+            return obj.__dict__
+        # Fallback — convert to string
+        return str(obj)
+    
+    def write(self, record: Dict[str, Any]) -> None:
+        """
+        Write a single execution record to the JSONL file.
+        """
+        
+        os.makedirs(os.path.dirname(self.file_path), exist_ok=True)
+        with open(self.file_path, "a", encoding="utf-8") as f:
+            f.write(
+                json.dumps(record, ensure_ascii=False, indent=self.indent, default=self._serialize)
+                + "\n"
+            )

datajet_compass/log_entry.py

@@ -0,0 +1,115 @@
+from threading import Lock
+from datetime import datetime
+from typing import Dict, Any
+
+class LogEntry:
+    date_fmt = "%Y-%m-%d %H:%M:%S"
+    
+    def __init__(self, testcase_id=None, run_id=None, request=None, type=None, function_name=None, url=None, order=1):
+        self._lock = Lock()
+        self.call_counter = 1
+        
+        self.testcase_id = testcase_id
+        self.run_id = run_id
+        self.name = function_name
+        self.type = type
+        self.target = url
+        self.order = order
+        self.input = request
+        self.output = None
+        self.status = "PROCESSING"
+        self.error = None
+        self.functionCalls = []
+        self.start_time = datetime.now().strftime(LogEntry.date_fmt)
+        self.end_time = None
+    
+    def get_count(self):
+        with self._lock:
+            self.call_counter += 1
+            self.functionCalls.append({})
+        
+        return self.call_counter - 1
+        
+    def set_response(self, output=None, status=None, error=None):
+        self.end_time = datetime.now().strftime(LogEntry.date_fmt)
+        self.output = output
+        self.status = status
+        self.error = error
+        
+    def to_dict(self):
+        
+        start = datetime.strptime(self.start_time, LogEntry.date_fmt)
+        end = datetime.strptime(self.end_time, LogEntry.date_fmt)
+        duration = end - start
+        
+        return_dict = {
+            "testcase_id": self.testcase_id,
+            "run_id": self.run_id,
+            "name": self.name,
+            "order": self.order,
+            "execution": {
+                "type": self.type,
+                "target": self.target,
+                "status": self.status,
+                "error": self.error,
+                "execution_time": f"{duration.total_seconds()} s"
+            },
+            "input": str(self.input),
+            "output": str(self.output),
+            "functionCall": [call.to_dict() for call in self.functionCalls if call.type == "function"],
+            "agentCall": [call.to_dict() for call in self.functionCalls if call.type == "agent"],
+            "toolCall": [call.to_dict() for call in self.functionCalls if call.type == "tool"]
+        }
+        
+        keys_to_remove = [
+            ('testcase_id', self.testcase_id),
+            ('run_id', self.run_id),
+            ('order', self.order),
+            ('target', self.target)
+        ]
+        for key, value in keys_to_remove:
+            if value is None:
+                return_dict.pop(key, None)
+                
+        return {k: v for k, v in return_dict.items() if v != []}
+    
+    def to_object_dict(self):
+        return {
+            "testcase_id": self.testcase_id,
+            "run_id": self.run_id,
+            "type": self.type,
+            "target": self.target,
+            "name": self.name,
+            "order": self.order,
+            "status": self.status,
+            "error": self.error,
+            "start_time": self.start_time,
+            "end_time": self.end_time,
+            "input": str(self.input),
+            "output": str(self.output),
+            "functionCalls": [call.to_object_dict() for call in self.functionCalls]
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'LogEntry':
+        data = dict(data)
+        calls_data = data.pop("functionCalls", [])
+        
+        # to_object_dict fields -> __init__ fields
+        obj = cls(
+            testcase_id=data.get("testcase_id", None),
+            run_id=data.get("run_id", None),
+            request=data.get("request", None),
+            type=data.get("type", None),
+            function_name=data.get("name", None),
+            url=data.get("target", None),
+            order=data.get("order", 1)
+        )
+        # restore live state fields
+        obj.status = data["status"]
+        obj.error = data["error"]
+        obj.output = data["output"]
+        obj.start_time = data["start_time"]
+        obj.end_time = data["end_time"]
+        obj.functionCalls = [LogEntry.from_dict(c) for c in calls_data]
+        return obj

datajet_compass/trace.py

@@ -0,0 +1,69 @@
+import threading
+from .context import _state_ctx
+from .file_writter import JSONLWriter
+from .log_entry import LogEntry
+from .config import LOG_PATH, ENABLE_ONLINE_COMPASS_EVAL
+
+
+class Trace:
+    
+    def __init__(self, call_type = "api"):
+        self.call_type = call_type
+        self.token = None
+        self.parent = None
+        self.current_state = None
+        self._lock = threading.Lock()
+        
+    def start_trace(self, function_name = None, request = None, url = None):
+        
+        if not ENABLE_ONLINE_COMPASS_EVAL:
+            return None, None
+        
+        try:
+            self.parent = _state_ctx.get()
+        except Exception as e:
+            self.parent = None
+        
+        with threading.Lock():
+            testcase_id = request.get("testcase_id", None) if isinstance(request, dict) else None
+            run_id = request.get("run_id", None) if isinstance(request, dict) else None
+            
+            # - 1 is added because a dummy of the current call has been added to the parent's functionCalls list on get_count()
+            order_count = self.parent.get_count() if self.parent else None
+            
+            state = LogEntry(testcase_id=testcase_id, run_id=run_id, request=request, type=self.call_type, function_name=function_name, url=url, order=order_count)
+            
+            if self.parent is not None:
+                self.parent.functionCalls[order_count - 1] = state
+        
+        self.token = _state_ctx.set(state)
+        self.current_state = state
+        
+        return self.token, self.current_state
+        
+    def consolidate_trace(self,output=None, status=None, error=None):
+        
+        if ENABLE_ONLINE_COMPASS_EVAL:
+            self.current_state.set_response(output=output, status=status, error=error)
+            
+            if self.parent is None and self.current_state.testcase_id is not None and self.current_state.run_id is not None:
+                print(self.current_state.to_dict())
+                return {"actual_response": output if error is None else error, "trace": self.current_state.to_dict()}
+            elif self.parent is None and self.current_state.testcase_id is None and self.current_state.run_id is None and LOG_PATH is not None:
+                with self._lock:
+                    writer = JSONLWriter(LOG_PATH)
+                    writer.write(self.current_state.to_dict())
+        
+        return output if error is None else error
+
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.token:
+            _state_ctx.reset(self.token)
+            self.token = None
+        return False
+
+    
+

datajet_compass-0.1.0.dist-info/METADATA

@@ -0,0 +1,21 @@
+Metadata-Version: 2.4
+Name: datajet-compass
+Version: 0.1.0
+Summary: Observability library for AI agents and tools
+Author-email: Basil K <Basil.Kalathodi@ibsplc.com>, Rose Baiju Malakaran <Rose.Malakaran@ibsplc.com>, Liana S <Liana.Salim@ibsplc.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/imaginebediscovered/datajet-compass
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+
+# DataJet Compass
+
+Observability library for AI agents and tools.
+
+## Installation
+
+```bash
+pip install datajet-compass

datajet_compass-0.1.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+datajet_compass/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+datajet_compass/config.py,sha256=W7CtH0ph-Rcy1CjgSsjv9XwkjlkblJMysE1ftzUkvHw,550
+datajet_compass/context.py,sha256=x6yFXkZINmzFGBhVRX1-X7KvA5F8dLbMXu1JhGFQvlk,76
+datajet_compass/decorators.py,sha256=UJhjsEXNSMl_f7eUR98f4mfm9pNGLlX4S_vtVYAjflU,1820
+datajet_compass/executors.py,sha256=_OAPi8p4AuxedtuOjqxJDPDQwr6-2c0GFs3nqtIyFy8,2547
+datajet_compass/file_writter.py,sha256=6JY5e2dWWzx2OImxgH_OhMnvB97PgpH_-W9mKU_jPNM,1183
+datajet_compass/log_entry.py,sha256=HeNqP5w5EEcnZW0G3K0dOVZ0oXCJmIw3hYOPcQ_UvTQ,4166
+datajet_compass/trace.py,sha256=jrMykjDqOOCjgK8rlQuK6nQGqlBmBQdNjduRQoxmHqE,2740
+datajet_compass-0.1.0.dist-info/METADATA,sha256=DRHA08ahbwInT4fcDrndK2CG9xcFC1JGPPWMkDxou0w,693
+datajet_compass-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+datajet_compass-0.1.0.dist-info/top_level.txt,sha256=qOGTjNmNfjeSXcmqLs4sowDpq-5fDxnIUg5ejdEhO6g,16
+datajet_compass-0.1.0.dist-info/RECORD,,

datajet_compass-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

datajet_compass-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+datajet_compass