|
| 1 | +import { ExtendedTool, ToolHandlers } from '../../utils/types' |
| 2 | +import { v2 } from '@datadog/datadog-api-client' |
| 3 | +import { createToolSchema } from '../../utils/tool' |
| 4 | +import { datadogConfig as config } from '../../utils/datadog' |
| 5 | +import { ListTracesZodSchema } from './schema' |
| 6 | + |
| 7 | +type TracesToolName = 'list_traces' |
| 8 | +type TracesTool = ExtendedTool<TracesToolName> |
| 9 | + |
| 10 | +export const TRACES_TOOLS: TracesTool[] = [ |
| 11 | + createToolSchema( |
| 12 | + ListTracesZodSchema, |
| 13 | + 'list_traces', |
| 14 | + 'Get APM traces from Datadog', |
| 15 | + ), |
| 16 | +] as const |
| 17 | + |
| 18 | +const API_INSTANCE = new v2.SpansApi(config) |
| 19 | + |
| 20 | +type TracesToolHandlers = ToolHandlers<TracesToolName> |
| 21 | + |
| 22 | +export const TRACES_HANDLERS: TracesToolHandlers = { |
| 23 | + list_traces: async (request) => { |
| 24 | + const { |
| 25 | + query, |
| 26 | + from, |
| 27 | + to, |
| 28 | + limit = 100, |
| 29 | + sort = '-timestamp', |
| 30 | + service, |
| 31 | + operation, |
| 32 | + } = request.params.arguments as { |
| 33 | + query: string |
| 34 | + from: number |
| 35 | + to: number |
| 36 | + limit?: number |
| 37 | + sort?: string |
| 38 | + service?: string |
| 39 | + operation?: string |
| 40 | + } |
| 41 | + |
| 42 | + const response = await API_INSTANCE.listSpans({ |
| 43 | + body: { |
| 44 | + data: { |
| 45 | + attributes: { |
| 46 | + filter: { |
| 47 | + query: [ |
| 48 | + query, |
| 49 | + ...(service ? [`service:${service}`] : []), |
| 50 | + ...(operation ? [`operation:${operation}`] : []), |
| 51 | + ].join(' '), |
| 52 | + from: new Date(from * 1000).toISOString(), |
| 53 | + to: new Date(to * 1000).toISOString(), |
| 54 | + }, |
| 55 | + sort: sort as 'timestamp' | '-timestamp', |
| 56 | + page: { limit }, |
| 57 | + }, |
| 58 | + type: 'search_request', |
| 59 | + }, |
| 60 | + }, |
| 61 | + }) |
| 62 | + |
| 63 | + if (!response.data) { |
| 64 | + throw new Error('No traces data returned') |
| 65 | + } |
| 66 | + |
| 67 | + return { |
| 68 | + content: [ |
| 69 | + { |
| 70 | + type: 'text', |
| 71 | + text: `Traces: ${JSON.stringify({ |
| 72 | + traces: response.data, |
| 73 | + count: response.data.length, |
| 74 | + })}`, |
| 75 | + }, |
| 76 | + ], |
| 77 | + } |
| 78 | + }, |
| 79 | +} as const |
0 commit comments