divi 0.0.1.dev28__tar.gz → 0.0.1.dev47__tar.gz

{divi-0.0.1.dev28 → divi-0.0.1.dev47}/PKG-INFO

@@ -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 → divi-0.0.1.dev47}/apps/graphql/src/datasources/datapark-api.ts

@@ -1,5 +1,5 @@
 import type { FetchResponse } from '@/types/response';
-import type { ChatInput, Span, Trace, UsageResult } from '@/types/types';
+import type { ChatInput, Score, Span, Trace, UsageResult } from '@/types/types';
 import { type AugmentedRequest, RESTDataSource } from '@apollo/datasource-rest';
 import type { KeyValueCache } from '@apollo/utils.keyvaluecache';
 
@@ -30,6 +30,12 @@ export class DataParkAPI extends RESTDataSource {
     return await this.get<FetchResponse<Span[]>>(`/api/trace/${traceId}/spans`);
   }
 
+  async getScores(traceId: string) {
+    return await this.get<FetchResponse<Score[]>>(
+      `/api/trace/${traceId}/scores`
+    );
+  }
+
   async getChatInput(spanId: string) {
     return await this.get<FetchResponse<ChatInput>>(
       `/api/v1/chat/completions/${spanId}/input`

{divi-0.0.1.dev28 → divi-0.0.1.dev47}/apps/graphql/src/resolvers.ts

@@ -26,6 +26,9 @@ export const resolvers: Resolvers = {
     spans: async (_, { trace_id }, { dataSources }) => {
       return (await dataSources.dataparkAPI.getSpans(trace_id)).data;
     },
+    scores: async (_, { trace_id }, { dataSources }) => {
+      return (await dataSources.dataparkAPI.getScores(trace_id)).data;
+    },
     chat_input: async (_, { span_id }, { dataSources }) => {
       return (await dataSources.dataparkAPI.getChatInput(span_id)).data;
     },

{divi-0.0.1.dev28 → divi-0.0.1.dev47}/apps/graphql/src/schema.graphql

@@ -20,6 +20,8 @@ type Query {
     end_time: Int
     group_by: GroupingKey
   ): [UsageResult!]
+  "Fetch all scores by trace id"
+  scores(trace_id: ID!): [Score!]
 }
 
 "GroupingKey is an enum that represents the key used to group usage results"
@@ -70,6 +72,7 @@ type Trace {
 enum Kind {
   SPAN_KIND_FUNCTION
   SPAN_KIND_LLM
+  SPAN_KIND_EVALUATION
 }
 
 "Span is a record of a single unit of work within a trace"
@@ -79,8 +82,8 @@ type Span {
   trace_id: ID!
   name: String!
   start_time: String!
-  end_time: String!
-  duration: Float!
+  end_time: NullTime!
+  duration: Float
   kind: Kind!
   metadata: [KeyValue!]
 }
@@ -113,6 +116,14 @@ type UsageResult {
   date: Int
 }
 
+"Score is a record of a score for a trace"
+type Score {
+  span_id: ID!
+  name: String!
+  score: Float!
+  representative_reasoning: String!
+}
+
 "Mutation is a collection of mutations that can be made to the API"
 type Mutation {
   "Auth Mutations"

{divi-0.0.1.dev28 → divi-0.0.1.dev47}/apps/graphql/src/types/types.d.ts

@@ -92,6 +92,7 @@ export type KeyValue = {
 
 /** Kind is an enum that represents the type of span */
 export enum Kind {
+  SpanKindEvaluation = 'SPAN_KIND_EVALUATION',
   SpanKindFunction = 'SPAN_KIND_FUNCTION',
   SpanKindLlm = 'SPAN_KIND_LLM'
 }
@@ -201,6 +202,8 @@ export type Query = {
   completion_usage?: Maybe<Array<UsageResult>>;
   /** Fetch current user */
   me: User;
+  /** Fetch all scores by trace id */
+  scores?: Maybe<Array<Score>>;
   /** Fetch all spans by trace id */
   spans?: Maybe<Array<Span>>;
   /** Fetch traces by session id */
@@ -224,6 +227,12 @@ export type QueryCompletion_UsageArgs = {
 };
 
 
+/** Query is a collection of queries that can be made to the API */
+export type QueryScoresArgs = {
+  trace_id: Scalars['ID']['input'];
+};
+
+
 /** Query is a collection of queries that can be made to the API */
 export type QuerySpansArgs = {
   trace_id: Scalars['ID']['input'];
@@ -249,11 +258,20 @@ export type RevokeApiKeyResponse = MutationResponse & {
   success: Scalars['Boolean']['output'];
 };
 
+/** Score is a record of a score for a trace */
+export type Score = {
+  __typename?: 'Score';
+  name: Scalars['String']['output'];
+  representative_reasoning: Scalars['String']['output'];
+  score: Scalars['Float']['output'];
+  span_id: Scalars['ID']['output'];
+};
+
 /** Span is a record of a single unit of work within a trace */
 export type Span = {
   __typename?: 'Span';
-  duration: Scalars['Float']['output'];
-  end_time: Scalars['String']['output'];
+  duration?: Maybe<Scalars['Float']['output']>;
+  end_time: NullTime;
   id: Scalars['ID']['output'];
   kind: Kind;
   metadata?: Maybe<Array<KeyValue>>;
@@ -405,6 +423,7 @@ export type ResolversTypes = {
   NullTime: ResolverTypeWrapper<NullTime>;
   Query: ResolverTypeWrapper<{}>;
   RevokeAPIKeyResponse: ResolverTypeWrapper<RevokeApiKeyResponse>;
+  Score: ResolverTypeWrapper<Score>;
   Span: ResolverTypeWrapper<Span>;
   String: ResolverTypeWrapper<Scalars['String']['output']>;
   Trace: ResolverTypeWrapper<Trace>;
@@ -433,6 +452,7 @@ export type ResolversParentTypes = {
   NullTime: NullTime;
   Query: {};
   RevokeAPIKeyResponse: RevokeApiKeyResponse;
+  Score: Score;
   Span: Span;
   String: Scalars['String']['output'];
   Trace: Trace;
@@ -536,6 +556,7 @@ export type QueryResolvers<ContextType = DataSourceContext, ParentType extends R
   chat_input?: Resolver<Maybe<ResolversTypes['ChatInput']>, ParentType, ContextType, RequireFields<QueryChat_InputArgs, 'span_id'>>;
   completion_usage?: Resolver<Maybe<Array<ResolversTypes['UsageResult']>>, ParentType, ContextType, RequireFields<QueryCompletion_UsageArgs, 'start_time'>>;
   me?: Resolver<ResolversTypes['User'], ParentType, ContextType>;
+  scores?: Resolver<Maybe<Array<ResolversTypes['Score']>>, ParentType, ContextType, RequireFields<QueryScoresArgs, 'trace_id'>>;
   spans?: Resolver<Maybe<Array<ResolversTypes['Span']>>, ParentType, ContextType, RequireFields<QuerySpansArgs, 'trace_id'>>;
   traces?: Resolver<Maybe<Array<ResolversTypes['Trace']>>, ParentType, ContextType, RequireFields<QueryTracesArgs, 'session_id'>>;
   user?: Resolver<ResolversTypes['User'], ParentType, ContextType, RequireFields<QueryUserArgs, 'id'>>;
@@ -548,9 +569,17 @@ export type RevokeApiKeyResponseResolvers<ContextType = DataSourceContext, Paren
   __isTypeOf?: IsTypeOfResolverFn<ParentType, ContextType>;
 };
 
+export type ScoreResolvers<ContextType = DataSourceContext, ParentType extends ResolversParentTypes['Score'] = ResolversParentTypes['Score']> = {
+  name?: Resolver<ResolversTypes['String'], ParentType, ContextType>;
+  representative_reasoning?: Resolver<ResolversTypes['String'], ParentType, ContextType>;
+  score?: Resolver<ResolversTypes['Float'], ParentType, ContextType>;
+  span_id?: Resolver<ResolversTypes['ID'], ParentType, ContextType>;
+  __isTypeOf?: IsTypeOfResolverFn<ParentType, ContextType>;
+};
+
 export type SpanResolvers<ContextType = DataSourceContext, ParentType extends ResolversParentTypes['Span'] = ResolversParentTypes['Span']> = {
-  duration?: Resolver<ResolversTypes['Float'], ParentType, ContextType>;
-  end_time?: Resolver<ResolversTypes['String'], ParentType, ContextType>;
+  duration?: Resolver<Maybe<ResolversTypes['Float']>, ParentType, ContextType>;
+  end_time?: Resolver<ResolversTypes['NullTime'], ParentType, ContextType>;
   id?: Resolver<ResolversTypes['ID'], ParentType, ContextType>;
   kind?: Resolver<ResolversTypes['Kind'], ParentType, ContextType>;
   metadata?: Resolver<Maybe<Array<ResolversTypes['KeyValue']>>, ParentType, ContextType>;
@@ -618,6 +647,7 @@ export type Resolvers<ContextType = DataSourceContext> = {
   NullTime?: NullTimeResolvers<ContextType>;
   Query?: QueryResolvers<ContextType>;
   RevokeAPIKeyResponse?: RevokeApiKeyResponseResolvers<ContextType>;
+  Score?: ScoreResolvers<ContextType>;
   Span?: SpanResolvers<ContextType>;
   Trace?: TraceResolvers<ContextType>;
   UpdateAPIKeyResponse?: UpdateApiKeyResponseResolvers<ContextType>;

{divi-0.0.1.dev28 → divi-0.0.1.dev47}/apps/web/app/dashboard/traces/(children)/[id]/components/Span.tsx

@@ -1,5 +1,10 @@
+import Highlighter from '@/components/Highter';
 import type { ExtendedSpan } from '@/lib/types/span';
 import { formatDate } from '@/lib/utils';
+import {
+  Kind,
+  type Score,
+} from '@workspace/graphql-client/src/types.generated';
 import {
   Accordion,
   AccordionContent,
@@ -16,7 +21,6 @@ import {
 } from '@workspace/ui/components/card';
 import { Timer } from 'lucide-react';
 import type * as React from 'react';
-import Highlighter from '@/components/Highter';
 
 interface SpanProps {
   span: ExtendedSpan;
@@ -30,11 +34,18 @@ export function Span({ span }: SpanProps) {
       <GeneralInfo span={span} />
       <Accordion
         type="multiple"
-        defaultValue={['properties', 'Input', 'Output']}
+        defaultValue={['properties', 'Input', 'Output', 'scores']}
       >
         <AccordionProperties span={span} />
-        <AccordionJsonCards name="Input" datas={messages} />
-        <AccordionJsonCards name="Output" datas={choices} />
+        {span.kind === Kind.SpanKindLlm && (
+          <>
+            <AccordionJsonCards name="Input" datas={messages} />
+            <AccordionJsonCards name="Output" datas={choices} />
+          </>
+        )}
+        {span.kind === Kind.SpanKindEvaluation && (
+          <AccordionScores scores={span.scores ?? []} />
+        )}
       </Accordion>
     </>
   );
@@ -57,10 +68,35 @@ function GeneralInfo({ span }: SpanProps) {
   );
 }
 
+function AccordionScores({ scores }: { scores: Score[] }) {
+  return (
+    <AccordionItem value="scores" className="px-4">
+      <AccordionTrigger className="hover:no-underline">Scores</AccordionTrigger>
+      <AccordionContent className="flex flex-col gap-3">
+        {scores.map((score) => (
+          <PropertyCard
+            title={score.name
+              .split('_')
+              .map((n) => n.charAt(0).toUpperCase() + n.slice(1).toLowerCase())
+              .join(' ')}
+            key={score.name}
+          >
+            <PropertyItem label="score" value={score.score} />
+            <PropertyItem
+              label="explanation"
+              value={score.representative_reasoning}
+            />
+          </PropertyCard>
+        ))}
+      </AccordionContent>
+    </AccordionItem>
+  );
+}
+
 function AccordionProperties({ span }: SpanProps) {
   // TODO select the properties to show
   const startTime = formatDate(span.start_time);
-  const endTime = formatDate(span.end_time);
+  const endTime = span.end_time.Valid ? formatDate(span.end_time.Time) : 'N/A';
   const usage = span.completion?.usage;
   // input should pop messages
   const { messages, ...inputs } = span.input ?? {};
@@ -128,9 +164,9 @@ function PropertyItem({
   value: string | number | boolean | null;
 }) {
   return (
-    <div className="flex items-center justify-between">
+    <div className="flex items-start justify-between gap-4">
       <span className="text-sm">{label}</span>
-      <span className="truncate font-light text-sm">{String(value)}</span>
+      <span className="break-words font-light text-sm">{String(value)}</span>
     </div>
   );
 }

{divi-0.0.1.dev28 → divi-0.0.1.dev47}/apps/web/app/dashboard/traces/(children)/[id]/components/trace-board.tsx

@@ -16,6 +16,11 @@ export function TraceBoard({ spans, direction }: TraceBoardProps) {
   const selectAction = (_data: ExtendedSpan, index: number) => {
     setIndex(index);
   };
+  if (!spans || spans.length === 0) {
+    return (
+      <div className="flex h-full items-center justify-center">No data</div>
+    );
+  }
 
   return (
     <ResponsiveResizable

{divi-0.0.1.dev28 → divi-0.0.1.dev47}/apps/web/app/dashboard/traces/(children)/[id]/components/trace-chart.tsx

@@ -32,6 +32,9 @@ const chartConfig = {
   SPAN_KIND_LLM: {
     color: 'var(--chart-2)',
   },
+  SPAN_KIND_EVALUATION: {
+    color: 'var(--chart-3)',
+  },
 } satisfies ChartConfig;
 
 interface TraceWaterfallChartProps {
@@ -107,7 +110,7 @@ export function TraceWaterfallChart({
           <LabelList
             dataKey="name"
             position="insideLeft"
-            className="fill-background"
+            className="fill-card"
           />
         </Bar>
       </BarChart>

{divi-0.0.1.dev28 → divi-0.0.1.dev47}/apps/web/app/dashboard/usages/actions.ts

@@ -8,7 +8,7 @@ import type {
   GroupingKey,
   UsageResult,
 } from '@workspace/graphql-client/src/types.generated';
-import { eachDayOfInterval, getUnixTime } from 'date-fns';
+import { addDays, eachDayOfInterval, getUnixTime, startOfDay } from 'date-fns';
 import { cache } from 'react';
 
 export const getCompletionUsage = cache(
@@ -25,7 +25,9 @@ export const getCompletionUsage = cache(
       query: GetCompletionUsageDocument,
       variables: {
         startTime: getUnixTime(startTime),
-        endTime: getUnixTime(endTime ?? new Date()),
+        endTime: getUnixTime(
+          endTime ? addDays(startOfDay(endTime), 1) : new Date()
+        ),
         groupBy,
       },
       context,

{divi-0.0.1.dev28 → divi-0.0.1.dev47}/apps/web/app/home/components/home-hero.tsx

@@ -9,7 +9,7 @@ export default function HomeHero() {
         <div className="flex flex-col items-center justify-center gap-8 py-20 lg:py-40">
           <div>
             <Button variant="secondary" size="sm" className="gap-4" asChild>
-              <Link href="https://docs.divine-agent.com">
+              <Link href="https://docs.divine-agent.com" target="_blank">
                 Read our launch article <MoveRight className="h-4 w-4" />
               </Link>
             </Button>

divi-0.0.1.dev47/apps/web/lib/server/evaluation.ts

@@ -0,0 +1,17 @@
+import { query } from '@/hooks/apolloClient';
+import { GetScoresDocument } from '@workspace/graphql-client/src/datapark/evaluation.generated';
+import { cache } from 'react';
+import { getAuthContext } from './auth';
+
+export const getScores = cache(async (traceId: string) => {
+  const context = await getAuthContext();
+  if (!context) {
+    return null;
+  }
+  const { data } = await query({
+    query: GetScoresDocument,
+    variables: { traceId },
+    context,
+  });
+  return data?.scores;
+});

{divi-0.0.1.dev28 → divi-0.0.1.dev47}/apps/web/lib/server/span.ts

@@ -2,9 +2,14 @@ import 'server-only';
 import { query } from '@/hooks/apolloClient';
 import type { Chat, ExtendedSpan } from '@/lib/types/span';
 import { GetSpansDocument } from '@workspace/graphql-client/src/datapark/traces.generated';
-import { Kind, type Span } from '@workspace/graphql-client/src/types.generated';
+import {
+  Kind,
+  type Score,
+  type Span,
+} from '@workspace/graphql-client/src/types.generated';
 import { cache } from 'react';
 import { getAuthContext } from './auth';
+import { getScores } from './evaluation';
 import { getChat } from './openai';
 
 /**
@@ -31,7 +36,10 @@ export const getSpans = cache(async (traceId: string) => {
 export const getTraceChartData = cache(
   async (traceId: string): Promise<ExtendedSpan[]> => {
     // get spans for a trace
-    let spans = await getSpans(traceId);
+    let [spans, scores] = await Promise.all([
+      getSpans(traceId),
+      getScores(traceId),
+    ]);
     if (!spans) {
       return [];
     }
@@ -39,12 +47,23 @@ export const getTraceChartData = cache(
     spans = sortSpans(spans);
     // get chat input and completion for LLM span
     const llmSpanIds = spans
-      .filter((span) => span.kind === Kind.SpanKindLlm)
+      .filter((span) => span.kind === Kind.SpanKindLlm && span.end_time.Valid)
       .map((s) => s.id);
     const chats = await Promise.all(llmSpanIds.map((id) => getChat(id)));
     const chatsMap = new Map<string, Chat>(
       chats.map((chat) => [chat.span_id, chat])
     );
+    const scoresMap = new Map<string, Score[]>();
+    // group scores by span_id
+    if (scores) {
+      for (const score of scores) {
+        const spanId = score.span_id;
+        if (!scoresMap.has(spanId)) {
+          scoresMap.set(spanId, []);
+        }
+        scoresMap.get(spanId)?.push(score);
+      }
+    }
     // calculate relative start_time with milliseconds
     const startTime = new Date(spans[0]?.start_time).getTime();
     return spans.map((span) => {
@@ -54,6 +73,11 @@ export const getTraceChartData = cache(
         relative_start_time: new Date(span.start_time).getTime() - startTime,
         input: chat?.input,
         completion: chat?.completion,
+        duration: span.end_time.Valid
+          ? span.duration
+          : Date.now() - new Date(span.start_time).getTime(),
+        scores:
+          span.kind === Kind.SpanKindEvaluation ? scoresMap.get(span.id) : [],
       };
     });
   }

{divi-0.0.1.dev28 → divi-0.0.1.dev47}/apps/web/lib/types/span.d.ts

@@ -1,4 +1,7 @@
-import type { Span } from '@workspace/graphql-client/src/types.generated';
+import type {
+  Score,
+  Span,
+} from '@workspace/graphql-client/src/types.generated';
 import type { ChatInput } from '@workspace/graphql-client/src/types.generated';
 import type { ChatCompletion } from 'openai/resources/index.mjs';
 
@@ -10,6 +13,7 @@ export interface ExtendedSpan extends Span {
   relative_start_time: number;
   input?: ChatInput;
   completion?: ChatCompletion;
+  scores?: Score[];
 }
 
 /**

{divi-0.0.1.dev28 → divi-0.0.1.dev47}/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-0.0.1.dev47/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-0.0.1.dev28 → divi-0.0.1.dev47}/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-0.0.1.dev47/divi/decorators/observable.py

@@ -0,0 +1,91 @@
+import functools
+from typing import (
+    Any,
+    Callable,
+    Generic,
+    Mapping,
+    Optional,
+    ParamSpec,
+    Protocol,
+    TypeVar,
+    Union,
+    overload,
+    runtime_checkable,
+)
+
+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.signals.span import Kind, Span
+
+R = TypeVar("R", covariant=True)
+P = ParamSpec("P")
+
+
+@runtime_checkable
+class WithSessionExtra(Protocol, Generic[P, R]):
+    def __call__(
+        self,
+        *args: P.args,
+        session_extra: Optional[SessionExtra] = None,  # type: ignore[valid-type]
+        **kwargs: P.kwargs,
+    ) -> R: ...
+
+
+@overload
+def observable(func: Callable[P, R]) -> WithSessionExtra[P, R]: ...
+
+
+@overload
+def observable(
+    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]]: ...
+
+
+def observable(
+    *args, **kwargs
+) -> Union[Callable, Callable[[Callable], Callable]]:
+    """Observable decorator factory."""
+
+    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
+            )
+
+            # 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
+
+        return wrapper
+
+    # Function Decorator
+    if len(args) == 1 and callable(args[0]) and not kwargs:
+        return decorator(args[0])
+    # Factory Decorator
+    return decorator

divi-0.0.1.dev47/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