From cdf8afac25c9bc18f7703752f6227fe2c3a51a35 Mon Sep 17 00:00:00 2001 From: Ashish Shubham Date: Tue, 13 May 2025 11:30:23 -0700 Subject: [PATCH] Split tools into multiple --- src/mcp-server.ts | 117 +++++++++++++++++++++++-- src/thoughtspot/thoughtspot-service.ts | 14 +++ 2 files changed, 123 insertions(+), 8 deletions(-) diff --git a/src/mcp-server.ts b/src/mcp-server.ts index c324790..5f665d6 100644 --- a/src/mcp-server.ts +++ b/src/mcp-server.ts @@ -5,7 +5,7 @@ import { zodToJsonSchema } from "zod-to-json-schema"; import { Props } from "./utils"; import { getRelevantData } from "./thoughtspot/relevant-data"; import { getThoughtSpotClient } from "./thoughtspot/thoughtspot-client"; -import { DataSource, getDataSources } from "./thoughtspot/thoughtspot-service"; +import { DataSource, fetchTMLAndCreateLiveboard, getAnswerForQuestion, getDataSources, getRelevantQuestions } from "./thoughtspot/thoughtspot-service"; const ToolInputSchema = ToolSchema.shape.inputSchema; @@ -13,6 +13,16 @@ type ToolInput = z.infer; const PingSchema = z.object({}); +const GetRelevantQuestionsSchema = z.object({ + query: z.string().describe("The query to get relevant data questions for, this could be a high level task or question the user is asking or hoping to get answered. You can pass the complete raw query as the system is smart to make sense of it."), + additionalContext: z.string() + .describe("Additional context to add to the query, this might be older data returned for previous questions or any other relevant context that might help the system generate better questions.") + .optional(), + datasourceId: z.string() + .describe("The datasource to get questions for, this is the id of the datasource to get data from") + .optional() +}); + const GetRelevantDataSchema = z.object({ query: z.string().describe("The query to get relevant data for, this could be a high level task or question the user is asking or hoping to get answered. You can pass the complete raw query as the system is smart to make sense of it."), datasourceId: z.string() @@ -20,9 +30,26 @@ const GetRelevantDataSchema = z.object({ .optional() }); +const GetAnswerSchema = z.object({ + question: z.string().describe("The question to get the answer for, these are generally the questions generated by the getRelevantQuestions tool."), + datasourceId: z.string() + .describe("The datasource to get the answer for, this is the id of the datasource to get data from") +}); + +const CreateLiveboardSchema = z.object({ + name: z.string().describe("The name of the liveboard to create"), + answers: z.array(z.object({ + question: z.string(), + session_identifier: z.string(), + generation_number: z.number(), + })).describe("The answers to create the liveboard from, these are the answers generated by the getAnswer tool."), +}); + enum ToolName { Ping = "ping", - GetRelevantData = "getRelevantData", + GetRelevantQuestions = "getRelevantQuestions", + GetAnswer = "getAnswer", + CreateLiveboard = "createLiveboard", } interface Context { @@ -54,9 +81,19 @@ export class MCPServer extends Server { inputSchema: zodToJsonSchema(PingSchema) as ToolInput, }, { - name: ToolName.GetRelevantData, - description: "Get relevant data from ThoughtSpot database", - inputSchema: zodToJsonSchema(GetRelevantDataSchema) as ToolInput, + name: ToolName.GetRelevantQuestions, + description: "Get relevant data questions from ThoughtSpot database", + inputSchema: zodToJsonSchema(GetRelevantQuestionsSchema) as ToolInput, + }, + { + name: ToolName.GetAnswer, + description: "Get the answer to a question from ThoughtSpot database", + inputSchema: zodToJsonSchema(GetAnswerSchema) as ToolInput, + }, + { + name: ToolName.CreateLiveboard, + description: "Create a liveboard from a list of answers", + inputSchema: zodToJsonSchema(CreateLiveboardSchema) as ToolInput, } ] }; @@ -95,7 +132,7 @@ export class MCPServer extends Server { The id of the datasource is ${sourceId}. - Use ThoughtSpot's getRelevantData tool to get data from this datasource for a question. + Use ThoughtSpot's getRelevantQuestions tool to get relevant questions for a query. And then use the getAnswer tool to get the answer for a question. `, }], }; @@ -120,8 +157,16 @@ export class MCPServer extends Server { }; } - case ToolName.GetRelevantData: { - return this.callGetRelevantData(request); + case ToolName.GetRelevantQuestions: { + return this.callGetRelevantQuestions(request); + } + + case ToolName.GetAnswer: { + return this.callGetAnswer(request); + } + + case ToolName.CreateLiveboard: { + return this.callCreateLiveboard(request); } default: @@ -130,6 +175,62 @@ export class MCPServer extends Server { }); } + + async callGetRelevantQuestions(request: z.infer) { + const { query, datasourceId: sourceId, additionalContext } = GetRelevantQuestionsSchema.parse(request.params.arguments); + const client = getThoughtSpotClient(this.ctx.props.instanceUrl, this.ctx.props.accessToken); + const progressToken = request.params._meta?.progressToken; + let progress = 0; + console.log("[DEBUG] Getting relevant questions for query: ", query, " and datasource: ", sourceId); + + const relevantQuestions = await getRelevantQuestions( + query, + sourceId!, + additionalContext, + client, + ); + + return { + content: [{ + type: "text", + text: relevantQuestions.map((question) => `- ${question}`).join("\n"), + }], + }; + } + + async callGetAnswer(request: z.infer) { + const { question, datasourceId: sourceId } = GetAnswerSchema.parse(request.params.arguments); + const client = getThoughtSpotClient(this.ctx.props.instanceUrl, this.ctx.props.accessToken); + const progressToken = request.params._meta?.progressToken; + let progress = 0; + console.log("[DEBUG] Getting answer for question: ", question, " and datasource: ", sourceId); + + const answer = await getAnswerForQuestion(question, sourceId, false, client); + + return { + content: [{ + type: "text", + text: answer.data, + }, { + type: "text", + text: `Question: ${question}\nSession Identifier: ${answer.session_identifier}\nGeneration Number: ${answer.generation_number} \n\nUse this information to create a liveboard with the createLiveboard tool, if the user asks.`, + }], + }; + } + + async callCreateLiveboard(request: z.infer) { + const { name, answers } = CreateLiveboardSchema.parse(request.params.arguments); + const client = getThoughtSpotClient(this.ctx.props.instanceUrl, this.ctx.props.accessToken); + const liveboardUrl = await fetchTMLAndCreateLiveboard(name, answers, client); + return { + content: [{ + type: "text", + text: `Liveboard created successfully, you can view it at ${liveboardUrl}`, + }], + }; + } + + async callGetRelevantData(request: z.infer) { const { query, datasourceId: sourceId } = GetRelevantDataSchema.parse(request.params.arguments); const client = getThoughtSpotClient(this.ctx.props.instanceUrl, this.ctx.props.accessToken); diff --git a/src/thoughtspot/thoughtspot-service.ts b/src/thoughtspot/thoughtspot-service.ts index 595429f..3441167 100644 --- a/src/thoughtspot/thoughtspot-service.ts +++ b/src/thoughtspot/thoughtspot-service.ts @@ -81,6 +81,20 @@ export async function getAnswerForQuestion(question: string, sourceId: string, s }; } +export async function fetchTMLAndCreateLiveboard(name: string, answers: any[], client: ThoughtSpotRestApi) { + const tmls = await Promise.all(answers.map((answer) => getAnswerTML({ + question: answer.question, + session_identifier: answer.session_identifier, + generation_number: answer.generation_number, + client, + }))); + answers.forEach((answer, idx) => { + answer.tml = tmls[idx]; + }); + + return createLiveboard(name, answers, client); +} + export async function createLiveboard(name: string, answers: any[], client: ThoughtSpotRestApi) { answers = answers.filter((answer) => answer.tml); const tml = {