divi 0.0.1.dev28__py3-none-any.whl → 0.0.1.dev47__py3-none-any.whl

divi/__init__.py

@@ -5,6 +5,7 @@ from .decorators import obs_openai, observable
 from .evaluation import Evaluator, Score
 from .services import Auth, Core, DataPark
 from .session import Session
+from .signals import Kind
 
 name: str = "divi"
 
@@ -14,5 +15,5 @@ _auth: Optional[Auth] = None
 _datapark: Optional[DataPark] = None
 _evaluator: Optional[Evaluator] = None
 
-__version__ = "0.0.1.dev28"
-__all__ = ["proto", "obs_openai", "observable", "Score"]
+__version__ = "0.0.1.dev47"
+__all__ = ["proto", "obs_openai", "observable", "Score", "Kind"]

divi/decorators/collect.py

@@ -0,0 +1,34 @@
+from typing import Any
+
+from google.protobuf.message import Error
+from openai.types.chat import ChatCompletion
+from typing_extensions import Dict
+
+import divi
+from divi.evaluation.evaluator import EvaluationScore
+from divi.signals.span import Span
+
+
+def collect(span: Span, input: Dict[str, Any], result: Any):
+    if not divi._datapark or span.trace_id is None:
+        raise Error("divi._datapark or span.trace_id is None")
+    # TODO: collect inputs and outputs for SPAN_KIND_FUNCTION
+
+    # collect inputs and outputs for SPAN_KIND_LLM
+    if isinstance(result, ChatCompletion):
+        divi._datapark.create_chat_completion(
+            span_id=span.span_id,
+            trace_id=span.trace_id,
+            inputs=input,
+            completion=result,
+        )
+
+    # collect inputs and outputs for SPAN_KIND_EVALUATION
+    if isinstance(result, list) and all(
+        isinstance(x, EvaluationScore) for x in result
+    ):
+        divi._datapark.create_scores(
+            span_id=span.span_id,
+            trace_id=span.trace_id,
+            scores=result,
+        )

divi/decorators/obs_openai.py

@@ -6,6 +6,7 @@ from typing_extensions import Optional
 
 from divi.decorators.observable import observable
 from divi.evaluation.scores import Score
+from divi.signals.span import Kind
 from divi.utils import is_async
 
 if TYPE_CHECKING:
@@ -21,7 +22,7 @@ def _get_observable_create(
 ) -> Callable:
     @functools.wraps(create)
     def observable_create(*args, stream: bool = False, **kwargs):
-        decorator = observable(kind="llm", name=name, scores=scores)
+        decorator = observable(kind=Kind.llm, name=name, scores=scores)
         return decorator(create)(*args, stream=stream, **kwargs)
 
     # TODO Async Observable Create
