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

divi/__init__.py

@@ -2,8 +2,10 @@ from typing import Optional
 
 from . import proto
 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"
 
@@ -11,6 +13,7 @@ _session: Optional[Session] = None
 _core: Optional[Core] = None
 _auth: Optional[Auth] = None
 _datapark: Optional[DataPark] = None
+_evaluator: Optional[Evaluator] = None
 
-__version__ = "0.0.1.dev23"
-__all__ = ["proto", "obs_openai", "observable"]
+__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

@@ -2,7 +2,11 @@ import functools
 from collections.abc import Callable
 from typing import TYPE_CHECKING, TypeVar, Union
 
+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:
@@ -11,22 +15,34 @@ if TYPE_CHECKING:
 C = TypeVar("C", bound=Union["OpenAI", "AsyncOpenAI"])
 
 
-def _get_observable_create(create: Callable) -> Callable:
+def _get_observable_create(
+    create: Callable,
+    name: Optional[str] = None,
+    scores: Optional[list[Score]] = None,
+) -> Callable:
     @functools.wraps(create)
     def observable_create(*args, stream: bool = False, **kwargs):
-        decorator = observable(kind="llm")
+        decorator = observable(kind=Kind.llm, name=name, scores=scores)
         return decorator(create)(*args, stream=stream, **kwargs)
 
     # TODO Async Observable Create
     return observable_create if not is_async(create) else create
 
 
-def obs_openai(client: C) -> C:
+def obs_openai(
+    client: C,
+    name: Optional[str] = "Agent",
+    scores: Optional[list[Score]] = None,
+) -> C:
     """Make OpenAI client observable."""
     client.chat.completions.create = _get_observable_create(
-        client.chat.completions.create
+        client.chat.completions.create,
+        name=name,
+        scores=scores,
     )
     client.completions.create = _get_observable_create(
-        client.completions.create
+        client.completions.create,
+        name=name,
+        scores=scores,
     )
     return client

divi/decorators/observable.py

@@ -1,11 +1,8 @@
-import contextvars
 import functools
-import inspect
 from typing import (
     Any,
     Callable,
     Generic,
-    List,
     Mapping,
     Optional,
     ParamSpec,
@@ -16,24 +13,15 @@ from typing import (
     runtime_checkable,
 )
 
-from openai.types.chat import ChatCompletion
-
-import divi
-from divi.proto.trace.v1.trace_pb2 import ScopeSpans
+from divi.decorators.observe import observe
+from divi.evaluation.evaluate import evaluate_scores
+from divi.evaluation.scores import Score
 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,9 +39,10 @@ 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,
     metadata: Optional[Mapping[str, Any]] = None,
 ) -> Callable[[Callable[P, R]], WithSessionExtra[P, R]]: ...
 
@@ -63,50 +52,33 @@ 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)
 
     def decorator(func):
         @functools.wraps(func)
         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_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:
-                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,
-                )
+
+            # 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/__init__.py

@@ -0,0 +1,4 @@
+from .evaluator import Evaluator
+from .scores import Score
+
+__all__ = ["Evaluator", "Score"]

divi/evaluation/evaluate.py

@@ -0,0 +1,61 @@
+import os
+from typing import Optional
+
+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"
+
+
+def init_evaluator(config: Optional[EvaluatorConfig] = None):
+    _config = config or EvaluatorConfig()
+    api_key = _config.api_key if _config.api_key else os.getenv(OPENAI_API_KEY)
+    base_url = (
+        _config.base_url if _config.base_url else os.getenv(OPENAI_BASE_URL)
+    )
+    if api_key is None:
+        raise ValueError("API key is required for evaluator")
+    _config.api_key = api_key
+    _config.base_url = base_url
+    evaluator = Evaluator(_config)
+    return evaluator
+
+
+def evaluate_scores(
+    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)
+
+    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

