divi 0.0.1.dev7__py3-none-any.whl → 0.0.1.dev14__py3-none-any.whl

divi/proto/trace/v1/trace_pb2.py

@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: divi/proto/trace/v1/trace.proto
+# Protobuf Python Version: 5.29.0
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    5,
+    29,
+    0,
+    '',
+    'divi/proto/trace/v1/trace.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from divi.proto.common.v1 import common_pb2 as divi_dot_proto_dot_common_dot_v1_dot_common__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1f\x64ivi/proto/trace/v1/trace.proto\x12\x13\x64ivi.proto.trace.v1\x1a!divi/proto/common/v1/common.proto\"\x83\x01\n\x05Trace\x12\x10\n\x08trace_id\x18\x01 \x01(\x0c\x12\x1c\n\x14start_time_unix_nano\x18\x02 \x01(\x06\x12\x1a\n\x12\x65nd_time_unix_nano\x18\x03 \x01(\x06\x12.\n\x05spans\x18\x04 \x03(\x0b\x32\x1f.divi.proto.trace.v1.ScopeSpans\"6\n\nScopeSpans\x12(\n\x05spans\x18\x02 \x03(\x0b\x32\x19.divi.proto.trace.v1.Span\"\xa4\x02\n\x04Span\x12\x10\n\x08trace_id\x18\x01 \x01(\x0c\x12\x0f\n\x07span_id\x18\x02 \x01(\x0c\x12\x16\n\x0eparent_span_id\x18\x03 \x01(\x0c\x12\x0c\n\x04name\x18\x04 \x01(\t\x12\x30\n\x04kind\x18\x05 \x01(\x0e\x32\".divi.proto.trace.v1.Span.SpanKind\x12\x1c\n\x14start_time_unix_nano\x18\x06 \x01(\x06\x12\x1a\n\x12\x65nd_time_unix_nano\x18\x07 \x01(\x06\x12\x30\n\x08metadata\x18\x08 \x03(\x0b\x32\x1e.divi.proto.common.v1.KeyValue\"5\n\x08SpanKind\x12\x16\n\x12SPAN_KIND_FUNCTION\x10\x00\x12\x11\n\rSPAN_KIND_LLM\x10\x01\x42\rZ\x0bservices/pbb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'divi.proto.trace.v1.trace_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'Z\013services/pb'
+  _globals['_TRACE']._serialized_start=92
+  _globals['_TRACE']._serialized_end=223
+  _globals['_SCOPESPANS']._serialized_start=225
+  _globals['_SCOPESPANS']._serialized_end=279
+  _globals['_SPAN']._serialized_start=282
+  _globals['_SPAN']._serialized_end=574
+  _globals['_SPAN_SPANKIND']._serialized_start=521
+  _globals['_SPAN_SPANKIND']._serialized_end=574
+# @@protoc_insertion_point(module_scope)

divi/proto/trace/v1/trace_pb2.pyi

