Expose Data sources as resources

src/mcp-server.ts

@@ -1,10 +1,11 @@
 import { Server } from "@modelcontextprotocol/sdk/server/index.js";
-import { CallToolRequestSchema, ListToolsRequestSchema, ToolSchema, Tool } from "@modelcontextprotocol/sdk/types.js";
+import { CallToolRequestSchema, ListToolsRequestSchema, ToolSchema, Tool, ListResourcesRequestSchema, ReadResourceRequestSchema } from "@modelcontextprotocol/sdk/types.js";
 import { z } from "zod";
 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";
 
 
 const ToolInputSchema = ToolSchema.shape.inputSchema;
@@ -13,7 +14,10 @@ type ToolInput = z.infer<typeof ToolInputSchema>;
 const PingSchema = z.object({});
 
 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")
+    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()
 });
 
 enum ToolName {
@@ -58,6 +62,46 @@ export class MCPServer extends Server {
             };
         });
 
+        this.setRequestHandler(ListResourcesRequestSchema, async () => {
+            const client = getThoughtSpotClient(this.ctx.props.instanceUrl, this.ctx.props.accessToken);
+            const sources = await this.getDatasources();
+            return {
+                resources: sources.list.map((s) => ({
+                    uri: `datasource:///${s.id}`,
+                    name: s.name,
+                    description: s.description,
+                    mimeType: "text/plain"
+                }))
+            }
+        });
+
+        this.setRequestHandler(ReadResourceRequestSchema, async (request: z.infer<typeof ReadResourceRequestSchema>) => {
+            const { uri } = request.params;
+            const sourceId = uri.split("///").pop();
+            if (!sourceId) {
+                throw new Error("Invalid datasource uri");
+            }
+            const { map: sourceMap } = await this.getDatasources();
+            const source = sourceMap.get(sourceId);
+            if (!source) {
+                throw new Error("Datasource not found");
+            }
+            return {
+                contents: [{
+                    uri: uri,
+                    mimeType: "text/plain",
+                    text: `
+                    ${source.description}
+
+                    The id of the datasource is ${sourceId}.
+
+                    Use ThoughtSpot's getRelevantData tool to get data from this datasource for a question.
+                    `,
+                }],
+            };
+        });
+
+
         // Handle call tool request
         this.setRequestHandler(CallToolRequestSchema, async (request: z.infer<typeof CallToolRequestSchema>) => {
             const { name } = request.params;
@@ -87,13 +131,15 @@ export class MCPServer extends Server {
     }
 
     async callGetRelevantData(request: z.infer<typeof CallToolRequestSchema>) {
-        const { query } = GetRelevantDataSchema.parse(request.params.arguments);
+        const { query, datasourceId: sourceId } = GetRelevantDataSchema.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 data for query: ", query, " and datasource: ", sourceId);
 
         const relevantData = await getRelevantData({
             query,
+            sourceId,
             shouldCreateLiveboard: true,
             notify: (data) => this.notification({
                 method: "notifications/progress",
@@ -113,8 +159,26 @@ export class MCPServer extends Server {
                 text: relevantData.allAnswers.map((answer) => `Question: ${answer.question}\nAnswer: ${answer.data}`).join("\n\n")
             }, {
                 type: "text",
-                text: `Dashboard Url: ${relevantData.liveboard}`,
+                text: `Dashboard Url: ${relevantData.liveboard} Use this url to view the dashboard in ThoughtSpot which contains visualizations for the generated data. Present this url to the user as a link to view the data as a reference.`,
             }],
         };
     }
+
+    private _sources: {
+        list: DataSource[];
+        map: Map<string, DataSource>;
+    } | null = null;
+    async getDatasources() {
+        if (this._sources) {
+            return this._sources;
+        }
+
+        const client = getThoughtSpotClient(this.ctx.props.instanceUrl, this.ctx.props.accessToken);
+        const sources = await getDataSources(client);
+        this._sources = {
+            list: sources,
+            map: new Map(sources.map(s => [s.id, s])),
+        }
+        return this._sources;
+    }
 }

src/thoughtspot/relevant-data.ts

@@ -1,10 +1,15 @@
 import { createLiveboard, getAnswerForQuestion, getRelevantQuestions } from "./thoughtspot-service";
 import { ThoughtSpotRestApi } from "@thoughtspot/rest-api-sdk";
-async function getAnswersForQuestions(questions: string[], shouldGetTML: boolean, notify: (data: string) => void, client: ThoughtSpotRestApi) {
+
+const DEFAULT_DATA_SOURCE_ID = "cd252e5c-b552-49a8-821d-3eadaa049cca";
+const DO_ADDITIONAL_QUESTIONS = false;
+
+
+async function getAnswersForQuestions(questions: string[], sourceId: string, shouldGetTML: boolean, notify: (data: string) => void, client: ThoughtSpotRestApi) {
     const answers = (await Promise.all(
         questions.map(async (question) => {
             try {
-                return await getAnswerForQuestion(question, shouldGetTML, client);
+                return await getAnswerForQuestion(question, sourceId, shouldGetTML, client);
             } catch (error) {
                 console.error(`Failed to get answer for question: ${question}`, error);
                 return null;
@@ -20,33 +25,39 @@ async function getAnswersForQuestions(questions: string[], shouldGetTML: boolean
 
 export const getRelevantData = async ({
     query,
+    sourceId,
     shouldCreateLiveboard,
     notify,
     client,
 }: {
     query: string;
+    sourceId?: string;
     shouldCreateLiveboard: boolean;
     notify: (data: string) => void;
     client: ThoughtSpotRestApi;
 }) => {
-    const questions = await getRelevantQuestions(query, "", client);
+    sourceId = sourceId || DEFAULT_DATA_SOURCE_ID;
+    const questions = await getRelevantQuestions(query, sourceId, "", client);
     notify(`#### Retrieving answers to these relevant questions:\n ${questions.map((q) => `- ${q}`).join("\n")}`);
 
-    const answers = await getAnswersForQuestions(questions, shouldCreateLiveboard, notify, client);
+    let answers = await getAnswersForQuestions(questions, sourceId, shouldCreateLiveboard, notify, client);
 
-    const additionalQuestions = await getRelevantQuestions(query, `
-        These questions have been answered already (with their csv data): ${answers.map((a) => `Question: ${a.question} \n CSV data: \n${a.data}`).join("\n\n ")}
+    if (DO_ADDITIONAL_QUESTIONS) {
+        const additionalQuestions = await getRelevantQuestions(query, sourceId, `
+            These questions have been answered already (with their csv data): ${answers.map((a) => `Question: ${a.question} \n CSV data: \n${a.data}`).join("\n\n ")}
         Look at the csv data of the above queries to see if you need additional related queries to be answered. You can also ask questions going deeper into the data returned by applying filters.
         Do NOT resend the same query already asked before.
     `, client);
-    notify(`#### Need to get answers to some of these additional questions:\n ${additionalQuestions.map((q) => `- ${q}`).join("\n")}`);
+        notify(`#### Need to get answers to some of these additional questions:\n ${additionalQuestions.map((q) => `- ${q}`).join("\n")}`);
+
+        const additionalAnswers = await getAnswersForQuestions(additionalQuestions, sourceId, shouldCreateLiveboard, notify, client);
 
-    const additionalAnswers = await getAnswersForQuestions(additionalQuestions, shouldCreateLiveboard, notify, client);
+        answers = [...answers, ...additionalAnswers];
+    }
 
-    const allAnswers = [...answers, ...additionalAnswers];
-    const liveboard = shouldCreateLiveboard ? await createLiveboard(query, allAnswers, client) : null;
+    const liveboard = shouldCreateLiveboard ? await createLiveboard(query, answers, client) : null;
     return {
-        allAnswers,
+        allAnswers: answers,
         liveboard,
     };
 };

src/thoughtspot/thoughtspot-service.ts

@@ -1,17 +1,14 @@
 import { ThoughtSpotRestApi } from "@thoughtspot/rest-api-sdk";
 
-const DATA_SOURCE_ID = "cd252e5c-b552-49a8-821d-3eadaa049cca";
-
-
-export async function getRelevantQuestions(query: string, additionalContext: string = '', client: ThoughtSpotRestApi): Promise<string[]> {
+export async function getRelevantQuestions(query: string, sourceId: string, additionalContext: string = '', client: ThoughtSpotRestApi): Promise<string[]> {
     const questions = await client.queryGetDecomposedQuery({
         nlsRequest: {
             query: query,
         },
         content: [
             additionalContext,
         ],
-        worksheetIds: [DATA_SOURCE_ID]
+        worksheetIds: [sourceId]
     })
     return questions.decomposedQueryResponse?.decomposedQueries?.map((q) => q.query!) || [];
 }
@@ -49,11 +46,11 @@ async function getAnswerTML({ question, session_identifier, generation_number, c
     }
 }
 
-export async function getAnswerForQuestion(question: string, shouldGetTML: boolean, client: ThoughtSpotRestApi) {
+export async function getAnswerForQuestion(question: string, sourceId: string, shouldGetTML: boolean, client: ThoughtSpotRestApi) {
     console.log("[DEBUG] Getting answer for question: ", question);
     const answer = await client.singleAnswer({
         query: question,
-        metadata_identifier: DATA_SOURCE_ID,
+        metadata_identifier: sourceId,
     })
 
     const { session_identifier, generation_number } = answer as any;
@@ -112,3 +109,30 @@ export async function createLiveboard(name: string, answers: any[], client: Thou
     return `${(client as any).instanceUrl}/#/pinboard/${resp[0].response.header.id_guid}`;
 }
 
+export interface DataSource {
+    name: string;
+    id: string;
+    description: string;
+}
+
+export async function getDataSources(client: ThoughtSpotRestApi): Promise<DataSource[]> {
+    const resp = await client.searchMetadata({
+        metadata: [{
+            type: "LOGICAL_TABLE",
+        }],
+        record_size: 1000,
+        sort_options: {
+            field_name: "LAST_ACCESSED",
+            order: "DESC",
+        }
+    });
+    return resp
+        .filter(d => d.metadata_header.type === "WORKSHEET")
+        .map(d => {
+            return {
+                name: d.metadata_header.name,
+                id: d.metadata_header.id,
+                description: d.metadata_header.description,
+            }
+        });
+}