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

… index (#1588)

* add reranking span

* add icon and color

* remove console.log

* add titleExtra

* rename reranking to reranker

* clean up

* revert demo notebook

* Revert "revert demo notebook"

This reverts commit 58cf84c.

* bump deps

* add rerank fixture, color reranked documents differently

---------

Co-authored-by: Mikyo King <mikyo@arize.com>

app/package.json

@@ -6,7 +6,7 @@
   "license": "None",
   "private": true,
   "dependencies": {
-    "@arizeai/components": "^0.17.4",
+    "@arizeai/components": "^0.17.5",
     "@arizeai/point-cloud": "^3.0.4",
     "@codemirror/autocomplete": "^6.9.2",
     "@codemirror/lang-json": "^6.0.1",

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

@@ -23,13 +23,15 @@ import {
   Icon,
   Icons,
   Label,
+  LabelProps,
   List,
   ListItem,
   TabbedCard,
   TabPane,
   Tabs,
   Text,
   View,
+  ViewProps,
   ViewStyleProps,
 } from "@arizeai/components";
 
@@ -52,6 +54,7 @@ import {
   MESSAGE_FUNCTION_CALL_NAME,
   MESSAGE_NAME,
   MESSAGE_ROLE,
+  RerankerAttributePostfixes,
   RetrievalAttributePostfixes,
   SemanticAttributePrefixes,
   ToolAttributePostfixes,
@@ -324,6 +327,12 @@ function SpanInfo({ span }: { span: Span }) {
       );
       break;
     }
+    case "reranker": {
+      content = (
+        <RerankerSpanInfo span={span} spanAttributes={attributesObject} />
+      );
+      break;
+    }
     case "embedding": {
       content = (
         <EmbeddingSpanInfo span={span} spanAttributes={attributesObject} />
@@ -560,7 +569,12 @@ function RetrieverSpanInfo(props: {
               {documents.map((document, idx) => {
                 return (
                   <li key={idx}>
-                    <DocumentItem document={document} />
+                    <DocumentItem
+                      document={document}
+                      borderColor={"seafoam-700"}
+                      backgroundColor={"seafoam-100"}
+                      labelColor="seafoam-1000"
+                    />
                   </li>
                 );
               })}
@@ -572,6 +586,110 @@ 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}
+                    borderColor={"seafoam-700"}
+                    backgroundColor={"seafoam-100"}
+                    labelColor="seafoam-1000"
+                  />
+                </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}
+                    borderColor={"celery-700"}
+                    backgroundColor={"celery-100"}
+                    labelColor="celery-1000"
+                  />
+                </li>
+              );
+            })}
+          </ul>
+        }
+      </Card>
+    </Flex>
+  );
+}
+
 function EmbeddingSpanInfo(props: {
   span: Span;
   spanAttributes: AttributeObject;
@@ -714,20 +832,30 @@ function ToolSpanInfo(props: { span: Span; spanAttributes: AttributeObject }) {
   );
 }
 
-function DocumentItem({ document }: { document: AttributeDocument }) {
+function DocumentItem({
+  document,
+  backgroundColor,
+  borderColor,
+  labelColor,
+}: {
+  document: AttributeDocument;
+  backgroundColor: ViewProps["backgroundColor"];
+  borderColor: ViewProps["borderColor"];
+  labelColor: LabelProps["color"];
+}) {
   const metadata = document[DOCUMENT_METADATA];
   return (
     <View
       borderRadius="medium"
-      backgroundColor="seafoam-100"
-      borderColor="seafoam-700"
+      backgroundColor={backgroundColor}
+      borderColor={borderColor}
       borderWidth="thin"
     >
       <Flex direction="column">
         <View
           width="100%"
           borderBottomWidth="thin"
-          borderBottomColor="seafoam-700"
+          borderBottomColor={borderColor}
         >
           <Flex
             direction="row"
@@ -740,7 +868,7 @@ function DocumentItem({ document }: { document: AttributeDocument }) {
               <Heading level={4}>document {document[DOCUMENT_ID]}</Heading>
             </Flex>
             {typeof document[DOCUMENT_SCORE] === "number" && (
-              <Label color="seafoam-1000">{`score ${numberFormatter(
+              <Label color={labelColor}>{`score ${numberFormatter(
                 document[DOCUMENT_SCORE]
               )}`}</Label>
             )}
@@ -757,7 +885,7 @@ function DocumentItem({ document }: { document: AttributeDocument }) {
         </pre>
         {metadata && (
           <>
-            <View borderColor="seafoam-700" borderTopWidth="thin">
+            <View borderColor={borderColor} borderTopWidth="thin">
               <CodeBlock value={JSON.stringify(metadata)} mimeType="json" />
             </View>
           </>