divi 0.0.1.dev21__py3-none-any.whl → 0.0.1.dev23__py3-none-any.whl

divi/__init__.py

@@ -12,5 +12,5 @@ _core: Optional[Core] = None
 _auth: Optional[Auth] = None
 _datapark: Optional[DataPark] = None
 
-__version__ = "0.0.1.dev21"
+__version__ = "0.0.1.dev23"
 __all__ = ["proto", "obs_openai", "observable"]

divi/decorators/obs_openai.py

@@ -18,7 +18,6 @@ def _get_observable_create(create: Callable) -> Callable:
         return decorator(create)(*args, stream=stream, **kwargs)
 
     # TODO Async Observable Create
-    print("Is async", is_async(create))
     return observable_create if not is_async(create) else create
 
 

divi/decorators/observable.py

@@ -85,7 +85,11 @@ def observable(
             # recover parent context
             _SESSION_EXTRA.reset(token)
 
-            # TODO: collect inputs and outputs
+            # get the trace to collect data
+            trace = session_extra.get("trace")
+            if not trace:
+                raise ValueError("Trace not found in session_extra")
+            # TODO: collect inputs and outputs for SPAN_KIND_FUNCTION
             inputs = extract_flattened_inputs(func, *args, **kwargs)
             # create the span if it is the root span
             if divi._datapark and span.trace_id:
@@ -94,43 +98,18 @@ def observable(
                 )
             # end the trace if it is the root span
             if divi._datapark and not span.parent_span_id:
-                trace = session_extra.get("trace")
-                if trace:
-                    trace.end()
-                    divi._datapark.upsert_traces(
-                        session_id=trace.session_id, traces=[trace.signal]
-                    )
+                trace.end()
             # create the chat completion if it is a chat completion
             if divi._datapark and isinstance(result, ChatCompletion):
                 divi._datapark.create_chat_completion(
-                    span_id=span.span_id, inputs=inputs, completion=result
+                    span_id=span.span_id,
+                    trace_id=trace.trace_id,
+                    inputs=inputs,
+                    completion=result,
                 )
-            return result
 
-        @functools.wraps(func)
-        def generator_wrapper(
-            *args, session_extra: Optional[SessionExtra] = None, **kwargs
-        ):
-            span = Span(
-                kind=kind, name=name or func.__name__, metadata=metadata
-            )
-            session_extra = setup(span, _SESSION_EXTRA.get() or session_extra)
-            # set current context
-            token = _SESSION_EXTRA.set(session_extra)
-            # execute the function
-            results: List[Any] = []
-            span.start()
-            for item in func(*args, **kwargs):
-                results.append(item)
-                yield item
-            span.end()
-
-            # recover parent context
-            _SESSION_EXTRA.reset(token)
-            # TODO: collect results
+            return result
 
-        if inspect.isgeneratorfunction(func):
-            return generator_wrapper
         return wrapper
 
     # Function Decorator

divi/services/datapark/datapark.py

@@ -1,16 +1,16 @@
-from typing import Any, Dict, TYPE_CHECKING
-from typing_extensions import Mapping
+from typing import Any, Dict
 
 from google.protobuf.json_format import MessageToDict
+from openai import NotGiven
 from openai.types.chat import ChatCompletion
 from pydantic import UUID4
+from typing_extensions import Mapping
 
 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
-from openai import NotGiven
 
 
 class DataPark(Service):
@@ -33,7 +33,11 @@ class DataPark(Service):
         if not isinstance(obj, Mapping):
             return obj
 
-        return {key: value for key, value in obj.items() if not isinstance(value, NotGiven)}
+        return {
+            key: value
+            for key, value in obj.items()
+            if not isinstance(value, NotGiven)
+        }
 
     def create_session(self, session: SessionSignal) -> None:
         self.post("/api/session/", payload=session)
@@ -49,6 +53,7 @@ class DataPark(Service):
     def create_chat_completion(
         self,
         span_id: bytes,
+        trace_id: UUID4,
         inputs: Dict[str, Any],
         completion: ChatCompletion,
     ) -> None:
@@ -61,6 +66,7 @@ class DataPark(Service):
                 },
                 "/api/v1/chat/completions": {
                     "span_id": hex_span_id,
+                    "trace_id": str(trace_id),
                     "data": completion.model_dump(),
                 },
             }

divi/session/session.py

@@ -9,13 +9,13 @@ class SessionExtra(TypedDict, total=False):
 
     session_name: Optional[str]
     """Name of the session"""
-    trace: Optional[Trace]
-    """Trace ID UUID4"""
+    trace: Trace
+    """Trace in session"""
     parent_span_id: Optional[bytes]
     """Parent Span ID fixed string(8)"""
 
 
-class SessionSignal(TypedDict, total=False):
+class SessionSignal(TypedDict):
     """Session request"""
 
     id: str

divi/session/setup.py

@@ -1,3 +1,4 @@
+from typing_extensions import Optional
 import divi
 from divi.services import init as init_services
 from divi.session import Session, SessionExtra