@@ -0,0 +1,52 @@
+from divi.proto.common.v1 import common_pb2 as _common_pb2
+from google.protobuf.internal import containers as _containers
+from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class Trace(_message.Message):
+    __slots__ = ("trace_id", "start_time_unix_nano", "end_time_unix_nano", "spans")
+    TRACE_ID_FIELD_NUMBER: _ClassVar[int]
+    START_TIME_UNIX_NANO_FIELD_NUMBER: _ClassVar[int]
+    END_TIME_UNIX_NANO_FIELD_NUMBER: _ClassVar[int]
+    SPANS_FIELD_NUMBER: _ClassVar[int]
+    trace_id: bytes
+    start_time_unix_nano: int
+    end_time_unix_nano: int
+    spans: _containers.RepeatedCompositeFieldContainer[ScopeSpans]
+    def __init__(self, trace_id: _Optional[bytes] = ..., start_time_unix_nano: _Optional[int] = ..., end_time_unix_nano: _Optional[int] = ..., spans: _Optional[_Iterable[_Union[ScopeSpans, _Mapping]]] = ...) -> None: ...
+
+class ScopeSpans(_message.Message):
+    __slots__ = ("spans",)
+    SPANS_FIELD_NUMBER: _ClassVar[int]
+    spans: _containers.RepeatedCompositeFieldContainer[Span]
+    def __init__(self, spans: _Optional[_Iterable[_Union[Span, _Mapping]]] = ...) -> None: ...
+
+class Span(_message.Message):
+    __slots__ = ("trace_id", "span_id", "parent_span_id", "name", "kind", "start_time_unix_nano", "end_time_unix_nano", "metadata")
+    class SpanKind(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+        __slots__ = ()
+        SPAN_KIND_FUNCTION: _ClassVar[Span.SpanKind]
+        SPAN_KIND_LLM: _ClassVar[Span.SpanKind]
+    SPAN_KIND_FUNCTION: Span.SpanKind
+    SPAN_KIND_LLM: Span.SpanKind
+    TRACE_ID_FIELD_NUMBER: _ClassVar[int]
+    SPAN_ID_FIELD_NUMBER: _ClassVar[int]
+    PARENT_SPAN_ID_FIELD_NUMBER: _ClassVar[int]
+    NAME_FIELD_NUMBER: _ClassVar[int]
+    KIND_FIELD_NUMBER: _ClassVar[int]
+    START_TIME_UNIX_NANO_FIELD_NUMBER: _ClassVar[int]
+    END_TIME_UNIX_NANO_FIELD_NUMBER: _ClassVar[int]
+    METADATA_FIELD_NUMBER: _ClassVar[int]
+    trace_id: bytes
+    span_id: bytes
+    parent_span_id: bytes
+    name: str
+    kind: Span.SpanKind
+    start_time_unix_nano: int
+    end_time_unix_nano: int
+    metadata: _containers.RepeatedCompositeFieldContainer[_common_pb2.KeyValue]
+    def __init__(self, trace_id: _Optional[bytes] = ..., span_id: _Optional[bytes] = ..., parent_span_id: _Optional[bytes] = ..., name: _Optional[str] = ..., kind: _Optional[_Union[Span.SpanKind, str]] = ..., start_time_unix_nano: _Optional[int] = ..., end_time_unix_nano: _Optional[int] = ..., metadata: _Optional[_Iterable[_Union[_common_pb2.KeyValue, _Mapping]]] = ...) -> None: ...

divi/services/__init__.py

@@ -0,0 +1,7 @@
+from .auth import Auth
+from .core import Core
+from .datapark import DataPark
+from .finish import finish
+from .init import init
+
+__all__ = ["init", "finish", "Core", "Auth", "DataPark"]

divi/services/auth/__init__.py

@@ -0,0 +1,4 @@
+from .auth import Auth
+from .init import init
+
+__all__ = ["Auth", "init"]

divi/services/auth/auth.py

@@ -0,0 +1,21 @@
+import requests
+
+from divi.services.auth.tokman import Token
+from divi.services.service import Service
+
+
+class Auth(Service):
+    def __init__(self, api_key: str, host="localhost", port=3000):
+        super().__init__(host, port)
+        self.api_key = api_key
+        self.token = Token(self)
+
+    def auth_with_api_key(self) -> str:
+        """Get the token with the API key."""
+        r = requests.post(
+            f"http://{self.target}/api/auth/api_key",
+            json={"api_key": self.api_key},
+        )
+        if r.status_code == 200:
+            return r.json()["data"]
+        raise ValueError(r.json()["message"])

divi/services/auth/init.py

@@ -0,0 +1,21 @@
+import os
+from typing import Optional
+
+from divi.services.auth import Auth
+
+DIVI_API_KEY = "DIVI_API_KEY"
+
+
+def init(api_key: Optional[str] = None) -> Auth:
+    key = api_key if api_key else os.getenv(DIVI_API_KEY)
+    if not key:
+        raise ValueError("API key is required")
+    return Auth(api_key=key)
+
+
+if __name__ == "__main__":
+    auth = init()
+    if not auth:
+        raise ValueError("Auth object is not available")
+    print("=== Auth ===")
+    print(auth.token)

divi/services/auth/tokman.py

@@ -0,0 +1,42 @@
+import time
+from weakref import ref
+
+import jwt
+
+
+class Token:
+    """JWT Manager Class."""
+
+    def __init__(self, auth) -> None:
+        self.auth = ref(auth)
+        self.claims: dict = {}
+        self.__token: str = ""
+
+    def __str__(self) -> str:
+        return self.token
+
+    @property
+    def exp(self) -> int:
+        """Return the expiration time."""
+        return self.claims.get("exp", 0)
+
+    @property
+    def token(self) -> str:
+        """Return the token string."""
+        # If the token is expired, get a new one
+        if not self.__token or self.exp - time.time() < 3600:
+            self._init_token()
+        return self.__token
+
+    def _init_token(self):
+        """Initialize the token."""
+        auth = self.auth()
+        if not auth:
+            raise ValueError("Auth object is not available")
+        self.__token = auth.auth_with_api_key()
+        self.claims = _decode_token(self.__token)
+
+
+def _decode_token(token: str) -> dict:
+    """Decode the token payload."""
+    return jwt.decode(token, options={"verify_signature": False})

divi/services/core/__init__.py

@@ -0,0 +1,5 @@
+from .core import Core
+from .finish import finish
+from .init import init
+
+__all__ = ["init", "finish", "Core"]

divi/services/core/core.py

@@ -0,0 +1,35 @@
+from subprocess import Popen
+from typing import Callable, List, Optional
+
+import grpc
+
+import divi
+from divi.proto.core.health.v1.health_service_pb2 import HealthCheckRequest
+from divi.proto.core.health.v1.health_service_pb2_grpc import HealthServiceStub
+from divi.services.service import Service
+
+
+class Core(Service):
+    """Core Runtime Class."""
+
+    def __init__(self, host="localhost", port=50051) -> None:
+        super().__init__(host, port)
+        self.process: Optional[Popen] = None
+        self.hooks: List[Callable[[], None]] = []
+
+    def check_health(self) -> bool:
+        """Check the health of the service."""
+        with grpc.insecure_channel(self.target) as channel:
+            stub = HealthServiceStub(channel)
+            response, call = stub.Check.with_call(
+                HealthCheckRequest(version=divi.__version__),
+                # Note: ((),) notice the `,` at the end of the tuple
+                metadata=(("version", divi.__version__),),
+            )
+        print(f"Health check: {response.message}")
+        for key, value in call.trailing_metadata():
+            print(
+                "python client received trailing metadata: key=%s value=%s"
+                % (key, value)
+            )
+        return response.status

divi/{core → services/core}/finish.py

@@ -4,12 +4,12 @@ import divi
 
 
 def finish():
-    """Clean up the run."""
-    run = divi.run
-    if run is None:
+    """Clean up the core."""
+    core = divi._core
+    if core is None:
         return
 
     # Clean up the hooks
-    for hook in run.hooks:
+    for hook in core.hooks:
         hook()
         atexit.unregister(hook)

divi/{core → services/core}/init.py

@@ -1,60 +1,50 @@
-import grpc
-import time
-import socket
 import atexit
+import socket
 import subprocess
+import time
 
-from typing import Union
-
-import divi
+import grpc
 
-from divi.core.run import Run
+from divi.services.core import Core
 from divi.utils import get_server_path
 
 
-def init(
-    host: Union[str, None] = None, port: Union[str, None] = None
-) -> Union[Run, None]:
-    divi.run = Run(host, port)
-    _start_server()
+def init(host="localhost", port=50051) -> Core:
+    core = Core(host=host, port=port)
+    _start_server(core)
+    return core
 
 
-def _start_server():
+def _start_server(core: Core):
     """Start the backend server."""
-    # get the run object
-    run = divi.run
-    if run is None:
-        return
-
     # start the server
     bin_path = get_server_path()
-    command = [bin_path]
-    run.process = subprocess.Popen(command)
+    command = [bin_path, "-port", str(core.port)]
+    core.process = subprocess.Popen(command)
 
     # Wait for the port to be open
-    if not _wait_for_port(run.host, run.port, 10):
-        run.process.terminate()
+    if not _wait_for_port(core.host, core.port, 10):
+        core.process.terminate()
         raise RuntimeError("Service failed to start: port not open")
 
     # Check if the gRPC channel is ready
-    channel = grpc.insecure_channel(run.target)
+    channel = grpc.insecure_channel(core.target)
     try:
         grpc.channel_ready_future(channel).result(timeout=10)
     except grpc.FutureTimeoutError:
-        run.process.terminate()
+        core.process.terminate()
         raise RuntimeError("gRPC channel not ready")
     finally:
         channel.close()
 
+    core.hooks.append(core.process.terminate)
+    atexit.register(core.process.terminate)
+
     # Health check
-    status = run.check_health()
+    status = core.check_health()
     if not status:
-        run.process.terminate()
         raise RuntimeError("Service failed health check")
 
-    run.hooks.append(run.process.terminate)
-    atexit.register(run.process.terminate)
-
 
 def _wait_for_port(host, port, timeout_seconds):
     """Wait until the specified port is open."""
@@ -78,3 +68,7 @@ def _is_port_open(host, port):
     except Exception as e:
         print(f"Error checking port: {e}")
         return False
+
+
+if __name__ == "__main__":
+    init()

divi/services/datapark/__init__.py

@@ -0,0 +1,4 @@
+from .datapark import DataPark
+from .init import init
+
+__all__ = ["DataPark", "init"]

divi/services/datapark/datapark.py

@@ -0,0 +1,63 @@
+import requests
+from google.protobuf.json_format import MessageToDict
+from openai.types.chat.chat_completion import ChatCompletion
+from pydantic import UUID4
+
+import divi
+from divi.proto.trace.v1.trace_pb2 import ScopeSpans
+from divi.services.service import Service
+from divi.session.session import SessionSignal
+from divi.signals.trace.trace import TraceSignal
+
+
+class DataPark(Service):
+    def __init__(self, host="localhost", port=3001):
+        super().__init__(host, port)
+        if not divi._auth:
+            raise ValueError("No auth service")
+        self.token = divi._auth.token
+
+    @property
+    def headers(self):
+        return {"Authorization": f"Bearer {self.token}"}
+
+    def create_session(self, session: SessionSignal):
+        r = requests.post(
+            f"http://{self.target}/api/session/",
+            headers=self.headers,
+            json=session,
+        )
+        if r.status_code != 201:
+            raise ValueError(r.json()["message"])
+
+    def upsert_traces(self, session_id: UUID4, traces: list[TraceSignal]):
+        r = requests.post(
+            f"http://{self.target}/api/session/{session_id}/traces",
+            headers=self.headers,
+            json=traces,
+        )
+        if r.status_code != 201:
+            raise ValueError(r.json()["message"])
+
+    def create_spans(self, trace_id: UUID4, spans: ScopeSpans):
+        r = requests.post(
+            f"http://{self.target}/api/trace/{trace_id}/spans",
+            headers=self.headers,
+            json=MessageToDict(spans),
+        )
+        if r.status_code != 201:
+            raise ValueError(r.json()["message"])
+
+    def create_chat_completion(
+        self, span_id: bytes, completion: ChatCompletion
+    ):
+        r = requests.post(
+            f"http://{self.target}/api/v1/chat/completions",
+            headers=self.headers,
+            json={
+                "span_id": span_id.hex(),
+                "data": completion.model_dump(),
+            },
+        )
+        if r.status_code != 201:
+            raise ValueError(r.json()["message"])

divi/services/datapark/init.py

@@ -0,0 +1,5 @@
+from divi.services.datapark import DataPark
+
+
+def init() -> DataPark:
+    return DataPark()

divi/services/finish.py

@@ -0,0 +1,9 @@
+import divi
+from divi.services.core import finish as clean_up_core
+
+
+def finish():
+    clean_up_core()
+    divi._auth = None
+    divi._core = None
+    divi._datapark = None

divi/services/init.py

@@ -0,0 +1,14 @@
+import divi
+from divi.services.auth import init as init_auth
+from divi.services.core import init as init_core
+from divi.services.datapark import init as init_datapark
+
+
+def init():
+    if not divi._auth:
+        divi._auth = init_auth()
+    if not divi._datapark:
+        divi._datapark = init_datapark()
+    # TODO - Uncomment this when the core service is ready
+    # if not divi._core:
+    #     divi._core = init_core(port=3002)

divi/services/service.py

@@ -0,0 +1,11 @@
+class Service:
+    """Service management class."""
+
+    def __init__(self, host: str, port: int):
+        self.host = host
+        self.port = port
+
+    @property
+    def target(self) -> str:
+        """Return the target string: host:port."""
+        return f"{self.host}:{self.port}"

divi/session/__init__.py

@@ -0,0 +1,3 @@
+from .session import Session, SessionExtra
+
+__all__ = ["Session", "SessionExtra"]

divi/session/session.py

@@ -0,0 +1,40 @@
+from typing import Optional, TypedDict
+from uuid import uuid4
+
+from divi.signals.trace.trace import Trace
+
+
+class SessionExtra(TypedDict, total=False):
+    """Extra information for Session"""
+
+    session_name: Optional[str]
+    """Name of the session"""
+    trace: Optional[Trace]
+    """Trace ID UUID4"""
+    parent_span_id: Optional[bytes]
+    """Parent Span ID fixed string(8)"""
+
+
+class SessionSignal(TypedDict, total=False):
+    """Session request"""
+
+    id: str
+    """Session ID UUID4"""
+    name: Optional[str]
+    """Session name"""
+
+
+class Session:
+    def __init__(
+        self,
+        name: Optional[str] = None,
+    ):
+        self.id = uuid4()
+        self.name = name
+
+    @property
+    def signal(self) -> SessionSignal:
+        return SessionSignal(
+            id=str(self.id),
+            name=self.name,
+        )

divi/session/setup.py

@@ -0,0 +1,62 @@
+import divi
+from divi.services import init as init_services
+from divi.session import Session, SessionExtra
+from divi.signals.trace import Span
+from divi.signals.trace.trace import Trace
+
+
+def init(session_extra: SessionExtra) -> Session:
+    """init initializes the services and the Run"""
+    init_services()
+    session = Session(name=session_extra.get("session_name"))
+    if divi._datapark:
+        divi._datapark.create_session(session.signal)
+    return session
+
+
+def setup(
+    span: Span,
+    session_extra: SessionExtra | None,
+):
+    """setup trace
+
+    Args:
+        span (Span): Span instance
+        session_extra (SessionExtra | None): Extra information from user input
+    """
+    # TOOD: merge run_extra input by user with the one from the context
+    # temp solution: Priority: run_extra_context.get() > run_extra
+    session_extra = session_extra or SessionExtra()
+
+    # init the session if not already initialized
+    if not divi._session:
+        divi._session = init(session_extra=session_extra)
+
+    # setup trace
+    # init current span
+    trace = session_extra.get("trace")
+    parent_span_id = session_extra.get("parent_span_id")
+    if trace and parent_span_id:
+        span._add_parent(trace.trace_id, parent_span_id)
+    else:
+        trace = Trace(divi._session.id)
+        trace.start()
+        span._as_root(trace.trace_id)
+        # create the trace
+        if divi._datapark:
+            divi._datapark.upsert_traces(
+                session_id=divi._session.id, traces=[trace.signal]
+            )
+
+    # update the session_extra with the current span
+    # session_extra["trace_id"] = span.trace_id
+    # session_extra["parent_span_id"] = span.span_id
+    session_extra = SessionExtra(
+        session_name=divi._session.name,
+        trace=trace,
+        # set the parent_span_id to the current span_id
+        parent_span_id=span.span_id,
+    )
+
+    # offer end hook to collect data at whe end of the span ?
+    return session_extra

divi/session/teardown.py

@@ -0,0 +1,7 @@
+import divi
+from divi.services import finish as clean_up_services
+
+
+def teardown():
+    clean_up_services()
+    divi._session = None

divi/signals/__init__.py

@@ -0,0 +1,3 @@
+from .trace import Span
+
+__all__ = ["Span"]

divi/signals/trace/__init__.py

@@ -0,0 +1,3 @@
+from .trace import Span
+
+__all__ = ["Span"]