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/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 "reranking":
+      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 "reranking":
+        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",
+  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