Split tools into multiple (#5)

Author: ashubham

src/mcp-server.ts

@@ -5,24 +5,51 @@ 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;
 type ToolInput = z.infer<typeof ToolInputSchema>;
 
 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()
         .describe("The datasource to get data from, this is the id of the datasource to get data from")
         .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<typeof CallToolRequestSchema>) {
+        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<typeof CallToolRequestSchema>) {
+        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<typeof CallToolRequestSchema>) {
+        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<typeof CallToolRequestSchema>) {
         const { query, datasourceId: sourceId } = GetRelevantDataSchema.parse(request.params.arguments);
         const client = getThoughtSpotClient(this.ctx.props.instanceUrl, this.ctx.props.accessToken);

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 = {