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

app/schema.graphql

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

app/src/components/trace/SpanKindIcon.tsx

@@ -164,6 +164,27 @@ const RetrieverSVG = () => (
   </svg>
 );
 
+const RerankerSVG = () => (
+  <svg
+    width="20"
+    height="20"
+    viewBox="0 0 20 20"
+    fill="none"
+    xmlns="http://www.w3.org/2000/svg"
+  >
+    <rect
+      x="0.5"
+      y="0.5"
+      width="19"
+      height="19"
+      rx="3.5"
+      stroke="currentColor"
+    />
+    <path d="M4.5359 10L8 4L11.4641 10H4.5359Z" stroke="currentColor" />
+    <path d="M8.5359 10L12 16L15.4641 10H8.5359Z" stroke="currentColor" />
+  </svg>
+);
+
 const ChainSVG = () => (
   <svg
     width="20"
@@ -219,6 +240,10 @@ export function SpanKindIcon({ spanKind }: { spanKind: string }) {
       color = "--ac-global-color-yellow-1200";
       icon = <ToolSVG />;
       break;
+    case "reranker":
+      color = "--ac-global-color-celery-1000";
+      icon = <RerankerSVG />;
+      break;
   }
 
   return (

app/src/components/trace/SpanKindLabel.tsx

@@ -16,6 +16,9 @@ export function SpanKindLabel(props: { spanKind: string }) {
       case "retriever":
         color = "seafoam-1000";
         break;
+      case "reranker":
+        color = "celery-1000";
+        break;
       case "embedding":
         color = "indigo-1000";
         break;

app/src/openInference/tracing/semanticConventions.ts

@@ -1,6 +1,7 @@
 export const SemanticAttributePrefixes = {
   llm: "llm",
   retrieval: "retrieval",
+  reranker: "reranker",
   messages: "messages",
   message: "message",
   document: "document",
@@ -26,6 +27,14 @@ export const RetrievalAttributePostfixes = {
   documents: "documents",
 } as const;
 
+export const RerankerAttributePostfixes = {
+  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,
+  RerankerAttributePostfixes,
   RetrievalAttributePostfixes,
   SemanticAttributePrefixes,
   ToolAttributePostfixes,
@@ -324,6 +325,12 @@ function SpanInfo({ span }: { span: Span }) {
       );
       break;
     }
+    case "reranker": {
+      content = (
+        <RerankerSpanInfo span={span} spanAttributes={attributesObject} />
+      );
+      break;
+    }
     case "embedding": {
       content = (
         <EmbeddingSpanInfo span={span} spanAttributes={attributesObject} />
@@ -572,6 +579,100 @@ function RetrieverSpanInfo(props: {
   );
 }
 
+function RerankerSpanInfo(props: {
+  span: Span;
+  spanAttributes: AttributeObject;
+}) {
+  const { spanAttributes } = props;
+  const rerankerAttributes = useMemo<AttributeObject | null>(() => {
+    const rerankerAttrs = spanAttributes[SemanticAttributePrefixes.reranker];
+    if (typeof rerankerAttrs === "object") {
+      return rerankerAttrs as AttributeObject;
+    }
+    return null;
+  }, [spanAttributes]);
+  const query = useMemo<string>(() => {
+    if (rerankerAttributes == null) {
+      return "";
+    }
+    return (rerankerAttributes[RerankerAttributePostfixes.query] ||
+      "") as string;
+  }, [rerankerAttributes]);
+  const input_documents = useMemo<AttributeDocument[]>(() => {
+    if (rerankerAttributes == null) {
+      return [];
+    }
+    return (rerankerAttributes[RerankerAttributePostfixes.input_documents] ||
+      []) as AttributeDocument[];
+  }, [rerankerAttributes]);
+  const output_documents = useMemo<AttributeDocument[]>(() => {
+    if (rerankerAttributes == null) {
+      return [];
+    }
+    return (rerankerAttributes[RerankerAttributePostfixes.output_documents] ||
+      []) as AttributeDocument[];
+  }, [rerankerAttributes]);
+
+  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"}
+        titleExtra={<Counter variant="light">{numInputDocuments}</Counter>}
+        {...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={"Output Documents"}
+        titleExtra={<Counter variant="light">{numOutputDocuments}</Counter>}
+        {...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,11 +2,16 @@
 
 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
@@ -33,6 +38,10 @@
     MESSAGE_ROLE,
     OUTPUT_MIME_TYPE,
     OUTPUT_VALUE,
+    RERANKER_INPUT_DOCUMENTS,
+    RERANKER_MODEL_NAME,
+    RERANKER_OUTPUT_DOCUMENTS,
+    RERANKER_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=NoOpExporter())
+    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
+    reranker_span = spans["reranking"]
+    assert reranker_span.span_kind == SpanKind.RERANKER
+    assert (
+        len(reranker_span.attributes[RERANKER_INPUT_DOCUMENTS])
+        == query_engine.retriever.similarity_top_k
+    )
+    assert len(reranker_span.attributes[RERANKER_OUTPUT_DOCUMENTS]) == cohere_rerank.top_n
+    assert reranker_span.attributes[RERANKER_TOP_K] == cohere_rerank.top_n
+    assert reranker_span.attributes[RERANKER_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
+    reranker = trace_schema.SpanKind.RERANKER
     unknown = trace_schema.SpanKind.UNKNOWN
 
     @classmethod

src/phoenix/trace/llama_index/callback.py

@@ -51,6 +51,11 @@
     MESSAGE_ROLE,
     OUTPUT_MIME_TYPE,
     OUTPUT_VALUE,
+    RERANKER_INPUT_DOCUMENTS,
+    RERANKER_MODEL_NAME,
+    RERANKER_OUTPUT_DOCUMENTS,
+    RERANKER_QUERY,
+    RERANKER_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,22 @@ 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[RERANKER_TOP_K] = payload[EventPayload.TOP_K]
+        if EventPayload.MODEL_NAME in payload:
+            attributes[RERANKER_MODEL_NAME] = payload[EventPayload.MODEL_NAME]
+        if EventPayload.QUERY_STR in payload:
+            attributes[RERANKER_QUERY] = payload[EventPayload.QUERY_STR]
+        if nodes := payload.get(EventPayload.NODES):
+            attributes[RERANKER_OUTPUT_DOCUMENTS if is_event_end else RERANKER_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 +238,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 +368,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.RERANKER,
     }.get(event_type, SpanKind.CHAIN)