feat(traces): add reranking span kind for document reranking in llama index

app/schema.graphql

@@ -576,6 +576,7 @@ enum SpanKind {
   retriever
   embedding
   agent
+  reranking
   unknown
 }
 

app/src/openInference/tracing/semanticConventions.ts

@@ -1,6 +1,7 @@
 export const SemanticAttributePrefixes = {
   llm: "llm",
   retrieval: "retrieval",
+  reranking: "reranking",
   messages: "messages",
   message: "message",
   document: "document",
@@ -26,6 +27,14 @@ export const RetrievalAttributePostfixes = {
   documents: "documents",
 } as const;
 
+export const RerankingAttributePostfixes = {
+  input_documents: "input_documents",
+  output_documents: "output_documents",
+  query: "query",
+  model_name: "model_name",
+  top_k: "top_k",
+} as const;
+
 export const EmbeddingAttributePostfixes = {
   embeddings: "embeddings",
   text: "text",

app/src/pages/trace/TracePage.tsx

@@ -52,6 +52,7 @@ import {
   MESSAGE_FUNCTION_CALL_NAME,
   MESSAGE_NAME,
   MESSAGE_ROLE,
+  RerankingAttributePostfixes,
   RetrievalAttributePostfixes,
   SemanticAttributePrefixes,
   ToolAttributePostfixes,
@@ -324,6 +325,12 @@ function SpanInfo({ span }: { span: Span }) {
       );
       break;
     }
+    case "reranking": {
+      content = (
+        <RerankingSpanInfo span={span} spanAttributes={attributesObject} />
+      );
+      break;
+    }
     case "embedding": {
       content = (
         <EmbeddingSpanInfo span={span} spanAttributes={attributesObject} />
@@ -572,6 +579,98 @@ function RetrieverSpanInfo(props: {
   );
 }
 
+function RerankingSpanInfo(props: {
+  span: Span;
+  spanAttributes: AttributeObject;
+}) {
+  const { spanAttributes } = props;
+  const rerankingAttributes = useMemo<AttributeObject | null>(() => {
+    const rerankingAttrs = spanAttributes[SemanticAttributePrefixes.reranking];
+    if (typeof rerankingAttrs === "object") {
+      return rerankingAttrs as AttributeObject;
+    }
+    return null;
+  }, [spanAttributes]);
+  const query = useMemo<string>(() => {
+    if (rerankingAttributes == null) {
+      return "";
+    }
+    return (rerankingAttributes[RerankingAttributePostfixes.query] ||
+      "") as string;
+  }, [rerankingAttributes]);
+  const input_documents = useMemo<AttributeDocument[]>(() => {
+    if (rerankingAttributes == null) {
+      return [];
+    }
+    return (rerankingAttributes[RerankingAttributePostfixes.input_documents] ||
+      []) as AttributeDocument[];
+  }, [rerankingAttributes]);
+  const output_documents = useMemo<AttributeDocument[]>(() => {
+    if (rerankingAttributes == null) {
+      return [];
+    }
+    return (rerankingAttributes[RerankingAttributePostfixes.output_documents] ||
+      []) as AttributeDocument[];
+  }, [rerankingAttributes]);
+
+  const numInputDocuments = input_documents.length;
+  const numOutputDocuments = output_documents.length;
+  return (
+    <Flex direction="column" gap="size-200">
+      <Card title="Query" {...defaultCardProps}>
+        <CodeBlock value={query} mimeType="text" />
+      </Card>
+      <Card
+        title={`Input Documents (${numInputDocuments})`}
+        {...defaultCardProps}
+        defaultOpen={false}
+      >
+        {
+          <ul
+            css={css`
+              padding: var(--ac-global-dimension-static-size-200);
+              display: flex;
+              flex-direction: column;
+              gap: var(--ac-global-dimension-static-size-200);
+            `}
+          >
+            {input_documents.map((document, idx) => {
+              return (
+                <li key={idx}>
+                  <DocumentItem document={document} />
+                </li>
+              );
+            })}
+          </ul>
+        }
+      </Card>
+      <Card
+        title={`Re-ranked Documents (${numOutputDocuments})`}
+        {...defaultCardProps}
+      >
+        {
+          <ul
+            css={css`
+              padding: var(--ac-global-dimension-static-size-200);
+              display: flex;
+              flex-direction: column;
+              gap: var(--ac-global-dimension-static-size-200);
+            `}
+          >
+            {output_documents.map((document, idx) => {
+              return (
+                <li key={idx}>
+                  <DocumentItem document={document} />
+                </li>
+              );
+            })}
+          </ul>
+        }
+      </Card>
+    </Flex>
+  );
+}
+
 function EmbeddingSpanInfo(props: {
   span: Span;
   spanAttributes: AttributeObject;

integration-tests/trace/llama_index/test_callback.py

@@ -2,16 +2,21 @@
 
 import pytest
 from gcsfs import GCSFileSystem
-from llama_index import ServiceContext, StorageContext, load_index_from_storage
+from llama_index import (
+    ServiceContext,
+    StorageContext,
+    load_index_from_storage,
+)
 from llama_index.agent import OpenAIAgent
 from llama_index.callbacks import CallbackManager
 from llama_index.embeddings.openai import OpenAIEmbedding
 from llama_index.graph_stores.simple import SimpleGraphStore
+from llama_index.indices.postprocessor.cohere_rerank import CohereRerank
 from llama_index.indices.vector_store import VectorStoreIndex
 from llama_index.llms import OpenAI
 from llama_index.query_engine import RetrieverQueryEngine
 from llama_index.tools import FunctionTool
-from phoenix.trace.exporter import NoOpExporter
+from phoenix.trace.exporter import HttpExporter, NoOpExporter
 from phoenix.trace.llama_index import OpenInferenceTraceCallbackHandler
 from phoenix.trace.schemas import SpanKind
 from phoenix.trace.semantic_conventions import (
@@ -33,6 +38,10 @@
     MESSAGE_ROLE,
     OUTPUT_MIME_TYPE,
     OUTPUT_VALUE,
+    RERANKING_INPUT_DOCUMENTS,
+    RERANKING_MODEL_NAME,
+    RERANKING_OUTPUT_DOCUMENTS,
+    RERANKING_TOP_K,
     TOOL_DESCRIPTION,
     TOOL_NAME,
     TOOL_PARAMETERS,
@@ -241,3 +250,30 @@ def add(a: int, b: int) -> int:
         "title": "multiply",
         "type": "object",
     }
+
+
+@pytest.mark.parametrize("model_name", ["text-davinci-003"], indirect=True)
+def test_cohere_rerank(index: VectorStoreIndex) -> None:
+    callback_handler = OpenInferenceTraceCallbackHandler(exporter=HttpExporter())
+    service_context = ServiceContext.from_defaults(
+        callback_manager=CallbackManager(handlers=[callback_handler])
+    )
+    cohere_rerank = CohereRerank(top_n=2)
+    query_engine = index.as_query_engine(
+        similarity_top_k=5,
+        node_postprocessors=[cohere_rerank],
+        service_context=service_context,
+    )
+    query_engine.query("How should timestamps be formatted?")
+
+    spans = {span.name: span for span in callback_handler.get_spans()}
+    assert "reranking" in spans
+    reranking_span = spans["reranking"]
+    assert reranking_span.span_kind == SpanKind.RERANKING
+    assert (
+        len(reranking_span.attributes[RERANKING_INPUT_DOCUMENTS])
+        == query_engine.retriever.similarity_top_k
+    )
+    assert len(reranking_span.attributes[RERANKING_OUTPUT_DOCUMENTS]) == cohere_rerank.top_n
+    assert reranking_span.attributes[RERANKING_TOP_K] == cohere_rerank.top_n
+    assert reranking_span.attributes[RERANKING_MODEL_NAME] == cohere_rerank.model

src/phoenix/server/api/types/Span.py

@@ -46,6 +46,7 @@ class SpanKind(Enum):
     retriever = trace_schema.SpanKind.RETRIEVER
     embedding = trace_schema.SpanKind.EMBEDDING
     agent = trace_schema.SpanKind.AGENT
+    reranking = trace_schema.SpanKind.RERANKING
     unknown = trace_schema.SpanKind.UNKNOWN
 
     @classmethod

src/phoenix/trace/llama_index/callback.py

@@ -51,6 +51,11 @@
     MESSAGE_ROLE,
     OUTPUT_MIME_TYPE,
     OUTPUT_VALUE,
+    RERANKING_INPUT_DOCUMENTS,
+    RERANKING_MODEL_NAME,
+    RERANKING_OUTPUT_DOCUMENTS,
+    RERANKING_QUERY,
+    RERANKING_TOP_K,
     RETRIEVAL_DOCUMENTS,
     TOOL_DESCRIPTION,
     TOOL_NAME,
@@ -82,6 +87,7 @@ class CBEventData(TypedDict, total=False):
 def payload_to_semantic_attributes(
     event_type: CBEventType,
     payload: Dict[str, Any],
+    is_event_end: bool = False,
 ) -> Dict[str, Any]:
     """
     Converts a LLMapp payload to a dictionary of semantic conventions compliant attributes.
@@ -95,10 +101,10 @@ def payload_to_semantic_attributes(
             {EMBEDDING_TEXT: text, EMBEDDING_VECTOR: vector}
             for text, vector in zip(payload[EventPayload.CHUNKS], payload[EventPayload.EMBEDDINGS])
         ]
-    if EventPayload.QUERY_STR in payload:
+    if event_type is not CBEventType.RERANKING and EventPayload.QUERY_STR in payload:
         attributes[INPUT_VALUE] = payload[EventPayload.QUERY_STR]
         attributes[INPUT_MIME_TYPE] = MimeType.TEXT
-    if EventPayload.NODES in payload:
+    if event_type is not CBEventType.RERANKING and EventPayload.NODES in payload:
         attributes[RETRIEVAL_DOCUMENTS] = [
             {
                 DOCUMENT_ID: node_with_score.node.node_id,
@@ -128,11 +134,24 @@ def payload_to_semantic_attributes(
             if (usage := getattr(raw, "usage", None)) is not None:
                 attributes.update(_get_token_counts(usage))
     if event_type is CBEventType.RERANKING:
-        ...  # TODO
-        # if EventPayload.TOP_K in payload:
-        #     attributes[RERANKING_TOP_K] = payload[EventPayload.TOP_K]
-        # if EventPayload.MODEL_NAME in payload:
-        #     attributes[RERANKING_MODEL_NAME] = payload[EventPayload.MODEL_NAME]
+        if EventPayload.TOP_K in payload:
+            attributes[RERANKING_TOP_K] = payload[EventPayload.TOP_K]
+        if EventPayload.MODEL_NAME in payload:
+            attributes[RERANKING_MODEL_NAME] = payload[EventPayload.MODEL_NAME]
+        if EventPayload.QUERY_STR in payload:
+            attributes[RERANKING_QUERY] = payload[EventPayload.QUERY_STR]
+        if nodes := payload.get(EventPayload.NODES):
+            attributes[
+                RERANKING_OUTPUT_DOCUMENTS if is_event_end else RERANKING_INPUT_DOCUMENTS
+            ] = [
+                {
+                    DOCUMENT_ID: node_with_score.node.node_id,
+                    DOCUMENT_SCORE: node_with_score.score,
+                    DOCUMENT_CONTENT: node_with_score.node.text,
+                    DOCUMENT_METADATA: node_with_score.node.metadata,
+                }
+                for node_with_score in nodes
+            ]
     if EventPayload.TOOL in payload:
         tool_metadata = cast(ToolMetadata, payload.get(EventPayload.TOOL))
         attributes[TOOL_NAME] = tool_metadata.name
@@ -221,7 +240,7 @@ def on_event_end(
         # Parse the payload to extract the parameters
         if payload is not None:
             event_data["attributes"].update(
-                payload_to_semantic_attributes(event_type, payload),
+                payload_to_semantic_attributes(event_type, payload, is_event_end=True),
             )
 
     def start_trace(self, trace_id: Optional[str] = None) -> None:
@@ -351,6 +370,7 @@ def _get_span_kind(event_type: CBEventType) -> SpanKind:
         CBEventType.RETRIEVE: SpanKind.RETRIEVER,
         CBEventType.FUNCTION_CALL: SpanKind.TOOL,
         CBEventType.AGENT_STEP: SpanKind.AGENT,
+        CBEventType.RERANKING: SpanKind.RERANKING,
     }.get(event_type, SpanKind.CHAIN)
 
 

src/phoenix/trace/schemas.py

@@ -41,6 +41,7 @@ class SpanKind(Enum):
     RETRIEVER = "RETRIEVER"
     EMBEDDING = "EMBEDDING"
     AGENT = "AGENT"
+    RERANKING = "RERANKING"
     UNKNOWN = "UNKNOWN"
 
     def __str__(self) -> str:

src/phoenix/trace/semantic_conventions.py

@@ -179,3 +179,24 @@ def _missing_(cls, v: Any) -> Optional["MimeType"]:
 """
 Document metadata as a string representing a JSON object
 """
+
+RERANKING_INPUT_DOCUMENTS = "reranking.input_documents"
+"""
+List of documents as input to the reranker
+"""
+RERANKING_OUTPUT_DOCUMENTS = "reranking.output_documents"
+"""
+List of documents as output from the reranker
+"""
+RERANKING_QUERY = "reranking.query"
+"""
+Query string for the reranker
+"""
+RERANKING_MODEL_NAME = "reranking.model_name"
+"""
+Model name of the reranker
+"""
+RERANKING_TOP_K = "reranking.top_k"
+"""
+Top K parameter of the reranker
+"""

tests/server/api/types/test_span.py

@@ -60,6 +60,13 @@ def test_nested_attributes() -> None:
         "retrieval": {
             "documents": ...,
         },
+        "reranking": {
+            "input_documents": ...,
+            "output_documents": ...,
+            "model_name": ...,
+            "top_k": ...,
+            "query": ...,
+        },
         "tool": {
             "description": ...,
             "name": ...,