@@ -30,7 +31,7 @@ def _get_observable_create(
 
 def obs_openai(
     client: C,
-    name: Optional[str] = "Chat Bot",
+    name: Optional[str] = "Agent",
     scores: Optional[list[Score]] = None,
 ) -> C:
     """Make OpenAI client observable."""

divi/decorators/observable.py

@@ -1,4 +1,3 @@
-import contextvars
 import functools
 from typing import (
     Any,
@@ -14,26 +13,15 @@ from typing import (
     runtime_checkable,
 )
 
-from openai.types.chat import ChatCompletion
-
-import divi
+from divi.decorators.observe import observe
 from divi.evaluation.evaluate import evaluate_scores
 from divi.evaluation.scores import Score
-from divi.proto.trace.v1.trace_pb2 import ScopeSpans
 from divi.session import SessionExtra
-from divi.session.setup import setup
-from divi.signals.trace import Span
-from divi.utils import extract_flattened_inputs
+from divi.signals.span import Kind, Span
 
 R = TypeVar("R", covariant=True)
 P = ParamSpec("P")
 
-# ContextVar to store the extra information
-# from the Session and parent Span
-_SESSION_EXTRA = contextvars.ContextVar[Optional[SessionExtra]](
-    "_SESSION_EXTRA", default=None
-)
-
 
 @runtime_checkable
 class WithSessionExtra(Protocol, Generic[P, R]):
@@ -51,7 +39,7 @@ def observable(func: Callable[P, R]) -> WithSessionExtra[P, R]: ...
 
 @overload
 def observable(
-    kind: str = "function",
+    kind: Kind = Kind.function,
     *,
     name: Optional[str] = None,
     scores: Optional[list[Score]] = None,
@@ -64,7 +52,7 @@ def observable(
 ) -> Union[Callable, Callable[[Callable], Callable]]:
     """Observable decorator factory."""
 
-    kind = kwargs.pop("kind", "function")
+    kind = kwargs.pop("kind", Kind.function)
     name = kwargs.pop("name", None)
     metadata = kwargs.pop("metadata", None)
     scores: list[Score] = kwargs.pop("scores", None)
@@ -74,44 +62,23 @@ def observable(
         def wrapper(
             *args, session_extra: Optional[SessionExtra] = None, **kwargs
         ):
+            # 1. init the span
             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
-            span.start()
-            result = func(*args, **kwargs)
-            span.end()
-            # recover parent context
-            _SESSION_EXTRA.reset(token)
-
-            # get the trace to collect data
-            trace = session_extra.get("trace")
-            if not trace:
-                raise ValueError("Trace not found in session context.")
-            # 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:
-                divi._datapark.create_spans(
-                    span.trace_id, ScopeSpans(spans=[span.signal])
-                )
-            # end the trace if it is the root span
-            if divi._datapark and not span.parent_span_id:
-                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,
-                    trace_id=trace.trace_id,
-                    inputs=inputs,
-                    completion=result,
-                )
-                # evaluate the scores if they are provided
-                if scores is not None and scores.__len__() > 0:
-                    evaluate_scores(inputs, outputs=result, scores=scores)
+
+            # 2. observe the function
+            result = observe(
+                *args,
+                func=func,
+                span=span,
+                session_extra=session_extra,
+                **kwargs,
+            )
+
+            # 3. evaluate the scores if they are provided
+            messages = kwargs.get("messages", [])
+            evaluate_scores(messages, outputs=result, scores=scores)
 
             return result
 

divi/decorators/observe.py

@@ -0,0 +1,47 @@
+import contextvars
+from typing import (
+    Callable,
+    Optional,
+)
+
+from divi.decorators.collect import collect
+from divi.session import SessionExtra
+from divi.session.setup import setup
+from divi.signals.span import Span
+from divi.utils import extract_flattened_inputs
+
+# ContextVar to store the extra information
+# from the Session and parent Span
+_SESSION_EXTRA = contextvars.ContextVar[Optional[SessionExtra]](
+    "_SESSION_EXTRA", default=None
+)
+
+
+def observe(
+    *args,
+    func: Callable,
+    span: Span,
+    session_extra: Optional[SessionExtra] = None,
+    **kwargs,
+):
+    session_extra = setup(span, _SESSION_EXTRA.get() or session_extra)
+    # set current context
+    token = _SESSION_EXTRA.set(session_extra)
+    # execute the function
+    span.start()
+    result = func(*args, **kwargs)
+    span.end()
+    # recover parent context
+    _SESSION_EXTRA.reset(token)
+
+    # get the trace to collect data
+    trace = session_extra.get("trace")
+    # end the trace if it is the root span
+    if trace and not span.parent_span_id:
+        trace.end()
+
+    # collect inputs and outputs
+    inputs = extract_flattened_inputs(func, *args, **kwargs)
+    collect(span, inputs, result)
+
+    return result

divi/evaluation/evaluate.py

@@ -1,14 +1,18 @@
-import copy
 import os
-from typing import Any, Dict, Optional
+from typing import Optional
 
-from openai.types.chat import ChatCompletion
+from openai.types.chat import (
+    ChatCompletion,
+    ChatCompletionMessageParam,
+)
 from typing_extensions import List
 
 import divi
+from divi.decorators.observe import observe
 from divi.evaluation import Evaluator
 from divi.evaluation.evaluator import EvaluatorConfig
 from divi.evaluation.scores import Score
+from divi.signals.span import Kind, Span
 
 OPENAI_API_KEY = "OPENAI_API_KEY"
 OPENAI_BASE_URL = "OPENAI_BASE_URL"
@@ -29,29 +33,29 @@ def init_evaluator(config: Optional[EvaluatorConfig] = None):
 
 
 def evaluate_scores(
-    inputs: Dict[str, Any],
-    outputs: ChatCompletion,
-    scores: List[Score],
+    messages: Optional[List[ChatCompletionMessageParam]],
+    outputs: Optional[ChatCompletion],
+    scores: Optional[List[Score]],
     config: Optional[EvaluatorConfig] = None,
 ):
+    if messages is None or scores is None or scores.__len__() == 0:
+        return
     if not divi._evaluator:
         divi._evaluator = init_evaluator(config)
 
-    # create conversation with result and inputs
-    input_messages = inputs.get("messages", None)
-    if input_messages is None:
-        raise ValueError("No messages found in inputs")
-    output_message = outputs.choices[0].message
-    if output_message is None:
-        raise ValueError("No message found in outputs")
-
-    conversations = copy.deepcopy(input_messages)
-    conversations.append(
-        {"role": output_message.role, "content": output_message.content}
-    )
-    evaluation_scores = divi._evaluator.evaluate(
-        "\n".join(f"{m['role']}: {m['content']}" for m in conversations), scores
-    )
-
-    # TODO: collect all evaluation scores and link them to span
-    print(evaluation_scores)
+    if isinstance(outputs, ChatCompletion):
+        output_message = outputs.choices[0].message.content
+        if not output_message:
+            return
+
+        evaluation_span = Span(kind=Kind.evaluation, name="Evaluation")
+        observe(
+            func=divi._evaluator.evaluate,
+            span=evaluation_span,
+            target=output_message,
+            conversation="\n".join(
+                f"{m.get('role', 'unknown')}: {m.get('content')}"
+                for m in messages
+            ),
+            scores=scores,
+        )

divi/evaluation/evaluator.py

@@ -50,34 +50,42 @@ class Evaluator:
         )
 
     @staticmethod
-    def generate_prompt(conversation: str, score: Score) -> str:
+    def generate_prompt(target: str, conversation: str, score: Score) -> str:
         return PROMPT_TEMPLATE.format(
-            requirements=PRESET_PROMPT[score.value], conversation=conversation
+            requirements=PRESET_PROMPT[score.value],
+            target=target,
+            conversation=conversation,
         )
 
     def _sync_evaluate_once(
-        self, conversation: str, score: Score
+        self, target: str, conversation: str, score: Score
     ) -> Optional[EvaluationResult]:
-        prompt = self.generate_prompt(conversation, score)
+        prompt = self.generate_prompt(target, conversation, score)
         response = self.sync_client.beta.chat.completions.parse(
             model=self.config.model,
             messages=[{"role": "user", "content": prompt}],
             temperature=self.config.temperature,
             response_format=EvaluationResult,
         )
-        return response.choices[0].message.parsed
+        result = response.choices[0].message.parsed
+        if result is not None:
+            result.name = score
+        return result
 
     async def _async_evaluate_once(
-        self, conversation: str, score: Score
+        self, target: str, conversation: str, score: Score
     ) -> Optional[EvaluationResult]:
-        prompt = self.generate_prompt(conversation, score)
+        prompt = self.generate_prompt(target, conversation, score)
         response = await self.async_client.beta.chat.completions.parse(
             model=self.config.model,
             messages=[{"role": "user", "content": prompt}],
             temperature=self.config.temperature,
             response_format=EvaluationResult,
         )
-        return response.choices[0].message.parsed
+        result = response.choices[0].message.parsed
+        if result is not None:
+            result.name = score
+        return result
 
     def _aggregate_result(
         self, name: Score, evaluations: List[EvaluationResult]
@@ -115,13 +123,15 @@ class Evaluator:
         return aggregated_results
 
     def evaluate_sync(
-        self, conversation: str, scores: list[Score], n_rounds: int
+        self, target: str, conversation: str, scores: list[Score], n_rounds: int
     ) -> List[EvaluationScore]:
         with concurrent.futures.ThreadPoolExecutor(
             max_workers=self.config.max_concurrency
         ) as executor:
             futures = [
-                executor.submit(self._sync_evaluate_once, conversation, score)
+                executor.submit(
+                    self._sync_evaluate_once, target, conversation, score
+                )
                 for _ in range(n_rounds)
                 for score in scores
             ]
@@ -133,13 +143,15 @@ class Evaluator:
         )
 
     async def evaluate_async(
-        self, conversation: str, scores: list[Score], n_rounds: int
+        self, target: str, conversation: str, scores: list[Score], n_rounds: int
     ) -> List[EvaluationScore]:
         semaphore = asyncio.Semaphore(self.config.max_concurrency)
 
         async def sem_task(score):
             async with semaphore:
-                return await self._async_evaluate_once(conversation, score)
+                return await self._async_evaluate_once(
+                    target, conversation, score
+                )
 
         tasks = [sem_task(score) for _ in range(n_rounds) for score in scores]
         evaluations = await asyncio.gather(*tasks)
@@ -149,6 +161,7 @@ class Evaluator:
 
     def evaluate(
         self,
+        target: str,
         conversation: str,
         scores: list[Score],
         n_rounds: int = 5,
@@ -156,6 +169,6 @@ class Evaluator:
     ) -> List[EvaluationScore]:
         if mode == "async":
             return asyncio.run(
-                self.evaluate_async(conversation, scores, n_rounds)
+                self.evaluate_async(target, conversation, scores, n_rounds)
             )
-        return self.evaluate_sync(conversation, scores, n_rounds)
+        return self.evaluate_sync(target, conversation, scores, n_rounds)

divi/evaluation/prompts.py

@@ -1,14 +1,19 @@
 PROMPT_TEMPLATE = (
-    "Your evaluation task is: {requirements}\n\n"
+    "The *requirements* of the evaluation task is: {requirements}\n\n"
+    "Below is the *context* of the conversation (for reference only):\n"
+    "{conversation}\n\n"
+    "Now, in view of both the requirements and the context, evaluate the assistant’s response:\n"
+    "{target}\n\n"
     "Please perform step-by-step reasoning to reach your judgment.\n\n"
     "Strictly output your answer in the following JSON format:\n"
-    '{{"judgment": bool, "reasoning": "string"}}\n\n'
-    "Do not output anything else.\n\n"
-    "Here is the conversation to evaluate:\n"
-    "{conversation}"
+    "{{\n"
+    '  "judgment": bool,        # true if the response meets all requirements\n'
+    '  "reasoning": "string"    # concise explanation, hitting only the key points\n'
+    "}}\n"
+    "Do not output anything else."
 )
 
 PRESET_PROMPT = {
-    "task_completion": "Evaluate whether the model's output completely fulfills the user's task requirements.",
-    "instruction_adherence": "Evaluate whether the model's output strictly follows the user's instructions without omissions, deviations, or hallucinations.",
+    "task_completion": "Assess whether the assistant response fulfills the user's task requirements.",
+    "instruction_adherence": "Assess whether the assistant response strictly follows every instruction given by the user, without omissions, deviations, or hallucinations.",
 }

divi/proto/trace/v1/trace.proto

@@ -31,6 +31,9 @@ message Span {
 
         // LLM represents a llm api call.
         SPAN_KIND_LLM = 1;
+
+        // EVALUATION represents an evaluation.
+        SPAN_KIND_EVALUATION = 2;
     }
 
     // The kind of the span.

divi/proto/trace/v1/trace_pb2.py

@@ -25,7 +25,7 @@ _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\"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')
+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\"6\n\nScopeSpans\x12(\n\x05spans\x18\x02 \x03(\x0b\x32\x19.divi.proto.trace.v1.Span\"\xbe\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\"O\n\x08SpanKind\x12\x16\n\x12SPAN_KIND_FUNCTION\x10\x00\x12\x11\n\rSPAN_KIND_LLM\x10\x01\x12\x18\n\x14SPAN_KIND_EVALUATION\x10\x02\x42\rZ\x0bservices/pbb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -36,7 +36,7 @@ if not _descriptor._USE_C_DESCRIPTORS:
   _globals['_SCOPESPANS']._serialized_start=91
   _globals['_SCOPESPANS']._serialized_end=145
   _globals['_SPAN']._serialized_start=148
-  _globals['_SPAN']._serialized_end=440
+  _globals['_SPAN']._serialized_end=466
   _globals['_SPAN_SPANKIND']._serialized_start=387
-  _globals['_SPAN_SPANKIND']._serialized_end=440
+  _globals['_SPAN_SPANKIND']._serialized_end=466
 # @@protoc_insertion_point(module_scope)

divi/proto/trace/v1/trace_pb2.pyi

@@ -19,8 +19,10 @@ class Span(_message.Message):
         __slots__ = ()
         SPAN_KIND_FUNCTION: _ClassVar[Span.SpanKind]
         SPAN_KIND_LLM: _ClassVar[Span.SpanKind]
+        SPAN_KIND_EVALUATION: _ClassVar[Span.SpanKind]
     SPAN_KIND_FUNCTION: Span.SpanKind
     SPAN_KIND_LLM: Span.SpanKind
+    SPAN_KIND_EVALUATION: Span.SpanKind
     TRACE_ID_FIELD_NUMBER: _ClassVar[int]
     SPAN_ID_FIELD_NUMBER: _ClassVar[int]
     PARENT_SPAN_ID_FIELD_NUMBER: _ClassVar[int]

divi/services/datapark/datapark.py

@@ -4,13 +4,14 @@ 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
+from typing_extensions import List, Mapping
 
 import divi
+from divi.evaluation.evaluator import EvaluationScore
 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 divi.signals.trace import TraceSignal
 
 
 class DataPark(Service):
@@ -58,6 +59,8 @@ class DataPark(Service):
         completion: ChatCompletion,
     ) -> None:
         hex_span_id = span_id.hex()
+        str_trace_id = str(trace_id)
+
         self.post_concurrent(
             {
                 "/api/v1/chat/completions/input": {
@@ -66,8 +69,23 @@ class DataPark(Service):
                 },
                 "/api/v1/chat/completions": {
                     "span_id": hex_span_id,
-                    "trace_id": str(trace_id),
+                    "trace_id": str_trace_id,
                     "data": completion.model_dump(),
                 },
             }
         )
+
+    def create_scores(
+        self,
+        span_id: bytes,
+        trace_id: UUID4,
+        scores: List[EvaluationScore],
+    ) -> None:
+        self.post(
+            "/api/v1/chat/completions/scores",
+            payload={
+                "span_id": span_id.hex(),
+                "trace_id": str(trace_id),
+                "data": [score.model_dump() for score in scores],
+            },
+        )

divi/services/init.py

@@ -7,7 +7,9 @@ from divi.services.datapark import init as init_datapark
 
 def init():
     if not divi._auth:
-        divi._auth = init_auth()
+        divi._auth = init_auth(
+            api_key="divi-aa31aef9-bb4c-4a98-aaad-7e12bdacec83"
+        )
     if not divi._datapark:
         divi._datapark = init_datapark()
     # TODO - Uncomment this when the core service is ready

divi/session/session.py

@@ -1,7 +1,7 @@
 from typing import Optional, TypedDict
 from uuid import uuid4
 
-from divi.signals.trace.trace import Trace
+from divi.signals.trace import Trace
 
 
 class SessionExtra(TypedDict, total=False):

divi/session/setup.py

@@ -3,8 +3,8 @@ from typing_extensions import Optional
 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
+from divi.signals.span import Span
+from divi.signals.trace import Trace
 
 
 def init_session(name: Optional[str] = None) -> Session:

divi/signals/__init__.py

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

divi/signals/span.py

@@ -0,0 +1,83 @@
+import os
+import time
+from enum import Enum
+from typing import Any, Mapping, Optional
+
+from pydantic import UUID4
+
+import divi
+from divi.proto.common.v1.common_pb2 import KeyValue
+from divi.proto.trace.v1.trace_pb2 import ScopeSpans
+from divi.proto.trace.v1.trace_pb2 import Span as SpanProto
+
+
+class Kind(int, Enum):
+    """Enum for the kind of span."""
+
+    function = SpanProto.SpanKind.SPAN_KIND_FUNCTION
+    llm = SpanProto.SpanKind.SPAN_KIND_LLM
+    evaluation = SpanProto.SpanKind.SPAN_KIND_EVALUATION
+
+
+class Span:
+    def __init__(
+        self,
+        kind: Kind = Kind.function,
+        name: Optional[str] = None,
+        metadata: Optional[Mapping[str, Any]] = None,
+    ):
+        # span_id is a FixedString(8)
+        self.span_id: bytes = self._generate_span_id()
+        self.name = name
+        self.kind = kind
+        self.metadata = metadata
+        self.start_time_unix_nano: int | None = None
+        self.end_time_unix_nano: int | None = None
+
+        self.trace_id: UUID4 | None = None
+        self.parent_span_id: bytes | None = None
+
+    @property
+    def signal(self) -> SpanProto:
+        signal: SpanProto = SpanProto(
+            name=self.name,
+            span_id=self.span_id,
+            kind=SpanProto.SpanKind.Name(self.kind),
+            start_time_unix_nano=self.start_time_unix_nano,
+            end_time_unix_nano=self.end_time_unix_nano,
+            trace_id=self.trace_id.bytes if self.trace_id else None,
+            parent_span_id=self.parent_span_id,
+        )
+        signal.metadata.extend(
+            KeyValue(key=k, value=v)
+            for k, v in (self.metadata or dict()).items()
+        )
+        return signal
+
+    @classmethod
+    def _generate_span_id(cls) -> bytes:
+        return os.urandom(8)
+
+    def start(self):
+        """Start the span by recording the current time in nanoseconds."""
+        self.start_time_unix_nano = time.time_ns()
+        self.upsert_span()
+
+    def end(self):
+        """End the span by recording the end time in nanoseconds."""
+        if self.start_time_unix_nano is None:
+            raise ValueError("Span must be started before ending.")
+        self.end_time_unix_nano = time.time_ns()
+        self.upsert_span()
+
+    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
+
+    def upsert_span(self):
+        """Upsert span with datapark."""
+        if divi._datapark and self.trace_id:
+            divi._datapark.create_spans(
+                self.trace_id, ScopeSpans(spans=[self.signal])
+            )

divi/signals/trace.py

@@ -0,0 +1,79 @@
+from datetime import UTC, datetime
+from typing import Optional
+from uuid import uuid4
+
+from pydantic import UUID4
+from typing_extensions import TypedDict
+
+import divi
+
+
+class NullTime(TypedDict, total=False):
+    """Null time"""
+
+    Time: str
+    """Time in iso format"""
+    Valid: bool
+    """Valid"""
+
+
+class TraceSignal(TypedDict, total=False):
+    """Trace request"""
+
+    id: str
+    """Trace ID UUID4"""
+    start_time: str
+    """Start time in iso format"""
+    end_time: NullTime
+    """End time in iso format"""
+    name: Optional[str]
+
+
+class Trace:
+    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:
+            raise ValueError("Trace must be started.")
+        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(
+                Time=self.end_time,
+                Valid=True,
+            )
+        return signal
+
+    @staticmethod
+    def unix_nano_to_iso(unix_nano: int) -> str:
+        return datetime.utcfromtimestamp(unix_nano / 1e9).isoformat()
+
+    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]
+            )

