From e37bd9236572ed4bca0dd52bf0e8fcee1353f046 Mon Sep 17 00:00:00 2001 From: Ashish Shubham Date: Tue, 13 May 2025 16:33:35 -0700 Subject: [PATCH 1/3] Add API server routes --- src/index.ts | 5 +- src/servers/api-server.ts | 72 ++++++++++++++++++++++++++ src/{ => servers}/mcp-server.ts | 8 +-- src/thoughtspot/thoughtspot-service.ts | 2 +- 4 files changed, 80 insertions(+), 7 deletions(-) create mode 100644 src/servers/api-server.ts rename src/{ => servers}/mcp-server.ts (97%) diff --git a/src/index.ts b/src/index.ts index 515f429..74b5519 100644 --- a/src/index.ts +++ b/src/index.ts @@ -2,8 +2,8 @@ import OAuthProvider from "@cloudflare/workers-oauth-provider"; import { McpAgent } from "agents/mcp"; import handler from "./handlers"; import { Props } from "./utils"; -import { MCPServer } from "./mcp-server"; - +import { MCPServer } from "./servers/mcp-server"; +import { apiServer } from "./servers/api-server"; export class ThoughtSpotMCP extends McpAgent { server = new MCPServer(this); @@ -17,6 +17,7 @@ export default new OAuthProvider({ apiHandlers: { "/mcp": ThoughtSpotMCP.serve("/mcp") as any, // TODO: Remove 'any' "/sse": ThoughtSpotMCP.serveSSE("/sse") as any, // TODO: Remove 'any' + "/api": apiServer as any, // TODO: Remove 'any' }, defaultHandler: handler as any, // TODO: Remove 'any' authorizeEndpoint: "/authorize", diff --git a/src/servers/api-server.ts b/src/servers/api-server.ts new file mode 100644 index 0000000..6242d04 --- /dev/null +++ b/src/servers/api-server.ts @@ -0,0 +1,72 @@ +import { Hono } from 'hono' +import { Props } from '../utils'; +import { createLiveboard, getAnswerForQuestion, getDataSources, getRelevantQuestions } from '../thoughtspot/thoughtspot-service'; +import { getThoughtSpotClient } from '../thoughtspot/thoughtspot-client'; + +const apiServer = new Hono<{ Bindings: Env & { props: Props } }>() + +apiServer.post("/api/tools/relevant-questions", async (c) => { + const { props } = c.executionCtx; + const { query, datasourceId, additionalContext } = await c.req.json(); + const client = getThoughtSpotClient(props.instanceUrl, props.accessToken); + const questions = await getRelevantQuestions(query, datasourceId, additionalContext || '', client); + return c.json(questions); +}); + +apiServer.post("/api/tools/get-answer", async (c) => { + const { props } = c.executionCtx; + const { question, datasourceId } = await c.req.json(); + const client = getThoughtSpotClient(props.instanceUrl, props.accessToken); + const answer = await getAnswerForQuestion(question, datasourceId, false, client); + return c.json(answer); +}); + +apiServer.post("/api/tools/create-liveboard", async (c) => { + const { props } = c.executionCtx; + const { name, answers } = await c.req.json(); + const client = getThoughtSpotClient(props.instanceUrl, props.accessToken); + const liveboardUrl = await createLiveboard(name, answers, client); + return c.text(liveboardUrl); +}); + +apiServer.get("/api/resources/datasources", async (c) => { + const { props } = c.executionCtx; + const client = getThoughtSpotClient(props.instanceUrl, props.accessToken); + const datasources = await getDataSources(client); + return c.json(datasources); +}); + +apiServer.post("/api/rest/2.0/*", async (c) => { + const { props } = c.executionCtx; + const path = c.req.path; + const method = c.req.method; + const body = await c.req.json(); + return fetch(props.instanceUrl + path, { + method, + headers: { + "Authorization": `Bearer ${props.accessToken}`, + "Accept": "application/json", + "Content-Type": "application/json", + "User-Agent": "ThoughtSpot-ts-client", + }, + body: JSON.stringify(body), + }); +}); + +apiServer.get("/api/rest/2.0/*", async (c) => { + const { props } = c.executionCtx; + const path = c.req.path; + const method = c.req.method; + return fetch(props.instanceUrl + path, { + method, + headers: { + "Authorization": `Bearer ${props.accessToken}`, + "Accept": "application/json", + "User-Agent": "ThoughtSpot-ts-client", + } + }); +}); + +export { + apiServer, +} \ No newline at end of file diff --git a/src/mcp-server.ts b/src/servers/mcp-server.ts similarity index 97% rename from src/mcp-server.ts rename to src/servers/mcp-server.ts index 0f61535..20d5ed3 100644 --- a/src/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -2,10 +2,10 @@ import { Server } from "@modelcontextprotocol/sdk/server/index.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, fetchTMLAndCreateLiveboard, getAnswerForQuestion, getDataSources, getRelevantQuestions } from "./thoughtspot/thoughtspot-service"; +import { Props } from "../utils"; +import { getRelevantData } from "../thoughtspot/relevant-data"; +import { getThoughtSpotClient } from "../thoughtspot/thoughtspot-client"; +import { DataSource, fetchTMLAndCreateLiveboard, getAnswerForQuestion, getDataSources, getRelevantQuestions } from "../thoughtspot/thoughtspot-service"; const ToolInputSchema = ToolSchema.shape.inputSchema; diff --git a/src/thoughtspot/thoughtspot-service.ts b/src/thoughtspot/thoughtspot-service.ts index 99dcb13..9bfef24 100644 --- a/src/thoughtspot/thoughtspot-service.ts +++ b/src/thoughtspot/thoughtspot-service.ts @@ -140,7 +140,7 @@ export async function getDataSources(client: ThoughtSpotRestApi): Promise Date: Tue, 13 May 2025 16:37:16 -0700 Subject: [PATCH 2/3] newline imports --- src/servers/api-server.ts | 7 ++++++- src/servers/mcp-server.ts | 16 ++++++++++++++-- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/src/servers/api-server.ts b/src/servers/api-server.ts index 6242d04..fcc9d0e 100644 --- a/src/servers/api-server.ts +++ b/src/servers/api-server.ts @@ -1,6 +1,11 @@ import { Hono } from 'hono' import { Props } from '../utils'; -import { createLiveboard, getAnswerForQuestion, getDataSources, getRelevantQuestions } from '../thoughtspot/thoughtspot-service'; +import { + createLiveboard, + getAnswerForQuestion, + getDataSources, + getRelevantQuestions +} from '../thoughtspot/thoughtspot-service'; import { getThoughtSpotClient } from '../thoughtspot/thoughtspot-client'; const apiServer = new Hono<{ Bindings: Env & { props: Props } }>() diff --git a/src/servers/mcp-server.ts b/src/servers/mcp-server.ts index 20d5ed3..87e3a47 100644 --- a/src/servers/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -1,11 +1,23 @@ import { Server } from "@modelcontextprotocol/sdk/server/index.js"; -import { CallToolRequestSchema, ListToolsRequestSchema, ToolSchema, Tool, ListResourcesRequestSchema, ReadResourceRequestSchema } from "@modelcontextprotocol/sdk/types.js"; +import { + CallToolRequestSchema, + ListToolsRequestSchema, + ToolSchema, + 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, fetchTMLAndCreateLiveboard, getAnswerForQuestion, getDataSources, getRelevantQuestions } from "../thoughtspot/thoughtspot-service"; +import { + DataSource, + fetchTMLAndCreateLiveboard, + getAnswerForQuestion, + getDataSources, + getRelevantQuestions +} from "../thoughtspot/thoughtspot-service"; const ToolInputSchema = ToolSchema.shape.inputSchema; From 944d5d58d4cfcdf2963305335438387cf4b610e9 Mon Sep 17 00:00:00 2001 From: Ashish Shubham Date: Wed, 14 May 2025 16:55:23 -0700 Subject: [PATCH 3/3] Remove relevant data + error handling --- package-lock.json | 9 +- package.json | 2 +- src/servers/api-server.ts | 4 +- src/servers/mcp-server.ts | 91 ++++++-------------- src/thoughtspot/relevant-data.ts | 63 -------------- src/thoughtspot/thoughtspot-service.ts | 110 ++++++++++++++++--------- 6 files changed, 104 insertions(+), 175 deletions(-) delete mode 100644 src/thoughtspot/relevant-data.ts diff --git a/package-lock.json b/package-lock.json index a2b9e99..573fa32 100644 --- a/package-lock.json +++ b/package-lock.json @@ -9,7 +9,7 @@ "version": "1.0.0", "license": "ISC", "dependencies": { - "@cloudflare/workers-oauth-provider": "^0.0.4", + "@cloudflare/workers-oauth-provider": "^0.0.5", "@modelcontextprotocol/sdk": "^1.10.2", "@thoughtspot/rest-api-sdk": "^2.13.1", "agents": "^0.0.75", @@ -733,9 +733,10 @@ } }, "node_modules/@cloudflare/workers-oauth-provider": { - "version": "0.0.4", - "resolved": "https://registry.npmjs.org/@cloudflare/workers-oauth-provider/-/workers-oauth-provider-0.0.4.tgz", - "integrity": "sha512-6VOfn+22VDHOYlE1kokZzQIxzCjwi9dha+vESSroULYAl5o/17w7S1IlRMcZlAapKGD/a3UdId0cM7tUDEaR5g==", + "version": "0.0.5", + "resolved": "https://registry.npmjs.org/@cloudflare/workers-oauth-provider/-/workers-oauth-provider-0.0.5.tgz", + "integrity": "sha512-t1x5KAzsubCvb4APnJ93z407X1x7SGj/ga5ziRnwIb/iLy4PMkT/hgd1y5z7Bbsdy5Fy6mywhCP4lym24bX66w==", + "license": "MIT", "dependencies": { "@cloudflare/workers-types": "^4.20250311.0" } diff --git a/package.json b/package.json index af0f854..97395e8 100644 --- a/package.json +++ b/package.json @@ -34,7 +34,7 @@ "wrangler": "^4.12.0" }, "dependencies": { - "@cloudflare/workers-oauth-provider": "^0.0.4", + "@cloudflare/workers-oauth-provider": "^0.0.5", "@modelcontextprotocol/sdk": "^1.10.2", "@thoughtspot/rest-api-sdk": "^2.13.1", "agents": "^0.0.75", diff --git a/src/servers/api-server.ts b/src/servers/api-server.ts index fcc9d0e..7820f52 100644 --- a/src/servers/api-server.ts +++ b/src/servers/api-server.ts @@ -12,9 +12,9 @@ const apiServer = new Hono<{ Bindings: Env & { props: Props } }>() apiServer.post("/api/tools/relevant-questions", async (c) => { const { props } = c.executionCtx; - const { query, datasourceId, additionalContext } = await c.req.json(); + const { query, datasourceIds, additionalContext } = await c.req.json(); const client = getThoughtSpotClient(props.instanceUrl, props.accessToken); - const questions = await getRelevantQuestions(query, datasourceId, additionalContext || '', client); + const questions = await getRelevantQuestions(query, datasourceIds, additionalContext || '', client); return c.json(questions); }); diff --git a/src/servers/mcp-server.ts b/src/servers/mcp-server.ts index 87e3a47..c3f690b 100644 --- a/src/servers/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -9,7 +9,6 @@ import { 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, @@ -30,16 +29,8 @@ const GetRelevantQuestionsSchema = z.object({ 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() + datasourceIds: z.array(z.string()) + .describe("The datasources to get questions for, this is the ids of the datasources to get data from") }); const GetAnswerSchema = z.object({ @@ -189,24 +180,29 @@ export class MCPServer extends Server { async callGetRelevantQuestions(request: z.infer) { - const { query, datasourceId: sourceId, additionalContext } = GetRelevantQuestionsSchema.parse(request.params.arguments); + const { query, datasourceIds: sourceIds, 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); + console.log("[DEBUG] Getting relevant questions for query: ", query, " and datasource: ", sourceIds); const relevantQuestions = await getRelevantQuestions( query, - sourceId!, + sourceIds!, additionalContext, client, ); + if (relevantQuestions.error) { + return { + isError: true, + content: [{ type: "text", text: "ERROR: " + relevantQuestions.error.message }], + }; + } + return { - content: [{ + content: relevantQuestions.questions.map(q => ({ type: "text", - text: relevantQuestions.map((question) => `- ${question}`).join("\n"), - }], + text: `Question: ${q.question}\nDatasourceId: ${q.datasourceId}`, + })), }; } @@ -218,6 +214,12 @@ export class MCPServer extends Server { console.log("[DEBUG] Getting answer for question: ", question, " and datasource: ", sourceId); const answer = await getAnswerForQuestion(question, sourceId, false, client); + if (answer.error) { + return { + isError: true, + content: [{ type: "text", text: "ERROR: " + answer.error.message }], + }; + } return { content: [{ @@ -233,60 +235,19 @@ export class MCPServer extends Server { 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} - - Provide this url to the user as a link to view the liveboard in ThoughtSpot.`, - }], - }; - } - - - 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); - 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", - params: { - message: data, - progressToken: progressToken, - progress: Math.min(progress++ * 10, 100), - total: 100, - }, - }), - client, - }); - - if (relevantData.allAnswers.length === 0) { + const liveboard = await fetchTMLAndCreateLiveboard(name, answers, client); + if (liveboard.error) { return { isError: true, - content: [{ - type: "text", - text: "No relevant data found, please make sure the datasource is correct, and you have data download privileges in ThoughtSpot.", - }], + content: [{ type: "text", text: "ERROR: " + liveboard.error.message }], }; } - return { content: [{ type: "text", - text: relevantData.allAnswers.map((answer) => `Question: ${answer.question}\nAnswer: ${answer.data}`).join("\n\n") - }, { - type: "text", - text: `Dashboard Url: ${relevantData.liveboard} + text: `Liveboard created successfully, you can view it at ${liveboard.url} - Use this url to view the dashboard/liveboard in ThoughtSpot which contains visualizations for the generated data. *Always* Present this url to the user as a link to view the data as a reference.`, + Provide this url to the user as a link to view the liveboard in ThoughtSpot.`, }], }; } diff --git a/src/thoughtspot/relevant-data.ts b/src/thoughtspot/relevant-data.ts deleted file mode 100644 index d700fb1..0000000 --- a/src/thoughtspot/relevant-data.ts +++ /dev/null @@ -1,63 +0,0 @@ -import { createLiveboard, getAnswerForQuestion, getRelevantQuestions } from "./thoughtspot-service"; -import { ThoughtSpotRestApi } from "@thoughtspot/rest-api-sdk"; - -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, sourceId, shouldGetTML, client); - } catch (error) { - console.error(`Failed to get answer for question: ${question}`, error); - return null; - } - }) - )).filter((answer): answer is NonNullable => answer !== null); - - notify(`\n\nRetrieved ${answers.length} answers using **ThoughtSpot Spotter**\n\n`); - return answers; -} - - - -export const getRelevantData = async ({ - query, - sourceId, - shouldCreateLiveboard, - notify, - client, -}: { - query: string; - sourceId?: string; - shouldCreateLiveboard: boolean; - notify: (data: string) => void; - client: ThoughtSpotRestApi; -}) => { - 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")}`); - - let answers = await getAnswersForQuestions(questions, sourceId, shouldCreateLiveboard, notify, client); - - 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")}`); - - const additionalAnswers = await getAnswersForQuestions(additionalQuestions, sourceId, shouldCreateLiveboard, notify, client); - - answers = [...answers, ...additionalAnswers]; - } - - const liveboard = shouldCreateLiveboard ? await createLiveboard(query, answers, client) : null; - return { - allAnswers: answers, - liveboard, - }; -}; \ No newline at end of file diff --git a/src/thoughtspot/thoughtspot-service.ts b/src/thoughtspot/thoughtspot-service.ts index 9bfef24..9759de7 100644 --- a/src/thoughtspot/thoughtspot-service.ts +++ b/src/thoughtspot/thoughtspot-service.ts @@ -1,21 +1,31 @@ import { ThoughtSpotRestApi } from "@thoughtspot/rest-api-sdk"; -export async function getRelevantQuestions(query: string, sourceId: string, additionalContext: string = '', client: ThoughtSpotRestApi): Promise { +export async function getRelevantQuestions(query: string, sourceIds: string[], additionalContext: string = '', client: ThoughtSpotRestApi): Promise<{ questions: { question: string, datasourceId: string }[], error: Error | null }> { try { - const questions = await client.queryGetDecomposedQuery({ + const resp = await client.queryGetDecomposedQuery({ nlsRequest: { query: query, }, content: [ additionalContext, ], - worksheetIds: [sourceId], - maxDecomposedQueries: 7, + worksheetIds: sourceIds, + maxDecomposedQueries: 5, }) - return questions.decomposedQueryResponse?.decomposedQueries?.map((q) => q.query!) || []; + const questions = resp.decomposedQueryResponse?.decomposedQueries?.map((q) => ({ + question: q.query!, + datasourceId: q.worksheetId!, + })) || []; + return { + questions, + error: null, + } } catch (error) { console.error("Error getting relevant questions: ", error); - return []; + return { + questions: [], + error: error as Error, + } } } @@ -34,7 +44,7 @@ async function getAnswerData({ question, session_identifier, generation_number, return csvData; } catch (error) { console.error("Error getting answer Data: ", error); - return null; + throw error; } } @@ -54,50 +64,70 @@ async function getAnswerTML({ question, session_identifier, generation_number, c 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: sourceId, - }) + try { + const answer = await client.singleAnswer({ + query: question, + metadata_identifier: sourceId, + }) - const { session_identifier, generation_number } = answer as any; + const { session_identifier, generation_number } = answer as any; - const [data, tml] = await Promise.all([ - getAnswerData({ - question, - session_identifier, - generation_number, - client - }), - shouldGetTML - ? getAnswerTML({ + const [data, tml] = await Promise.all([ + getAnswerData({ question, session_identifier, generation_number, client - }) - : Promise.resolve(null) - ]) + }), + shouldGetTML + ? getAnswerTML({ + question, + session_identifier, + generation_number, + client + }) + : Promise.resolve(null) + ]) - return { - question, - ...answer, - data, - tml, - }; + return { + question, + ...answer, + data, + tml, + error: null, + }; + } catch (error) { + console.error("Error getting answer for question: ", question, " and sourceId: ", sourceId, " and shouldGetTML: ", shouldGetTML, " and error: ", error); + return { + error: error as Error, + }; + } } 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]; - }); + try { + 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); + const liveboardUrl = await createLiveboard(name, answers, client); + return { + url: liveboardUrl, + error: null, + } + } catch (error) { + console.error("Error fetching TML and creating liveboard: ", error); + return { + liveboardUrl: null, + error: error as Error, + } + } } export async function createLiveboard(name: string, answers: any[], client: ThoughtSpotRestApi) {