From 607a9dad71f56e82f2ac87dd2e6db25511835295 Mon Sep 17 00:00:00 2001 From: Ashish Shubham Date: Mon, 12 May 2025 12:02:47 -0700 Subject: [PATCH] Expose Data sources as resources --- src/mcp-server.ts | 72 ++++++++++++++++++++++++-- src/thoughtspot/relevant-data.ts | 33 ++++++++---- src/thoughtspot/thoughtspot-service.ts | 38 +++++++++++--- 3 files changed, 121 insertions(+), 22 deletions(-) diff --git a/src/mcp-server.ts b/src/mcp-server.ts index e6fe512..fa51ca6 100644 --- a/src/mcp-server.ts +++ b/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; 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) => { + 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) => { const { name } = request.params; @@ -87,13 +131,15 @@ export class MCPServer extends Server { } async callGetRelevantData(request: z.infer) { - 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; + } | 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; + } } diff --git a/src/thoughtspot/relevant-data.ts b/src/thoughtspot/relevant-data.ts index 95c58ec..d700fb1 100644 --- a/src/thoughtspot/relevant-data.ts +++ b/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, }; }; \ No newline at end of file diff --git a/src/thoughtspot/thoughtspot-service.ts b/src/thoughtspot/thoughtspot-service.ts index 9cac861..2952579 100644 --- a/src/thoughtspot/thoughtspot-service.ts +++ b/src/thoughtspot/thoughtspot-service.ts @@ -1,9 +1,6 @@ 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 { +export async function getRelevantQuestions(query: string, sourceId: string, additionalContext: string = '', client: ThoughtSpotRestApi): Promise { const questions = await client.queryGetDecomposedQuery({ nlsRequest: { query: query, @@ -11,7 +8,7 @@ export async function getRelevantQuestions(query: string, additionalContext: str 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 { + 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, + } + }); +}