@@ -5,10 +6,10 @@ from divi.signals.trace import Span
 from divi.signals.trace.trace import Trace
 
 
-def init(session_extra: SessionExtra) -> Session:
+def init_session(name: Optional[str]=None) -> Session:
     """init initializes the services and the Run"""
     init_services()
-    session = Session(name=session_extra.get("session_name"))
+    session = Session(name=name)
     if divi._datapark:
         divi._datapark.create_session(session.signal)
     return session
@@ -24,39 +25,21 @@ def setup(
         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)
+        divi._session = init_session(name=session_extra.get('session_name') or span.name)
 
     # setup trace
-    # init current span
-    trace = session_extra.get("trace")
+    trace = session_extra.get("trace") or Trace(divi._session.id, span.name)
     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(
+    span._add_node(trace.trace_id, parent_span_id)
+
+    # update the session_extra with the current trace and span
+    return 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/signals/trace/trace.py

@@ -7,6 +7,7 @@ from uuid import uuid4
 from pydantic import UUID4
 from typing_extensions import TypedDict
 
+import divi
 from divi.proto.common.v1.common_pb2 import KeyValue
 from divi.proto.trace.v1.trace_pb2 import Span as SpanProto
 
@@ -29,15 +30,19 @@ class TraceSignal(TypedDict, total=False):
     """Start time in iso format"""
     end_time: NullTime
     """End time in iso format"""
+    name: Optional[str]
 
 
 class Trace:
-    def __init__(self, session_id: UUID4):
+    def __init__(self, session_id: UUID4, name: Optional[str] = None):
         self.trace_id: UUID4 = uuid4()
         self.start_time: str | None = None
         self.end_time: str | None = None
+        self.name: Optional[str] = name
         self.session_id: UUID4 = session_id
 
+        self.start()
+
     @property
     def signal(self) -> TraceSignal:
         if self.start_time is None:
@@ -45,6 +50,7 @@ class Trace:
         signal = TraceSignal(
             id=str(self.trace_id),
             start_time=self.start_time,
+            name=self.name,
         )
         if self.end_time is not None:
             signal["end_time"] = NullTime(
@@ -60,12 +66,21 @@ class Trace:
     def start(self):
         """Start the trace by recording the current time in nanoseconds."""
         self.start_time = datetime.now(UTC).isoformat()
+        self.upsert_trace()
 
     def end(self):
         """End the trace by recording the end time in nanoseconds."""
         if self.start_time is None:
             raise ValueError("Span must be started before ending.")
         self.end_time = datetime.now(UTC).isoformat()
+        self.upsert_trace()
+
+    def upsert_trace(self):
+        """Upsert trace with datapark."""
+        if divi._datapark:
+            divi._datapark.upsert_traces(
+                session_id=self.session_id, traces=[self.signal]
+            )
 
 
 class Span:
@@ -102,7 +117,6 @@ class Span:
             trace_id=self.trace_id.bytes if self.trace_id else None,
             parent_span_id=self.parent_span_id,
         )
-        print(self.kind)
         signal.metadata.extend(
             KeyValue(key=k, value=v)
             for k, v in (self.metadata or dict()).items()
@@ -131,21 +145,7 @@ class Span:
             raise ValueError("Span must be started before ending.")
         self.end_time_unix_nano = time.time_ns()
 
-    def _as_root(self, trace_id: UUID4):
-        """Set the span as a root span."""
-        self.trace_id = trace_id
-        print("as root")
-        print(f"name: {self.name}")
-        print(f"trace_id: {self.trace_id}")
-        print(f"span_id: {self.span_id}")
-
-    def _add_parent(self, trace_id: UUID4, parent_id: bytes):
-        """Set the parent span ID."""
+    def _add_node(self, trace_id: UUID4, parent_id: Optional[bytes] = None):
+        """Add node for obs tree."""
         self.trace_id = trace_id
         self.parent_span_id = parent_id
-
-        print("add parent")
-        print(f"name: {self.name}")
-        print(f"trace_id: {trace_id}")
-        print(f"span_id: {self.span_id}")
-        print(f"parent_span_id: {parent_id}")

{divi-0.0.1.dev21.dist-info → divi-0.0.1.dev23.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: divi
-Version: 0.0.1.dev21
+Version: 0.0.1.dev23
 Summary: The Agent Platform for Observability & Evaluation
 License-File: LICENSE
 Requires-Python: >=3.11

{divi-0.0.1.dev21.dist-info → divi-0.0.1.dev23.dist-info}/RECORD

@@ -1,9 +1,9 @@
-divi/__init__.py,sha256=-XjMcUmlsDSB_v0_rf5owoldPj1mXFoAgUaVZpyfwqk,396
+divi/__init__.py,sha256=lcpSVEjIcXy-e9CuJB8V1izwaoEK8YpH8FtkzH0IqJc,396
 divi/utils.py,sha256=fXkjoyo_Lh8AZliKICOP460m0czUcNQjcEcceJbaOVA,1439
 divi/config/config.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 divi/decorators/__init__.py,sha256=HkyWdC1ctTsVFucCWCkj57JB4NmwONus1d2S2dUbvs4,110
-divi/decorators/obs_openai.py,sha256=4LtQPRgJPPyTOw1318Fi04WYqjIwGX8EP0YiWrTW0dk,1003
-divi/decorators/observable.py,sha256=OZ0Ry6Q2sdZPAT_T5j2wSsM7g0nwDroxeSbP3nLMDvA,4250
+divi/decorators/obs_openai.py,sha256=GI9c6gFArA6pTfa3EPednMtRqV2zIlofuTDdUYtS_x4,963
+divi/decorators/observable.py,sha256=xKapTyHL50mUyIDhDSYIXR1yZIdHTuw6WtI4H0j1xDE,3464
 divi/proto/common/v1/common.proto,sha256=Rx8wr0_tOtQ1NseTMnsav4ApD1MDALzQDBA2IvLRTU0,1775
 divi/proto/common/v1/common_pb2.py,sha256=br61OHQVAi6SI3baFcb5xJv2Xd-AZ04A19xeSjLNMXo,2442
 divi/proto/common/v1/common_pb2.pyi,sha256=LmTpFFLxHg2a_qPIdNXXwGEMkbiDcTJdarR9eC-6Fq8,2133
@@ -30,16 +30,16 @@ divi/services/core/core.py,sha256=PRwPtLgrgmCrejUfKf7HJNrAhGS0paFNZ7JwDToEUAk,12
 divi/services/core/finish.py,sha256=dIGQpVXcJY4-tKe7A1_VV3yoSHNCDPfOlUltvzvk6VI,231
 divi/services/core/init.py,sha256=e7-fgpOPglBXyEoPkgOAnpJk2ApdFbo7LPupxOb8N-w,1966
 divi/services/datapark/__init__.py,sha256=GbV1mwHE07yutgOlCIYHykSEL5KJ-ApgLutGMzu2eUE,86
-divi/services/datapark/datapark.py,sha256=1CoWJI2embzvijCmIIPwTLnliU4kSolUginGuSWHsBw,2182
+divi/services/datapark/datapark.py,sha256=d2pbrzVJtR3mNW1eQpbm-Wca-SvcfJqT7IuaQy7yHT0,2285
 divi/services/datapark/init.py,sha256=C32f9t3eLsxcYNqEyheh6nW455G2oR0YhhdqBcbN3ec,92
 divi/session/__init__.py,sha256=6lYemv21VQCIHx-xIdi7BxXcPNxVdvE60--8ArReUew,82
-divi/session/session.py,sha256=ID7bQ4CuRPUcZ9S-DSjndiduslYFJukRLuCtChYm3wk,838
-divi/session/setup.py,sha256=jlc3ICR5hRDdrcxoSxi66FYS2TYxqrIJOMgI4i1tEJs,1972
+divi/session/session.py,sha256=LlB2W2qGo0Vf-0L0CTQoXfzg_gCGpf0MTFsXQW7E6i4,817
+divi/session/setup.py,sha256=NeCxCb-uYhkKnOEiw8dBQHz0DEL8j1oxzQY3cBAmHbo,1380
 divi/session/teardown.py,sha256=YiBz_3yCiljMFEofZ60VmRL5sb8WA5GT7EYF8nFznZ4,133
 divi/signals/__init__.py,sha256=K1PaTAMwyBDsK6jJUg4QWy0xVJ_5MA6dlWiUyJeiSQA,44
 divi/signals/trace/__init__.py,sha256=K1PaTAMwyBDsK6jJUg4QWy0xVJ_5MA6dlWiUyJeiSQA,44
-divi/signals/trace/trace.py,sha256=y8Xio8Tp_e1tO8l_DAkUuzMrmENNKfIdklF0iWrjGYk,4547
-divi-0.0.1.dev21.dist-info/METADATA,sha256=xyAsE6Uj3sA8T7I83LVLtCEyum_hMQT4ULien2Q-0NU,497
-divi-0.0.1.dev21.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-divi-0.0.1.dev21.dist-info/licenses/LICENSE,sha256=5OJuZ4wMMEV0DgF0tofhAlS_KLkaUsZwwwDS2U_GwQ0,1063
-divi-0.0.1.dev21.dist-info/RECORD,,
+divi/signals/trace/trace.py,sha256=OsfrZPHp241_NN8W79U4O69HsHQajez_d3rz6yJRN9s,4508
+divi-0.0.1.dev23.dist-info/METADATA,sha256=KdspT1iEra11ixCEtl94Mz_fCziMVOOER08rDE4yCWU,497
+divi-0.0.1.dev23.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+divi-0.0.1.dev23.dist-info/licenses/LICENSE,sha256=5OJuZ4wMMEV0DgF0tofhAlS_KLkaUsZwwwDS2U_GwQ0,1063
+divi-0.0.1.dev23.dist-info/RECORD,,