411 lines
15 KiB
Python
411 lines
15 KiB
Python
|
import logging
|
||
|
|
import uuid
|
||
|
|
from typing import TYPE_CHECKING, Dict, List, Optional, Union, cast
|
||
|
|
|
||
|
|
from llama_index.core.base.llms.types import ChatMessage, MessageRole
|
||
|
|
from llama_index.core.base.response.schema import Response, StreamingResponse
|
||
|
|
from llama_index.core.instrumentation.event_handlers import BaseEventHandler
|
||
|
|
from llama_index.core.instrumentation.events import BaseEvent
|
||
|
|
from llama_index.core.instrumentation.events.agent import (
|
||
|
|
AgentChatWithStepEndEvent,
|
||
|
|
AgentChatWithStepStartEvent,
|
||
|
|
AgentRunStepEndEvent,
|
||
|
|
AgentRunStepStartEvent,
|
||
|
|
)
|
||
|
|
from llama_index.core.instrumentation.events.embedding import (
|
||
|
|
EmbeddingEndEvent,
|
||
|
|
EmbeddingStartEvent,
|
||
|
|
)
|
||
|
|
from llama_index.core.instrumentation.events.llm import (
|
||
|
|
LLMChatEndEvent,
|
||
|
|
LLMChatStartEvent,
|
||
|
|
)
|
||
|
|
from llama_index.core.instrumentation.events.query import QueryEndEvent, QueryStartEvent
|
||
|
|
from llama_index.core.instrumentation.events.retrieval import (
|
||
|
|
RetrievalEndEvent,
|
||
|
|
RetrievalStartEvent,
|
||
|
|
)
|
||
|
|
from llama_index.core.instrumentation.events.synthesis import SynthesizeEndEvent
|
||
|
|
from llama_index.core.schema import NodeWithScore, QueryBundle, TextNode
|
||
|
|
from openai.types import CompletionUsage
|
||
|
|
from openai.types.chat import ChatCompletion, ChatCompletionChunk
|
||
|
|
from pydantic import PrivateAttr
|
||
|
|
|
||
|
|
from literalai.context import active_thread_var
|
||
|
|
from literalai.instrumentation.llamaindex.span_handler import LiteralSpanHandler
|
||
|
|
from literalai.observability.generation import (
|
||
|
|
ChatGeneration,
|
||
|
|
GenerationMessage,
|
||
|
|
GenerationMessageRole,
|
||
|
|
)
|
||
|
|
from literalai.observability.step import Step, StepType
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from literalai.client import LiteralClient
|
||
|
|
|
||
|
|
|
||
|
|
def convert_message_role(role: MessageRole) -> GenerationMessageRole:
|
||
|
|
mapping = {
|
||
|
|
MessageRole.SYSTEM: "system",
|
||
|
|
MessageRole.USER: "user",
|
||
|
|
MessageRole.ASSISTANT: "assistant",
|
||
|
|
MessageRole.FUNCTION: "function",
|
||
|
|
MessageRole.TOOL: "tool",
|
||
|
|
MessageRole.CHATBOT: "assistant",
|
||
|
|
MessageRole.MODEL: "assistant",
|
||
|
|
}
|
||
|
|
|
||
|
|
return cast(GenerationMessageRole, mapping.get(role, "user"))
|
||
|
|
|
||
|
|
|
||
|
|
def extract_query_from_bundle(str_or_query_bundle: Union[str, QueryBundle]):
|
||
|
|
if isinstance(str_or_query_bundle, QueryBundle):
|
||
|
|
return str_or_query_bundle.query_str
|
||
|
|
|
||
|
|
return str_or_query_bundle
|
||
|
|
|
||
|
|
|
||
|
|
def extract_document_info(nodes: List[NodeWithScore]):
|
||
|
|
if not nodes:
|
||
|
|
return []
|
||
|
|
|
||
|
|
return [
|
||
|
|
{
|
||
|
|
"id_": node.id_,
|
||
|
|
"metadata": node.metadata,
|
||
|
|
"text": node.get_text(),
|
||
|
|
"mimetype": node.node.mimetype,
|
||
|
|
"start_char_idx": node.node.start_char_idx,
|
||
|
|
"end_char_idx": node.node.end_char_idx,
|
||
|
|
"char_length": (
|
||
|
|
node.node.end_char_idx - node.node.start_char_idx
|
||
|
|
if node.node.end_char_idx is not None
|
||
|
|
and node.node.start_char_idx is not None
|
||
|
|
else None
|
||
|
|
),
|
||
|
|
"score": node.get_score(),
|
||
|
|
}
|
||
|
|
for node in nodes
|
||
|
|
if isinstance(node.node, TextNode)
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def build_message_dict(message: ChatMessage):
|
||
|
|
message_dict: GenerationMessage = {
|
||
|
|
"role": convert_message_role(message.role),
|
||
|
|
"content": message.content,
|
||
|
|
}
|
||
|
|
|
||
|
|
kwargs = message.additional_kwargs
|
||
|
|
|
||
|
|
if kwargs:
|
||
|
|
if kwargs.get("tool_call_id", None):
|
||
|
|
message_dict["tool_call_id"] = kwargs.get("tool_call_id")
|
||
|
|
if kwargs.get("name", None):
|
||
|
|
message_dict["name"] = kwargs.get("name")
|
||
|
|
tool_calls = kwargs.get("tool_calls", [])
|
||
|
|
if len(tool_calls) > 0:
|
||
|
|
message_dict["tool_calls"] = [
|
||
|
|
tool_call.to_dict() for tool_call in tool_calls
|
||
|
|
]
|
||
|
|
return message_dict
|
||
|
|
|
||
|
|
|
||
|
|
def create_generation(event: LLMChatStartEvent):
|
||
|
|
model_dict = event.model_dict
|
||
|
|
|
||
|
|
return ChatGeneration(
|
||
|
|
provider=model_dict.get("class_name"),
|
||
|
|
model=model_dict.get("model"),
|
||
|
|
settings={
|
||
|
|
"model": model_dict.get("model"),
|
||
|
|
"temperature": model_dict.get("temperature"),
|
||
|
|
"max_tokens": model_dict.get("max_tokens"),
|
||
|
|
"logprobs": model_dict.get("logprobs"),
|
||
|
|
"top_logprobs": model_dict.get("top_logprobs"),
|
||
|
|
},
|
||
|
|
messages=[build_message_dict(message) for message in event.messages],
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def extract_query(x: Union[str, QueryBundle]):
|
||
|
|
return x.query_str if isinstance(x, QueryBundle) else x
|
||
|
|
|
||
|
|
|
||
|
|
class LiteralEventHandler(BaseEventHandler):
|
||
|
|
"""This class handles events coming from LlamaIndex."""
|
||
|
|
|
||
|
|
_client: "LiteralClient" = PrivateAttr()
|
||
|
|
_span_handler: "LiteralSpanHandler" = PrivateAttr()
|
||
|
|
runs: Dict[str, List[Step]] = {}
|
||
|
|
streaming_run_ids: List[str] = []
|
||
|
|
_standalone_step_id: Optional[str] = None
|
||
|
|
open_runs: List[Step] = []
|
||
|
|
|
||
|
|
class Config:
|
||
|
|
arbitrary_types_allowed = True
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
literal_client: "LiteralClient",
|
||
|
|
llama_index_span_handler: "LiteralSpanHandler",
|
||
|
|
):
|
||
|
|
super().__init__()
|
||
|
|
object.__setattr__(self, "_client", literal_client)
|
||
|
|
object.__setattr__(self, "_span_handler", llama_index_span_handler)
|
||
|
|
|
||
|
|
def _convert_message(self, message: ChatMessage):
|
||
|
|
tool_calls = message.additional_kwargs.get("tool_calls")
|
||
|
|
msg: GenerationMessage = {
|
||
|
|
"name": getattr(message, "name", None),
|
||
|
|
"role": convert_message_role(message.role),
|
||
|
|
"content": message.content,
|
||
|
|
"tool_calls": (
|
||
|
|
[tool_call.to_dict() for tool_call in tool_calls]
|
||
|
|
if tool_calls
|
||
|
|
else None
|
||
|
|
),
|
||
|
|
}
|
||
|
|
|
||
|
|
return msg
|
||
|
|
|
||
|
|
def handle(self, event: BaseEvent, **kwargs) -> None:
|
||
|
|
"""Logic for handling event."""
|
||
|
|
try:
|
||
|
|
thread_id = self._span_handler.get_thread_id(event.span_id)
|
||
|
|
run_id = self._span_handler.get_run_id(event.span_id)
|
||
|
|
|
||
|
|
# AgentChatWithStep wraps several AgentRunStep events
|
||
|
|
# as the agent may want to perform multiple tool calls in a row.
|
||
|
|
if isinstance(event, AgentChatWithStepStartEvent) or isinstance(
|
||
|
|
event, AgentRunStepStartEvent
|
||
|
|
):
|
||
|
|
run_name = (
|
||
|
|
"Agent Chat"
|
||
|
|
if isinstance(event, AgentChatWithStepStartEvent)
|
||
|
|
else "Agent Step"
|
||
|
|
)
|
||
|
|
parent_run_id = None
|
||
|
|
if len(self.open_runs) > 0:
|
||
|
|
parent_run_id = self.open_runs[-1].id
|
||
|
|
|
||
|
|
agent_run_id = str(uuid.uuid4())
|
||
|
|
|
||
|
|
run = self._client.start_step(
|
||
|
|
name=run_name,
|
||
|
|
type="run",
|
||
|
|
id=agent_run_id,
|
||
|
|
parent_id=parent_run_id,
|
||
|
|
)
|
||
|
|
|
||
|
|
self.open_runs.append(run)
|
||
|
|
|
||
|
|
if isinstance(event, AgentChatWithStepEndEvent) or isinstance(
|
||
|
|
event, AgentRunStepEndEvent
|
||
|
|
):
|
||
|
|
try:
|
||
|
|
step = self.open_runs.pop()
|
||
|
|
except IndexError:
|
||
|
|
logging.error(
|
||
|
|
"[Literal] Error in Llamaindex instrumentation: AgentRunStepEndEvent called without an open run."
|
||
|
|
)
|
||
|
|
if step:
|
||
|
|
step.end()
|
||
|
|
|
||
|
|
if isinstance(event, QueryStartEvent):
|
||
|
|
active_thread = active_thread_var.get()
|
||
|
|
query = extract_query(event.query)
|
||
|
|
|
||
|
|
if not active_thread or not active_thread.name:
|
||
|
|
self._client.api.upsert_thread(id=thread_id, name=query)
|
||
|
|
|
||
|
|
self._client.message(
|
||
|
|
name="User query",
|
||
|
|
id=str(event.id_),
|
||
|
|
type="user_message",
|
||
|
|
thread_id=thread_id,
|
||
|
|
content=query,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Retrieval wraps the Embedding step in LlamaIndex
|
||
|
|
if isinstance(event, RetrievalStartEvent):
|
||
|
|
run = self._client.start_step(
|
||
|
|
name="RAG",
|
||
|
|
type="run",
|
||
|
|
id=run_id,
|
||
|
|
thread_id=thread_id,
|
||
|
|
)
|
||
|
|
|
||
|
|
self.store_step(run_id=run_id, step=run)
|
||
|
|
|
||
|
|
retrieval_step = self._client.start_step(
|
||
|
|
type="retrieval",
|
||
|
|
name="Retrieval",
|
||
|
|
parent_id=run_id,
|
||
|
|
thread_id=thread_id,
|
||
|
|
)
|
||
|
|
|
||
|
|
self.store_step(run_id=run_id, step=retrieval_step)
|
||
|
|
|
||
|
|
if isinstance(event, EmbeddingStartEvent):
|
||
|
|
retrieval_step = self.get_first_step_of_type(
|
||
|
|
run_id=run_id, step_type="retrieval"
|
||
|
|
)
|
||
|
|
|
||
|
|
if run_id and retrieval_step:
|
||
|
|
embedding_step = self._client.start_step(
|
||
|
|
type="embedding",
|
||
|
|
name="Embedding",
|
||
|
|
parent_id=retrieval_step.id,
|
||
|
|
thread_id=thread_id,
|
||
|
|
)
|
||
|
|
embedding_step.metadata = event.model_dict
|
||
|
|
self.store_step(run_id=run_id, step=embedding_step)
|
||
|
|
|
||
|
|
if isinstance(event, EmbeddingEndEvent):
|
||
|
|
embedding_step = self.get_first_step_of_type(
|
||
|
|
run_id=run_id, step_type="embedding"
|
||
|
|
)
|
||
|
|
|
||
|
|
if run_id and embedding_step:
|
||
|
|
embedding_step.input = {"query": event.chunks}
|
||
|
|
embedding_step.output = {"embeddings": event.embeddings}
|
||
|
|
embedding_step.end()
|
||
|
|
|
||
|
|
if isinstance(event, RetrievalEndEvent):
|
||
|
|
retrieval_step = self.get_first_step_of_type(
|
||
|
|
run_id=run_id, step_type="retrieval"
|
||
|
|
)
|
||
|
|
|
||
|
|
if run_id and retrieval_step:
|
||
|
|
retrieved_documents = extract_document_info(event.nodes)
|
||
|
|
query = extract_query(event.str_or_query_bundle)
|
||
|
|
|
||
|
|
retrieval_step.input = {"query": query}
|
||
|
|
retrieval_step.output = {"retrieved_documents": retrieved_documents}
|
||
|
|
retrieval_step.end()
|
||
|
|
|
||
|
|
# Only event where we create LLM steps
|
||
|
|
if isinstance(event, LLMChatStartEvent):
|
||
|
|
if run_id:
|
||
|
|
self._client.step()
|
||
|
|
llm_step = self._client.start_step(
|
||
|
|
type="llm",
|
||
|
|
parent_id=run_id,
|
||
|
|
thread_id=thread_id,
|
||
|
|
)
|
||
|
|
self.store_step(run_id=run_id, step=llm_step)
|
||
|
|
llm_step = self.get_first_step_of_type(run_id=run_id, step_type="llm")
|
||
|
|
|
||
|
|
if not run_id and not llm_step:
|
||
|
|
self._standalone_step_id = str(uuid.uuid4())
|
||
|
|
llm_step = self._client.start_step(
|
||
|
|
name=event.model_dict.get("model", "LLM"),
|
||
|
|
type="llm",
|
||
|
|
id=self._standalone_step_id,
|
||
|
|
# Remove thread_id for standalone runs
|
||
|
|
)
|
||
|
|
self.store_step(run_id=self._standalone_step_id, step=llm_step)
|
||
|
|
|
||
|
|
if llm_step:
|
||
|
|
generation = create_generation(event=event)
|
||
|
|
llm_step.generation = generation
|
||
|
|
llm_step.name = event.model_dict.get("model")
|
||
|
|
|
||
|
|
# Actual creation of the event happens upon ending the event
|
||
|
|
if isinstance(event, LLMChatEndEvent):
|
||
|
|
llm_step = self.get_first_step_of_type(run_id=run_id, step_type="llm")
|
||
|
|
if not llm_step and self._standalone_step_id:
|
||
|
|
llm_step = self.get_first_step_of_type(
|
||
|
|
run_id=self._standalone_step_id, step_type="llm"
|
||
|
|
)
|
||
|
|
|
||
|
|
response = event.response
|
||
|
|
|
||
|
|
if llm_step and response:
|
||
|
|
chat_completion = response.raw
|
||
|
|
|
||
|
|
# ChatCompletionChunk needed for chat stream methods
|
||
|
|
if isinstance(chat_completion, ChatCompletion) or isinstance(
|
||
|
|
chat_completion, ChatCompletionChunk
|
||
|
|
):
|
||
|
|
usage = chat_completion.usage
|
||
|
|
|
||
|
|
if isinstance(usage, CompletionUsage):
|
||
|
|
llm_step.generation.input_token_count = usage.prompt_tokens
|
||
|
|
llm_step.generation.output_token_count = (
|
||
|
|
usage.completion_tokens
|
||
|
|
)
|
||
|
|
|
||
|
|
if self._standalone_step_id:
|
||
|
|
llm_step.generation.message_completion = (
|
||
|
|
self._convert_message(response.message)
|
||
|
|
)
|
||
|
|
|
||
|
|
llm_step.end()
|
||
|
|
self._standalone_step_id = None
|
||
|
|
|
||
|
|
if isinstance(event, SynthesizeEndEvent):
|
||
|
|
llm_step = self.get_first_step_of_type(run_id=run_id, step_type="llm")
|
||
|
|
run = self.get_first_step_of_type(run_id=run_id, step_type="run")
|
||
|
|
|
||
|
|
if llm_step and run:
|
||
|
|
synthesized_response = event.response
|
||
|
|
text_response = ""
|
||
|
|
|
||
|
|
if isinstance(synthesized_response, StreamingResponse):
|
||
|
|
text_response = str(synthesized_response.get_response())
|
||
|
|
if isinstance(synthesized_response, Response):
|
||
|
|
text_response = str(synthesized_response)
|
||
|
|
|
||
|
|
llm_step.generation.message_completion = {
|
||
|
|
"role": "assistant",
|
||
|
|
"content": text_response,
|
||
|
|
}
|
||
|
|
|
||
|
|
llm_step.end()
|
||
|
|
run.end()
|
||
|
|
|
||
|
|
self._client.message(
|
||
|
|
type="assistant_message",
|
||
|
|
thread_id=thread_id,
|
||
|
|
content=text_response,
|
||
|
|
)
|
||
|
|
|
||
|
|
if isinstance(event, QueryEndEvent):
|
||
|
|
if run_id in self.runs:
|
||
|
|
del self.runs[run_id]
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logging.error(
|
||
|
|
"[Literal] Error in Llamaindex instrumentation : %s",
|
||
|
|
str(e),
|
||
|
|
exc_info=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
def store_step(self, run_id: str, step: Step):
|
||
|
|
if run_id not in self.runs:
|
||
|
|
self.runs[run_id] = []
|
||
|
|
|
||
|
|
self.runs[run_id].append(step)
|
||
|
|
|
||
|
|
def get_first_step_of_type(
|
||
|
|
self, run_id: Optional[str], step_type: StepType
|
||
|
|
) -> Optional[Step]:
|
||
|
|
if not run_id:
|
||
|
|
return None
|
||
|
|
|
||
|
|
if run_id not in self.runs:
|
||
|
|
return None
|
||
|
|
|
||
|
|
for step in self.runs[run_id]:
|
||
|
|
if step.type == step_type:
|
||
|
|
return step
|
||
|
|
|
||
|
|
return None
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def class_name(cls) -> str:
|
||
|
|
"""Class name."""
|
||
|
|
return "LiteralEventHandler"
|