{divi-0.0.1.dev28.dist-info → divi-0.0.1.dev47.dist-info}/METADATA

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

{divi-0.0.1.dev28.dist-info → divi-0.0.1.dev47.dist-info}/RECORD

@@ -1,12 +1,14 @@
-divi/__init__.py,sha256=IcRE_HzskERpnkps5ABPTGxtsLGcRN8eRTjozDU5OMU,485
+divi/__init__.py,sha256=e3T9Znmiwfihr_X4Ly0Z0yoW23xGNQowPz6lCfXcXPI,519
 divi/utils.py,sha256=fXkjoyo_Lh8AZliKICOP460m0czUcNQjcEcceJbaOVA,1439
 divi/decorators/__init__.py,sha256=HkyWdC1ctTsVFucCWCkj57JB4NmwONus1d2S2dUbvs4,110
-divi/decorators/obs_openai.py,sha256=CGX60JgOKsN3E06yCq2L8baN9ZLiJJm3ukBsaN3IBBg,1324
-divi/decorators/observable.py,sha256=Ko_Pmw8l_T9mNp0corXkWqawwNEmQN7gf7rcAq-DB9M,3829
+divi/decorators/collect.py,sha256=5iUxAnbHYx4ISkFg64IK_4miGdrWgbOXLJxKz8lGIv8,1074
+divi/decorators/obs_openai.py,sha256=ouw3GYDFg6S27tcUzY0dIqz8JX_JM8IOXttzo7HK7nk,1359
+divi/decorators/observable.py,sha256=isUS3P_07wbZBj2UcRAoYNDceQTIn6zdein3-PWVsi8,2289
+divi/decorators/observe.py,sha256=I2RVsp2WQep6iTLSxkAlMP8wiRsSYiiYrxR2hJzPxcI,1211
 divi/evaluation/__init__.py,sha256=3qMHWu_zBh6FJa6-1dZZEWiAblQZurn5doa0OjGvDGs,93
