Skip to content

Commit

Permalink
feat(observe): add raw prompt for streamlit agent (#161)
Browse files Browse the repository at this point in the history
Signed-off-by: Milan Gallas <[email protected]>
Co-authored-by: Milan Gallas <[email protected]>
  • Loading branch information
GALLLASMILAN and Milan Gallas authored Dec 20, 2024
1 parent 2b25a66 commit 586e66f
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 31 deletions.
54 changes: 23 additions & 31 deletions src/modules/chat/trace/useBuildTraceData.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import {
getLastNewTokenSpan,
getGeneratedTokenCountSafe,
getExecutionTime,
getRawPrompt,
} from './utils';
import { useEffect, useMemo, useState } from 'react';
import { GENERATE_EVENT_TOOL_START, TraceData } from './types';
Expand Down Expand Up @@ -79,35 +80,24 @@ export function useBuildTraceData({ enabled, threadId, runId }: Props): {
if (!spans) return null;
const iterations = spans
.filter((span) => span.name.includes('iteration-'))
.map((span) => {
return {
span,
children: loadNestedSpans(span, spans).flat(),
};
})
.map((span) => ({
span,
children: loadNestedSpans(span, spans).flat(),
}))
.filter((span) => span.children.length >= 1)
.map((iteration) => {
const customEventData = iteration.children.find(
(span) => span.attributes.name === 'startCustom',
)?.attributes.data;
const rawPrompt = customEventData?.rawPrompt
? String(customEventData.rawPrompt)
: undefined;

return {
groupId: iteration.span.name,
executionTime: getExecutionTime(iteration.children),
tokenCount: getGeneratedTokenCountSafe(
getLastNewTokenSpan(iteration.children),
),
rawPrompt,
type: iteration.children.find(
(span) => span.attributes.name === GENERATE_EVENT_TOOL_START,
)
? InterationType.TOOL
: InterationType.FINAL_ANSWER,
};
});
.map((iteration) => ({
groupId: iteration.span.name,
executionTime: getExecutionTime(iteration.children),
tokenCount: getGeneratedTokenCountSafe(
getLastNewTokenSpan(iteration.children),
),
rawPrompt: getRawPrompt(iteration.children),
type: iteration.children.find(
(span) => span.attributes.name === GENERATE_EVENT_TOOL_START,
)
? InterationType.TOOL
: InterationType.FINAL_ANSWER,
}));

const executionTime = getExecutionTime(spans);

Expand All @@ -118,9 +108,11 @@ export function useBuildTraceData({ enabled, threadId, runId }: Props): {
0,
) || getGeneratedTokenCountSafe(getLastNewTokenSpan(spans));

const rawPrompt = iterations.find(
(iteration) => iteration.type === InterationType.FINAL_ANSWER,
)?.rawPrompt;
// load rawPrompt from iterations if they exist. Otherwise, compute it from all spans.
const rawPrompt =
iterations.find(
(iteration) => iteration.type === InterationType.FINAL_ANSWER,
)?.rawPrompt || getRawPrompt(spans);

return { executionTime, tokenCount, rawPrompt, iterations };
}, [spans]);
Expand Down
10 changes: 10 additions & 0 deletions src/modules/chat/trace/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -51,3 +51,13 @@ export function getExecutionTime(spans: TraceSpan[]) {
new Date(firstItem.start_time).getTime()
);
}

export function getRawPrompt(spans: TraceSpan[]): string | undefined {
const customSpanData = spans.find(
(span) => span.attributes.name === 'startCustom',
)?.attributes.data;

return customSpanData?.rawPrompt
? String(customSpanData.rawPrompt)
: undefined;
}

0 comments on commit 586e66f

Please sign in to comment.