Skip to content

Commit

Permalink
create generic llmobs decorator
Browse files Browse the repository at this point in the history
  • Loading branch information
ncybul committed Jan 3, 2025
1 parent 6613185 commit c7fa2f0
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 0 deletions.
18 changes: 18 additions & 0 deletions ddtrace/llmobs/_llmobs.py
Original file line number Diff line number Diff line change
Expand Up @@ -572,6 +572,24 @@ def retrieval(
log.warning(SPAN_START_WHILE_DISABLED_WARNING)
return cls._instance._start_span("retrieval", name=name, session_id=session_id, ml_app=ml_app)

@classmethod
def undefined_kind(
cls, name: Optional[str] = None, session_id: Optional[str] = None, ml_app: Optional[str] = None
) -> Span:
"""
Trace a generic operation. Used when the span kind is not specified,
:param str name: The name of the traced operation. If not provided, a default value of "workflow" will be set.
:param str session_id: The ID of the underlying user session. Required for tracking sessions.
:param str ml_app: The name of the ML application that the agent is orchestrating. If not provided, the default
value will be set to the value of `DD_LLMOBS_ML_APP`.
:returns: The Span object representing the traced operation.
"""
if cls.enabled is False:
log.warning(SPAN_START_WHILE_DISABLED_WARNING)
return cls._instance._start_span("undefined", name=name, session_id=session_id, ml_app=ml_app)

@classmethod
def annotate(
cls,
Expand Down
36 changes: 36 additions & 0 deletions ddtrace/llmobs/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -253,6 +253,41 @@ def wrapper(*args, **kwargs):

return decorator

def _generic_decorator(operation_kind):
def decorator(
original_func: Optional[Callable] = None,
kind: Optional[str] = None,
model_name: Optional[str] = None,
model_provider: Optional[str] = None,
name: Optional[str] = None,
session_id: Optional[str] = None,
ml_app: Optional[str] = None,
_automatic_io_annotation: bool = True,
):
if kind in MODEL_DECORATORS:
decorator = _model_decorator(kind)
return decorator(
original_func,
model_name,
model_provider,
name,
session_id,
ml_app,
)
else:
decorator = _llmobs_decorator(kind or operation_kind)
return decorator(
original_func,
name,
session_id,
ml_app,
_automatic_io_annotation,
)

return decorator


MODEL_DECORATORS = ["llm", "embedding"]

llm = _model_decorator("llm")
embedding = _model_decorator("embedding")
Expand All @@ -261,3 +296,4 @@ def wrapper(*args, **kwargs):
tool = _llmobs_decorator("tool")
retrieval = _llmobs_decorator("retrieval")
agent = _llmobs_decorator("agent")
observe = _generic_decorator("undefined_kind")

0 comments on commit c7fa2f0

Please sign in to comment.