-divi/evaluation/evaluate.py,sha256=4Li8NciQXcgmuAYRGFaB5C3EzMZR79n8IW0ubJBVrbQ,1789
-divi/evaluation/evaluator.py,sha256=22URc2pqg6WSYT-xUavxI8eT5w6t-mdVyTDZy-Qe4LU,5395
-divi/evaluation/prompts.py,sha256=yRFP3QjtcMflZ-jsFqXSu9RHlQA1cEVgDXo5D4VC6lo,662
+divi/evaluation/evaluate.py,sha256=lVMCw5vHGa5sJvUyhVDZ9m3Sgl4baCjWhw2OKazhvgM,1861
+divi/evaluation/evaluator.py,sha256=ulTyfSg2JXxzCCL7hRsn-EBb9UKcpQFA6rVT42mouVQ,5819
+divi/evaluation/prompts.py,sha256=qiv7TljwV8NTy0iLS2GEWIDFFNXhHKUlgVb-WoZhm4Q,970
 divi/evaluation/scores.py,sha256=ZgSxfve-ZivX3WU4TGcgPOSpUQVMbG5a15IQNPeq_bQ,173
 divi/proto/common/v1/common.proto,sha256=Rx8wr0_tOtQ1NseTMnsav4ApD1MDALzQDBA2IvLRTU0,1775
 divi/proto/common/v1/common_pb2.py,sha256=br61OHQVAi6SI3baFcb5xJv2Xd-AZ04A19xeSjLNMXo,2442
