def with_logfire(cls):
"""Wraps a pydantic class with a Logfire span."""
if hasattr(cls, "call"):
setattr(cls, "call", mirascope_logfire_span(cls.call))
if hasattr(cls, "stream"):
setattr(cls, "stream", mirascope_logfire_span(cls.stream))
if hasattr(cls, "call_async"):
setattr(cls, "call_async", mirascope_logfire_span(cls.call_async))
if hasattr(cls, "stream_async"):
setattr(cls, "stream_async", mirascope_logfire_span(cls.stream_async))
if hasattr(cls, "extract"):
setattr(cls, "extract", mirascope_logfire_span(cls.extract))
if hasattr(cls, "extract_async"):
setattr(cls, "extract_async", mirascope_logfire_span(cls.extract_async))
if hasattr(cls, "retrieve"):
setattr(cls, "retrieve", mirascope_logfire_span(cls.retrieve))
if hasattr(cls, "add"):
setattr(cls, "add", mirascope_logfire_span(cls.add))
if get_parent_class_name(cls, "OpenAI"):
if hasattr(cls, "call_params"):
cls.call_params.logfire = logfire.instrument_openai
if hasattr(cls, "embedding_params"):
cls.embedding_params.logfire = logfire.instrument_openai
else:
# TODO: Use instrument instead when they are integrated into logfire
if hasattr(cls, "call_params"):
cls.call_params.logfire = mirascope_logfire
cls.call_params.logfire_async = mirascope_logfire_async
if hasattr(cls, "vectorstore_params"):
# Wraps class methods rather than calls directly
cls.vectorstore_params.logfire = mirascope_logfire_span
cls.vectorstore_params.logfire_async = mirascope_logfire_span
if hasattr(cls, "embedding_params"):
cls.embedding_params.logfire = mirascope_logfire
cls.embedding_params.logfire_async = mirascope_logfire_async
return cls