@@ -0,0 +1,174 @@
+import asyncio
+import concurrent.futures
+import random
+from typing import List, Literal, Optional
+
+import openai
+from pydantic import BaseModel
+
+from divi.evaluation.prompts import PRESET_PROMPT, PROMPT_TEMPLATE
+from divi.evaluation.scores import Score
+
+
+class EvaluatorConfig:
+    def __init__(
+        self,
+        model: str = "gpt-4.1-nano",
+        temperature: float = 0.5,
+        max_concurrency: int = 10,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+    ):
+        self.model = model
+        self.api_key = api_key
+        self.base_url = base_url
+        self.temperature = temperature
+        self.max_concurrency = max_concurrency
+
+
+class EvaluationResult(BaseModel):
+    name: Score
+    judgment: bool
+    reasoning: str
+
+
+class EvaluationScore(BaseModel):
+    name: Score
+    score: float
+    representative_reasoning: str
+    all_evaluations: List[EvaluationResult]
+
+
+class Evaluator:
+    def __init__(self, config: Optional[EvaluatorConfig] = None):
+        self.config = config or EvaluatorConfig()
+        self.async_client = openai.AsyncOpenAI(
+            api_key=self.config.api_key, base_url=self.config.base_url
+        )
+        self.sync_client = openai.OpenAI(
+            api_key=self.config.api_key, base_url=self.config.base_url
+        )
+
+    @staticmethod
+    def generate_prompt(target: str, conversation: str, score: Score) -> str:
+        return PROMPT_TEMPLATE.format(
+            requirements=PRESET_PROMPT[score.value],
+            target=target,
+            conversation=conversation,
+        )
+
+    def _sync_evaluate_once(
+        self, target: str, conversation: str, score: Score
+    ) -> Optional[EvaluationResult]:
+        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,
+        )
+        result = response.choices[0].message.parsed
+        if result is not None:
+            result.name = score
+        return result
+
+    async def _async_evaluate_once(
+        self, target: str, conversation: str, score: Score
+    ) -> Optional[EvaluationResult]:
+        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,
+        )
+        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]
+    ) -> EvaluationScore:
+        n = len(evaluations)
+        true_count = sum(1 for e in evaluations if e.judgment is True)
+        score = true_count / n
+        majority_judgment = True if true_count >= (n / 2) else False
+        majority_reasons = [
+            e.reasoning for e in evaluations if e.judgment == majority_judgment
+        ]
+        representative_reasoning = (
+            random.choice(majority_reasons) if majority_reasons else ""
+        )
+        return EvaluationScore(
+            name=name,
+            score=score,
+            representative_reasoning=representative_reasoning,
+            all_evaluations=evaluations,
+        )
+
+    def _aggregate_results(
+        self, evaluations: List[EvaluationResult]
+    ) -> List[EvaluationScore]:
+        scores = {}
+        for evaluation in evaluations:
+            if evaluation.name not in scores:
+                scores[evaluation.name] = []
+            scores[evaluation.name].append(evaluation)
+
+        aggregated_results = [
+            self._aggregate_result(name, evals)
+            for name, evals in scores.items()
+        ]
+        return aggregated_results
+
+    def evaluate_sync(
+        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, target, conversation, score
+                )
+                for _ in range(n_rounds)
+                for score in scores
+            ]
+            evaluations = [
+                f.result() for f in concurrent.futures.as_completed(futures)
+            ]
+        return self._aggregate_results(
+            [e for e in evaluations if e is not None]
+        )
+
+    async def evaluate_async(
+        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(
+                    target, conversation, score
+                )
+
+        tasks = [sem_task(score) for _ in range(n_rounds) for score in scores]
+        evaluations = await asyncio.gather(*tasks)
+        return self._aggregate_results(
+            [e for e in evaluations if e is not None]
+        )
+
+    def evaluate(
+        self,
+        target: str,
+        conversation: str,
+        scores: list[Score],
+        n_rounds: int = 5,
+        mode: Literal["sync", "async"] = "sync",
+    ) -> List[EvaluationScore]:
+        if mode == "async":
+            return asyncio.run(
+                self.evaluate_async(target, conversation, scores, n_rounds)
+            )
+        return self.evaluate_sync(target, conversation, scores, n_rounds)

divi/evaluation/prompts.py

@@ -0,0 +1,19 @@
+PROMPT_TEMPLATE = (
+    "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"
+    "{{\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": "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/evaluation/scores.py

@@ -0,0 +1,8 @@
+from enum import Enum
+
+
+class Score(str, Enum):
+    """Enum for score types."""
+
+    task_completion = "task_completion"
+    instruction_adherence = "instruction_adherence"

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

@@ -1,12 +1,15 @@
 import divi
 from divi.services.auth import init as init_auth
-from divi.services.core import init as init_core
+
+# 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()
+        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

@@ -1,12 +1,13 @@
 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:
+def init_session(name: Optional[str] = None) -> Session:
     """init initializes the services and the Run"""
     init_services()
     session = Session(name=name)
@@ -29,7 +30,9 @@ def setup(
 
     # init the session if not already initialized
     if not divi._session:
-        divi._session = init_session(name=session_extra.get('session_name') or span.name)
+        divi._session = init_session(
+            name=session_extra.get("session_name") or span.name
+        )
 
     # setup trace
     trace = session_extra.get("trace") or Trace(divi._session.id, span.name)

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.dev23.dist-info → divi-0.0.1.dev47.dist-info}/METADATA

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

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

@@ -1,9 +1,15 @@
-divi/__init__.py,sha256=lcpSVEjIcXy-e9CuJB8V1izwaoEK8YpH8FtkzH0IqJc,396
+divi/__init__.py,sha256=e3T9Znmiwfihr_X4Ly0Z0yoW23xGNQowPz6lCfXcXPI,519
 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=GI9c6gFArA6pTfa3EPednMtRqV2zIlofuTDdUYtS_x4,963
-divi/decorators/observable.py,sha256=xKapTyHL50mUyIDhDSYIXR1yZIdHTuw6WtI4H0j1xDE,3464
+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=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
 divi/proto/common/v1/common_pb2.pyi,sha256=LmTpFFLxHg2a_qPIdNXXwGEMkbiDcTJdarR9eC-6Fq8,2133
@@ -14,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=JVzRQ1m1DTHXFVGUMYnsv-vRvzCO8XFdR6MjIwOL_NY,433
+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
@@ -30,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=NeCxCb-uYhkKnOEiw8dBQHz0DEL8j1oxzQY3cBAmHbo,1380
+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.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,,
+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