@@ -18,12 +20,12 @@ divi/proto/core/health/v1/health_service_pb2_grpc.py,sha256=YmlO94d-G71YBW1XZDSb
 divi/proto/metric/v1/metric.proto,sha256=YHRMLUW-MtakHuibR3PJ0s2w5KgV12kc4737iHw0DTk,585
 divi/proto/metric/v1/metric_pb2.py,sha256=uvBhyy8QpaES3Jl82yVfsGazW5654XpRnsdGlpVgIRE,1974
 divi/proto/metric/v1/metric_pb2.pyi,sha256=S7ipsojkD7QZAYefDE4b3PO99Yzc6mOdtSLxH3-b67A,1304
-divi/proto/trace/v1/trace.proto,sha256=mh1nzEgufzRTJx3p8NNute-ozEwEYwClWJTdWUGGVA8,1284
-divi/proto/trace/v1/trace_pb2.py,sha256=CuTkSSvhxCa1bk3Ku7tgLqRSovp_Gi52CZ0zLcLP2Ew,2327
-divi/proto/trace/v1/trace_pb2.pyi,sha256=rPo2Oa3NWrINE_dyOVU9HUYHo5LY82Bm5TMenj5dnK8,2136
+divi/proto/trace/v1/trace.proto,sha256=tPRIgBZB5KOKj7AoD3NoDZvLwoiJkbLiLqW53Ah-2-0,1367
+divi/proto/trace/v1/trace_pb2.py,sha256=zMuQO5mN2xl11USHkhi0lLwBAPlYXRU_UG1r0Uu3mJg,2369
+divi/proto/trace/v1/trace_pb2.pyi,sha256=k4dHYKAusH4I-XSW9KP3maogSWdRL7hVy8HCHhqFWzM,2231
 divi/services/__init__.py,sha256=TcVJ_gKxyPIcwhT9GgttqHeyk0icW44uE285KmUiyh4,185
 divi/services/finish.py,sha256=XKPKGJ5cWd5H95G_VpIOlOZOLrcf9StoTs7ayRic2jY,173
