diff --git a/app/api/traces/route.ts b/app/api/traces/route.ts index feab6ff9..c41ab4c9 100644 --- a/app/api/traces/route.ts +++ b/app/api/traces/route.ts @@ -10,7 +10,7 @@ export async function POST(req: NextRequest) { try { const session = await getServerSession(authOptions); const apiKey = req.headers.get("x-api-key"); - let { page, pageSize, projectId, filters, group } = await req.json(); + let { page, pageSize, projectId, filters, group, keyword } = await req.json(); // check if user is logged in or has an api key if (!session || !session.user) { @@ -90,7 +90,8 @@ export async function POST(req: NextRequest) { page, pageSize, filters, - group + group, + keyword ); return NextResponse.json( diff --git a/components/project/traces/traces-table.tsx b/components/project/traces/traces-table.tsx index 563e4bd6..c0bcd940 100644 --- a/components/project/traces/traces-table.tsx +++ b/components/project/traces/traces-table.tsx @@ -84,6 +84,7 @@ export function TracesTable({ inputs: true, outputs: true, status: newMode === "trace", + session_id: newMode === "trace", namespace: newMode === "trace", user_ids: newMode === "trace", prompt_ids: newMode === "trace", @@ -306,7 +307,7 @@ export function TracesTable({ )} {!loading && data && data.length > 0 && ( - +
{table.getHeaderGroups().map((headerGroup) => ( diff --git a/components/project/traces/traces.tsx b/components/project/traces/traces.tsx index 9d71b52c..055fe2c9 100644 --- a/components/project/traces/traces.tsx +++ b/components/project/traces/traces.tsx @@ -21,6 +21,7 @@ import { Checkbox } from "../../ui/checkbox"; import { Switch } from "../../ui/switch"; import TraceFilter from "./trace-filter"; import { TracesTable } from "./traces-table"; +import { Input } from "@/components/ui/input"; export default function Traces({ email }: { email: string }) { const project_id = useParams()?.project_id as string; @@ -38,6 +39,7 @@ export default function Traces({ email }: { email: string }) { const [model, setModel] = useState(""); const [expandedView, setExpandedView] = useState(false); const [group, setGroup] = useState(true); + const [keyword, setKeyword] = useState(""); useEffect(() => { setCurrentData([]); @@ -105,6 +107,18 @@ export default function Traces({ email }: { email: string }) { ); }, }, + { + accessorKey: "session_id", + header: "Session ID", + cell: ({ row }) => { + const session_id = row.getValue("session_id") as string; + return ( +
+

{session_id}

+
+ ); + }, + }, { accessorKey: "namespace", header: "Namespace", @@ -349,7 +363,7 @@ export default function Traces({ email }: { email: string }) { async (pageNum: number) => { const apiEndpoint = "/api/traces"; - const body = { + const body: any = { page: pageNum, pageSize: PAGE_SIZE, projectId: project_id, @@ -360,6 +374,10 @@ export default function Traces({ email }: { email: string }) { group: true, }; + if (keyword !== "") { + body.keyword = keyword; + } + const response = await fetch(apiEndpoint, { method: "POST", headers: { @@ -582,6 +600,20 @@ export default function Traces({ email }: { email: string }) { + setKeyword(e.target.value)} + onKeyDown={(e) => { + if (e.key === "Enter") { + // reset everything + setPage(1); + setFilters([]); + setEnableFetch(true); + } + }} + /> string; GetFilteredSpansAttributesQuery: ( tableName: string, filters: Filter, pageSize: number, offset: number, - lastNHours?: number + lastNHours?: number, + keyword?: string ) => string; GetFilteredSpanAttributesSpanById: ( tableName: string, spanId: string, - filters: Filter + filters: Filter, + keyword?: string ) => string; CountFilteredTraceAttributesQuery: ( tableName: string, - filters: Filter + filters: Filter, + keyword?: string ) => string; GetFilteredTraceAttributesQuery: ( tableName: string, filters: Filter, pageSize: number, - offset: number + offset: number, + keyword?: string ) => string; GetFilteredTraceAttributesTraceById: ( tableName: string, traceId: string, - filters: Filter + filters: Filter, + keyword?: string ) => string; } @@ -102,9 +108,14 @@ export class QueryBuilderService implements IQueryBuilderService { return condition; } - CountFilteredSpanAttributesQuery(tableName: string, filters: Filter): string { + CountFilteredSpanAttributesQuery( + tableName: string, + filters: Filter, + keyword?: string + ): string { let baseQuery = `COUNT(DISTINCT span_id) AS total_spans FROM ${tableName}`; let whereConditions: string[] = []; + keyword = keyword?.toLowerCase(); filters.filters.forEach((filter) => { // if it's a property filter @@ -128,16 +139,25 @@ export class QueryBuilderService implements IQueryBuilderService { baseQuery += ` WHERE (${whereConditions.join(` ${filters.operation} `)})`; } + if (keyword !== "" && keyword !== undefined) { + if (baseQuery.includes("WHERE")) { + baseQuery += ` AND (position(lower(CAST(attributes AS String)), '${keyword}') > 0 OR arrayExists(x -> position(lower(x), '${keyword}') > 0, JSONExtractArrayRaw(events)))`; + } else { + baseQuery += ` WHERE (position(lower(CAST(attributes AS String)), '${keyword}') > 0 OR arrayExists(x -> position(lower(x), '${keyword}') > 0, JSONExtractArrayRaw(events)))`; + } + } + return baseQuery; } CountFilteredTraceAttributesQuery( tableName: string, - filters: Filter + filters: Filter, + keyword?: string ): string { let baseQuery = `COUNT(DISTINCT trace_id) AS total_traces FROM ${tableName}`; let whereConditions: string[] = []; - + keyword = keyword?.toLowerCase(); filters.filters.forEach((filter) => { // if it's a property filter if ("key" in filter) { @@ -160,6 +180,14 @@ export class QueryBuilderService implements IQueryBuilderService { baseQuery += ` WHERE (${whereConditions.join(` ${filters.operation} `)})`; } + if (keyword !== "" && keyword !== undefined) { + if (baseQuery.includes("WHERE")) { + baseQuery += ` AND (position(lower(CAST(attributes AS String)), '${keyword}') > 0 OR arrayExists(x -> position(lower(x), '${keyword}') > 0, JSONExtractArrayRaw(events)))`; + } else { + baseQuery += ` WHERE (position(lower(CAST(attributes AS String)), '${keyword}') > 0 OR arrayExists(x -> position(lower(x), '${keyword}') > 0, JSONExtractArrayRaw(events)))`; + } + } + return baseQuery; } @@ -168,11 +196,12 @@ export class QueryBuilderService implements IQueryBuilderService { filters: Filter, pageSize: number = 10, offset: number = 0, - lastNHours?: number + lastNHours?: number, + keyword?: string ): string { let baseQuery = `* FROM ${tableName}`; let whereConditions: string[] = []; - + keyword = keyword?.toLowerCase(); filters.filters.forEach((filter) => { // if it's a property filter if ("key" in filter) { @@ -195,6 +224,14 @@ export class QueryBuilderService implements IQueryBuilderService { baseQuery += ` WHERE (${whereConditions.join(` ${filters.operation} `)})`; } + if (keyword !== "" && keyword !== undefined) { + if (baseQuery.includes("WHERE")) { + baseQuery += ` AND (match(CAST(attributes AS String), '${keyword}') OR arrayExists(x -> position(x, '${keyword}') > 0, JSONExtractArrayRaw(events)))`; + } else { + baseQuery += ` WHERE (match(CAST(attributes AS String), '${keyword}') OR arrayExists(x -> position(x, '${keyword}') > 0, JSONExtractArrayRaw(events)))`; + } + } + if (lastNHours) { const startTime = getFormattedTime(lastNHours); baseQuery += ` AND start_time >= '${startTime}'`; @@ -210,10 +247,12 @@ export class QueryBuilderService implements IQueryBuilderService { tableName: string, filters: Filter, pageSize: number, - offset: number + offset: number, + keyword?: string ): string { let baseQuery = `trace_id, MIN(start_time) AS earliest_start_time FROM ${tableName}`; let whereConditions: string[] = []; + keyword = keyword?.toLowerCase(); filters.filters.forEach((filter) => { // if it's a property filter if ("key" in filter) { @@ -236,6 +275,14 @@ export class QueryBuilderService implements IQueryBuilderService { baseQuery += ` WHERE (${whereConditions.join(` ${filters.operation} `)})`; } + if (keyword !== "" && keyword !== undefined) { + if (baseQuery.includes("WHERE")) { + baseQuery += ` AND (position(lower(CAST(attributes AS String)), '${keyword}') > 0 OR arrayExists(x -> position(lower(x), '${keyword}') > 0, JSONExtractArrayRaw(events)))`; + } else { + baseQuery += ` WHERE (position(lower(CAST(attributes AS String)), '${keyword}') > 0 OR arrayExists(x -> position(lower(x), '${keyword}') > 0, JSONExtractArrayRaw(events)))`; + } + } + baseQuery += ` GROUP BY trace_id ORDER BY earliest_start_time DESC LIMIT ${pageSize} OFFSET ${offset}`; return baseQuery; } @@ -243,11 +290,12 @@ export class QueryBuilderService implements IQueryBuilderService { GetFilteredSpanAttributesSpanById( tableName: string, spanId: string, - filters: Filter + filters: Filter, + keyword?: string ): string { let baseQuery = `* FROM ${tableName} WHERE span_id = '${spanId}'`; let whereConditions: string[] = []; - + keyword = keyword?.toLowerCase(); filters.filters.forEach((filter) => { // if it's a property filter if ("key" in filter) { @@ -270,16 +318,26 @@ export class QueryBuilderService implements IQueryBuilderService { baseQuery += ` AND (${whereConditions.join(` ${filters.operation} `)})`; } + if (keyword !== "" && keyword !== undefined) { + if (baseQuery.includes("WHERE")) { + baseQuery += ` AND (position(lower(CAST(attributes AS String)), '${keyword}') > 0 OR arrayExists(x -> position(lower(x), '${keyword}') > 0, JSONExtractArrayRaw(events)))`; + } else { + baseQuery += ` WHERE (position(lower(CAST(attributes AS String)), '${keyword}') > 0 OR arrayExists(x -> position(lower(x), '${keyword}') > 0, JSONExtractArrayRaw(events)))`; + } + } + return baseQuery; } GetFilteredTraceAttributesTraceById( tableName: string, traceId: string, - filters: Filter + filters: Filter, + keyword?: string ): string { let baseQuery = `* FROM ${tableName} WHERE trace_id = '${traceId}'`; let whereConditions: string[] = []; + keyword = keyword?.toLowerCase(); filters.filters.forEach((filter) => { // if it's a property filter if ("key" in filter) { @@ -302,6 +360,14 @@ export class QueryBuilderService implements IQueryBuilderService { baseQuery += ` AND (${whereConditions.join(` ${filters.operation} `)})`; } + if (keyword !== "" && keyword !== undefined) { + if (baseQuery.includes("WHERE")) { + baseQuery += ` AND (position(lower(CAST(attributes AS String)), '${keyword}') > 0 OR arrayExists(x -> position(lower(x), '${keyword}') > 0, JSONExtractArrayRaw(events)))`; + } else { + baseQuery += ` WHERE (position(lower(CAST(attributes AS String)), '${keyword}') > 0 OR arrayExists(x -> position(lower(x), '${keyword}') > 0, JSONExtractArrayRaw(events)))`; + } + } + return baseQuery; } } diff --git a/lib/services/trace_service.ts b/lib/services/trace_service.ts index 5118a0ec..9cf19974 100644 --- a/lib/services/trace_service.ts +++ b/lib/services/trace_service.ts @@ -660,7 +660,8 @@ export class TraceService implements ITraceService { page: number, pageSize: number, filters: Filter = { operation: "AND", filters: [] }, - group: boolean = true + group: boolean = true, + keyword?: string ): Promise> { try { // check if the table exists @@ -678,7 +679,8 @@ export class TraceService implements ITraceService { const getTotalTracesPerProjectQuery = sql.select( this.queryBuilderService.CountFilteredTraceAttributesQuery( project_id, - filters + filters, + keyword ) ); const result = await this.client.find(getTotalTracesPerProjectQuery); @@ -699,7 +701,8 @@ export class TraceService implements ITraceService { project_id, filters, pageSize, - (page - 1) * pageSize + (page - 1) * pageSize, + keyword ) ); const spans: Span[] = await this.client.find(getTraceIdsQuery); @@ -720,7 +723,8 @@ export class TraceService implements ITraceService { this.queryBuilderService.GetFilteredTraceAttributesTraceById( project_id, span.trace_id, - filters + filters, + keyword ) ); let trace = await this.client.find(getTraceByIdQuery); diff --git a/lib/trace_util.ts b/lib/trace_util.ts index 94c8c1ff..0eb42b17 100644 --- a/lib/trace_util.ts +++ b/lib/trace_util.ts @@ -4,6 +4,7 @@ import { calculatePriceFromUsage } from "./utils"; export interface Trace { id: string; status: string; + session_id: string; namespace: string; user_ids: string[]; prompt_ids: string[]; @@ -30,6 +31,7 @@ export function processTrace(trace: any): Trace { const traceHierarchy = convertTracesToHierarchy(trace); const totalTime = calculateTotalTime(trace); const startTime = trace[0].start_time; + let session_id = ""; let tokenCounts: any = {}; let models: string[] = []; let vendors: string[] = []; @@ -55,6 +57,11 @@ export function processTrace(trace: any): Trace { attributes = JSON.parse(span.attributes); let vendor = ""; + // get session.id from the attributes + if (attributes["session.id"]) { + session_id = attributes["session.id"]; + } + // get the service name from the attributes if (attributes["langtrace.service.name"]) { vendor = attributes["langtrace.service.name"].toLowerCase(); @@ -250,6 +257,7 @@ export function processTrace(trace: any): Trace { const result: Trace = { id: trace[0]?.trace_id, status: status, + session_id: session_id, namespace: traceHierarchy[0].name, user_ids: userIds, prompt_ids: promptIds, diff --git a/lib/ts_sdk_constants.ts b/lib/ts_sdk_constants.ts index 086bc76b..a4145228 100644 --- a/lib/ts_sdk_constants.ts +++ b/lib/ts_sdk_constants.ts @@ -18,7 +18,8 @@ export type VendorsLocal = Vendor; export const SpanAttributes = extractPropertyNames( DatabaseSpanAttributeNames, FrameworkSpanAttributeNames, - LLMSpanAttributeNames + LLMSpanAttributeNames, + ["session.id"], ); export const TracedFunctionsByVendorsLocal = TracedFunctionsByVendor;