169 lines
4.7 KiB
Python
169 lines
4.7 KiB
Python
import uuid
|
|
from typing import Any, Dict, Optional
|
|
|
|
from llama_index.core.instrumentation.span import SimpleSpan
|
|
from llama_index.core.instrumentation.span_handlers.base import BaseSpanHandler
|
|
from llama_index.core.query_engine import RetrieverQueryEngine
|
|
from typing_extensions import TypedDict
|
|
|
|
from literalai.context import active_thread_var
|
|
|
|
literalai_uuid_namespace = uuid.UUID("05f6b2b5-a912-47bd-958f-98a9c4496322")
|
|
|
|
|
|
class SpanEntry(TypedDict):
|
|
id: str
|
|
parent_id: Optional[str]
|
|
root_id: Optional[str]
|
|
is_run_root: bool
|
|
|
|
|
|
class LiteralSpanHandler(BaseSpanHandler[SimpleSpan]):
|
|
"""This class handles spans coming from LlamaIndex."""
|
|
|
|
spans: Dict[str, SpanEntry] = {}
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def new_span(
|
|
self,
|
|
id_: str,
|
|
bound_args: Any,
|
|
instance: Optional[Any] = None,
|
|
parent_span_id: Optional[str] = None,
|
|
tags: Optional[Dict[str, Any]] = None,
|
|
**kwargs: Any,
|
|
):
|
|
self.spans[id_] = {
|
|
"id": id_,
|
|
"parent_id": parent_span_id,
|
|
"root_id": None,
|
|
"is_run_root": self.is_run_root(instance, parent_span_id),
|
|
}
|
|
|
|
if parent_span_id is not None:
|
|
self.spans[id_]["root_id"] = self.get_root_span_id(parent_span_id)
|
|
else:
|
|
self.spans[id_]["root_id"] = id_
|
|
|
|
def prepare_to_exit_span(
|
|
self,
|
|
id_: str,
|
|
bound_args: Any,
|
|
instance: Optional[Any] = None,
|
|
result: Optional[Any] = None,
|
|
**kwargs: Any,
|
|
):
|
|
"""Logic for preparing to exit a span."""
|
|
if id_ in self.spans:
|
|
del self.spans[id_]
|
|
|
|
def prepare_to_drop_span(
|
|
self,
|
|
id_: str,
|
|
bound_args: Any,
|
|
instance: Optional[Any] = None,
|
|
err: Optional[BaseException] = None,
|
|
**kwargs: Any,
|
|
):
|
|
"""Logic for preparing to drop a span."""
|
|
if id_ in self.spans:
|
|
del self.spans[id_]
|
|
|
|
def is_run_root(
|
|
self, instance: Optional[Any], parent_span_id: Optional[str]
|
|
) -> bool:
|
|
"""Returns True if the span is of type RetrieverQueryEngine, and it has no run root in its parent chain"""
|
|
if not isinstance(instance, RetrieverQueryEngine):
|
|
return False
|
|
|
|
# Span is of correct type, we check that it doesn't have a run root in its parent chain
|
|
while parent_span_id:
|
|
parent_span = self.spans.get(parent_span_id)
|
|
|
|
if not parent_span:
|
|
parent_span_id = None
|
|
continue
|
|
|
|
if parent_span["is_run_root"]:
|
|
return False
|
|
|
|
parent_span_id = parent_span["parent_id"]
|
|
|
|
return True
|
|
|
|
def get_root_span_id(self, span_id: Optional[str]):
|
|
"""Finds the root span and returns its ID"""
|
|
if not span_id:
|
|
return None
|
|
|
|
current_span = self.spans.get(span_id)
|
|
|
|
if current_span is None:
|
|
return None
|
|
|
|
while current_span["parent_id"] is not None:
|
|
current_span = self.spans.get(current_span["parent_id"])
|
|
if current_span is None:
|
|
return None
|
|
|
|
return current_span["id"]
|
|
|
|
def get_run_id(self, span_id: Optional[str]):
|
|
"""Go up the span chain to find a run_root, return its ID (or None)"""
|
|
if not span_id:
|
|
return None
|
|
|
|
current_span = self.spans.get(span_id)
|
|
|
|
if current_span is None:
|
|
return None
|
|
|
|
while current_span:
|
|
if current_span["is_run_root"]:
|
|
return str(uuid.uuid5(literalai_uuid_namespace, current_span["id"]))
|
|
|
|
parent_id = current_span["parent_id"]
|
|
|
|
if parent_id:
|
|
current_span = self.spans.get(parent_id)
|
|
else:
|
|
current_span = None
|
|
|
|
return None
|
|
|
|
def get_thread_id(self, span_id: Optional[str]):
|
|
"""Returns the root span ID as a thread ID"""
|
|
active_thread = active_thread_var.get()
|
|
|
|
if active_thread:
|
|
return active_thread.id
|
|
|
|
if span_id is None:
|
|
return None
|
|
|
|
current_span = self.spans.get(span_id)
|
|
|
|
if current_span is None:
|
|
return None
|
|
|
|
root_id = current_span["root_id"]
|
|
|
|
if not root_id:
|
|
return None
|
|
|
|
root_span = self.spans.get(root_id)
|
|
|
|
if root_span is None:
|
|
# span is already the root, uuid its own id
|
|
return str(uuid.uuid5(literalai_uuid_namespace, span_id))
|
|
else:
|
|
# uuid the id of the root span
|
|
return str(uuid.uuid5(literalai_uuid_namespace, root_span["id"]))
|
|
|
|
@classmethod
|
|
def class_name(cls) -> str:
|
|
"""Class name."""
|
|
return "LiteralSpanHandler"
|