-divi/services/init.py,sha256=BFSsW6UzeujfHba8pvVPnEYmW0Y8VVFQLuIXI5Z2Mu4,436
+divi/services/init.py,sha256=dwXXXbf1-V0iAHZOETiv527TZQ07-waMIR5cSiU3QjI,509
 divi/services/service.py,sha256=539MhcYfMvsVGjDdu0UtYSZnL2cloaPeYeOSMl2eUy8,1532
 divi/services/auth/__init__.py,sha256=PIQ9rQ0jcRqcy03a3BOY7wbzwluIRG_4kI_H4J4mRFk,74
 divi/services/auth/auth.py,sha256=eRcE6Kq8jbBr6YL93HCGDIoga90SoZf3ogOAKeza9WY,445
@@ -34,16 +36,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=d2pbrzVJtR3mNW1eQpbm-Wca-SvcfJqT7IuaQy7yHT0,2285
+divi/services/datapark/datapark.py,sha256=f-qE2kmkLAniIj9mOP3nCbI3A3qkfIUnoVekwQ5w0QE,2781
 divi/services/datapark/init.py,sha256=C32f9t3eLsxcYNqEyheh6nW455G2oR0YhhdqBcbN3ec,92
 divi/session/__init__.py,sha256=6lYemv21VQCIHx-xIdi7BxXcPNxVdvE60--8ArReUew,82
