Add API server routes (#6)

* Add API server routes

* newline imports

* Remove relevant data + error handling

Author: ashubham

package-lock.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",

package.json

@@ -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<Env, any, Props> {
     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",

src/index.ts

@@ -0,0 +1,77 @@
+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, datasourceIds, additionalContext } = await c.req.json();
+    const client = getThoughtSpotClient(props.instanceUrl, props.accessToken);
+    const questions = await getRelevantQuestions(query, datasourceIds, 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,
+}

src/servers/api-server.ts

@@ -1,11 +1,22 @@
 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 { Props } from "../utils";
+import { getThoughtSpotClient } from "../thoughtspot/thoughtspot-client";
+import {
+    DataSource,
+    fetchTMLAndCreateLiveboard,
+    getAnswerForQuestion,
+    getDataSources,
+    getRelevantQuestions
+} from "../thoughtspot/thoughtspot-service";
 
 
 const ToolInputSchema = ToolSchema.shape.inputSchema;
@@ -18,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({
@@ -177,24 +180,29 @@ export class MCPServer extends Server {
 
 
     async callGetRelevantQuestions(request: z.infer<typeof CallToolRequestSchema>) {
-        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}`,
+            })),
         };
     }
 
@@ -206,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: [{
@@ -221,60 +235,19 @@ export class MCPServer extends Server {
     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}
-                
-                Provide this url to the user as a link to view the liveboard in ThoughtSpot.`,
-            }],
-        };
-    }
-
-
-    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);
-        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.`,
             }],
         };
     }