-divi/session/session.py,sha256=LlB2W2qGo0Vf-0L0CTQoXfzg_gCGpf0MTFsXQW7E6i4,817
-divi/session/setup.py,sha256=XU4_wXqdIg91dZSttjmjlC6u3hIyV0Nn_xqRw8db7os,1405
+divi/session/session.py,sha256=QxtEezI447PbtKG2U6cxL1ACae55e8nFfTufAY8pEYI,811
+divi/session/setup.py,sha256=SHNzCuvOzlrlBJj34_jbzhfa6SXX3oaXrcG8bN0-Xvo,1398
 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=OsfrZPHp241_NN8W79U4O69HsHQajez_d3rz6yJRN9s,4508
-divi-0.0.1.dev28.dist-info/METADATA,sha256=I2KiV8mzi0AnMFpUaWYHGKE0ETSDKvS6wcMMx08vLVw,497
-divi-0.0.1.dev28.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-divi-0.0.1.dev28.dist-info/licenses/LICENSE,sha256=5OJuZ4wMMEV0DgF0tofhAlS_KLkaUsZwwwDS2U_GwQ0,1063
-divi-0.0.1.dev28.dist-info/RECORD,,
+divi/signals/__init__.py,sha256=wfSkkCwkRsFP4aLj8aGHk_k6Y50P5yN44WWlO3XyW18,43
+divi/signals/span.py,sha256=FQWql6ivAeXGk1HPZCsCjL5mXW6S6Nn9SmOiKH4aXik,2629
+divi/signals/trace.py,sha256=IoYeTfd6x_Xmxcp4HbFSEne0d48hol4ng2Mb_AO8hZw,2144
+divi-0.0.1.dev47.dist-info/METADATA,sha256=3QEVpc6O2YUEyMtFJ8kJn3rd8y_xMHGc67ObLg20vfs,497
+divi-0.0.1.dev47.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+divi-0.0.1.dev47.dist-info/licenses/LICENSE,sha256=5OJuZ4wMMEV0DgF0tofhAlS_KLkaUsZwwwDS2U_GwQ0,1063
+divi-0.0.1.dev47.dist-info/RECORD,,

divi/signals/trace/__init__.py

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

divi/signals/trace/trace.py

@@ -1,151 +0,0 @@
-import os
-import time
-from datetime import UTC, datetime
-from typing import Any, Mapping, Optional
-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
-
-
-class NullTime(TypedDict, total=False):
-    """Null time"""
-
-    Time: str
-    """Time in iso format"""
-    Valid: bool
-    """Valid"""
-
-
-class TraceSignal(TypedDict, total=False):
-    """Trace request"""
-
-    id: str
-    """Trace ID UUID4"""
-    start_time: str
-    """Start time in iso format"""
-    end_time: NullTime
-    """End time in iso format"""
-    name: Optional[str]
-
-
-class Trace:
-    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:
-            raise ValueError("Trace must be started.")
-        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(
-                Time=self.end_time,
-                Valid=True,
-            )
-        return signal
-
-    @staticmethod
-    def unix_nano_to_iso(unix_nano: int) -> str:
-        return datetime.utcfromtimestamp(unix_nano / 1e9).isoformat()
-
-    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:
-    KIND_MAP = {
-        "function": SpanProto.SpanKind.SPAN_KIND_FUNCTION,
-        "llm": SpanProto.SpanKind.SPAN_KIND_LLM,
-    }
-
-    def __init__(
-        self,
-        kind: str = "function",
-        name: Optional[str] = None,
-        metadata: Optional[Mapping[str, Any]] = None,
-    ):
-        # span_id is a FixedString(8)
-        self.span_id: bytes = self._generate_span_id()
-        self.name = name
-        self.kind = kind
-        self.metadata = metadata
-        self.start_time_unix_nano: int | None = None
-        self.end_time_unix_nano: int | None = None
-
-        self.trace_id: UUID4 | None = None
-        self.parent_span_id: bytes | None = None
-
-    @property
-    def signal(self) -> SpanProto:
-        signal: SpanProto = SpanProto(
-            name=self.name,
-            span_id=self.span_id,
-            kind=self._get_kind(self.kind),
-            start_time_unix_nano=self.start_time_unix_nano,
-            end_time_unix_nano=self.end_time_unix_nano,
-            trace_id=self.trace_id.bytes if self.trace_id else None,
-            parent_span_id=self.parent_span_id,
-        )
-        signal.metadata.extend(
-            KeyValue(key=k, value=v)
-            for k, v in (self.metadata or dict()).items()
-        )
-        return signal
-
-    @classmethod
-    def _get_kind(cls, kind: str) -> SpanProto.SpanKind:
-        if (k := cls.KIND_MAP.get(kind)) is None:
-            raise ValueError(
-                f"Unknown kind: {kind}. Now allowed: {cls.KIND_MAP.keys()}"
-            )
-        return k
-
-    @classmethod
-    def _generate_span_id(cls) -> bytes:
-        return os.urandom(8)
-
-    def start(self):
-        """Start the span by recording the current time in nanoseconds."""
-        self.start_time_unix_nano = time.time_ns()
-
-    def end(self):
-        """End the span by recording the end time in nanoseconds."""
-        if self.start_time_unix_nano is None:
-            raise ValueError("Span must be started before ending.")
-        self.end_time_unix_nano = time.time_ns()
-
-    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