Datasource discovery tool (#59)

* [WIP] Add getDataSourceSuggestions tool

* make sure datasource discovery tool is invoked only if cluster version is > 10.13.0.cl

* address comments

* add test for number of tools returned

Author: shikharbsar

src/servers/mcp-server-base.ts

@@ -63,6 +63,14 @@ export abstract class BaseMCPServer extends Server {
         });
     }
 
+    /**
+     * Check if data source discovery is available
+     */
+    protected isDatasourceDiscoveryAvailable(): boolean {
+        const enableSpotterDataSourceDiscovery = this.sessionInfo?.enableSpotterDataSourceDiscovery;
+        return !!enableSpotterDataSourceDiscovery;
+    };
+
     /**
      * Initialize span with common attributes (user_guid and instance_url)
      */

src/servers/mcp-server.ts

@@ -66,11 +66,17 @@ const CreateLiveboardSchema = z.object({
 
 });
 
+const GetDataSourceSuggestionsSchema = z.object({
+    query: z.string().describe(`The query to get data source suggestions for, this could be a high level task or question the user is asking or hoping to get answered.
+         There can be multiple data sources. Each data source can be used to get the data for the user's query using the other tools getRelevantQuestions and getAnswer.`),
+});
+
 enum ToolName {
     Ping = "ping",
     GetRelevantQuestions = "getRelevantQuestions",
     GetAnswer = "getAnswer",
     CreateLiveboard = "createLiveboard",
+    GetDataSourceSuggestions = "getDataSourceSuggestions",
 }
 
 export class MCPServer extends BaseMCPServer {
@@ -120,7 +126,17 @@ export class MCPServer extends BaseMCPServer {
                         readOnlyHint: true,
                         destructiveHint: false,
                     },
-                }
+                },
+                ...(this.isDatasourceDiscoveryAvailable() ? [{
+                    name: ToolName.GetDataSourceSuggestions,
+                    description: "Get data source suggestions for a query. Use this tool only if there is not datasource id provided in the context or the users query. If mulitple data sources are returned, and the confidence difference between the top two data sources is less than 0.3, ask the user to select the most relevant data source. Otherwise use the data source with the highest confidence to get the relevant questions and answers for the query.",
+                    inputSchema: zodToJsonSchema(GetDataSourceSuggestionsSchema) as ToolInput,
+                    annotations: {
+                        title: "Get Data Source Suggestions for a Query",
+                        readOnlyHint: true,
+                        destructiveHint: false,
+                    },
+                }] : []),
             ]
         };
     }
@@ -188,6 +204,10 @@ export class MCPServer extends BaseMCPServer {
                 return this.callCreateLiveboard(request);
             }
 
+            case ToolName.GetDataSourceSuggestions: {
+                return this.callGetDataSourceSuggestions(request);
+            }
+
             default:
                 throw new Error(`Unknown tool: ${name}`);
         }
@@ -255,6 +275,25 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`;
         return this.createSuccessResponse(successMessage, "Liveboard created successfully");
     }
 
+    @WithSpan('call-get-data-source-suggestions')
+    async callGetDataSourceSuggestions(request: z.infer<typeof CallToolRequestSchema>) {
+        const { query } = GetDataSourceSuggestionsSchema.parse(request.params.arguments);
+        const dataSources = await this.getThoughtSpotService().getDataSourceSuggestions(query);
+
+        if (!dataSources || dataSources.length === 0) {
+            return this.createErrorResponse("No data source suggestions found", "No data source suggestions found");
+        }
+
+        // Return information for all suggested data sources
+        const dataSourcesInfo = dataSources.map(ds => ({
+            header: ds.header,
+            confidence: ds.confidence,
+            llmReasoning: ds.llmReasoning
+        }));
+
+        return this.createSuccessResponse(JSON.stringify(dataSourcesInfo), `${dataSources.length} data source suggestion(s) found`);
+    }
+
     private _sources: {
         list: DataSource[];
         map: Map<string, DataSource>;

src/servers/openai-mcp-server.ts

@@ -120,10 +120,25 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer {
 
             return this.createStructuredContentSuccessResponse({ results }, "Relevant questions found");
         }
-
         // Search for datasources in case the query is not of the form "datasource:<id> <query-with-spaces>"
-        // TODO: Implement this
-        return this.createStructuredContentSuccessResponse({ results: [] }, "No relevant questions found");
+        if (!this.isDatasourceDiscoveryAvailable()) {
+            return this.createStructuredContentSuccessResponse({ results: [] }, "No relevant questions found");
+        }
+        const dataSources = await this.getThoughtSpotService().getDataSourceSuggestions(queryWithoutDatasourceId);
+        if (!dataSources || dataSources.length === 0) {
+            return this.createSuccessResponse("No relevant data sources found, please provide a datasource id in the query");
+        }
+        const results = dataSources.map(d => ({
+            id: `datasource:///${d.header.guid}`,
+            title: d.header.displayName,
+            text: `Datasource Description: ${d.header.description}. Confidence that this datasource is relevant to the query: ${d.confidence}. Reasoning for the confidence: ${d.llmReasoning}. 
+            Use this datasource to search for relevant questions and to get answers for the questions. 
+            Use the search tool to search for relevant questions with the format "datasource:<id> <query-with-spaces>" and the fetch tool to get answers for the questions.`,
+        }));
+
+        return this.createStructuredContentSuccessResponse({ results }, "Relevant questions found");
+
+
     }
 
     @WithSpan('call-fetch')

src/thoughtspot/thoughtspot-client.ts

@@ -3,7 +3,7 @@ import type { RequestContext, ResponseContext } from "@thoughtspot/rest-api-sdk"
 import YAML from "yaml";
 import type { Observable } from "rxjs";
 import { of } from "rxjs";
-import type { SessionInfo } from "./types";
+import type { SessionInfo, DataSourceSuggestionResponse } from "./types";
 
 export const getThoughtSpotClient = (instanceUrl: string, bearerToken: string) => {
     const config = createBearerAuthenticationConfig(
@@ -27,6 +27,7 @@ export const getThoughtSpotClient = (instanceUrl: string, bearerToken: string) =
     (client as any).instanceUrl = instanceUrl;
     addExportUnsavedAnswerTML(client, instanceUrl, bearerToken);
     addGetSessionInfo(client, instanceUrl, bearerToken);
+    addQueryGetDataSourceSuggestions(client, instanceUrl, bearerToken);
     return client;
 }
 
@@ -105,4 +106,50 @@ async function addGetSessionInfo(client: any, instanceUrl: string, token: string
         const info = data.info;
         return info;
     };
+}
+
+const getDataSourceSuggestionsQuery = `
+query QueryGetDataSourceSuggestions($request: Input_eureka_DataSourceSuggestionRequest) {
+  queryGetDataSourceSuggestions(request: $request) {
+    dataSources {
+      confidence
+      header {
+        description
+        displayName
+        guid
+      }
+      llmReasoning
+    }
+  }
+}`;
+
+// This is a workaround until we get the public API for this
+function addQueryGetDataSourceSuggestions(client: any, instanceUrl: string, token: string) {
+    (client as any).queryGetDataSourceSuggestions = async (query: string): Promise<DataSourceSuggestionResponse> => {
+        const endpoint = "/prism/";
+        const fetchOptions = {
+            method: "POST",
+            headers: {
+                "Content-Type": "application/json",
+                "Accept": "application/json",
+                "user-agent": "ThoughtSpot-ts-client",
+                "Authorization": `Bearer ${token}`,
+            },
+            body: JSON.stringify({
+                query: getDataSourceSuggestionsQuery,
+                variables: {
+                    request: {
+                        query: query
+                    }
+                },
+                operationName: "QueryGetDataSourceSuggestions"
+            })
+        };
+        console.log("fetchOptions", fetchOptions);
+        const response = await fetch(`${instanceUrl}${endpoint}`, fetchOptions);
+        console.log("response", response);
+        const data = await response.json() as any;
+        console.log("data", data);
+        return data.data.queryGetDataSourceSuggestions;
+    };
 }

src/thoughtspot/thoughtspot-service.ts

@@ -1,7 +1,7 @@
 import type { ThoughtSpotRestApi } from "@thoughtspot/rest-api-sdk";
 import { SpanStatusCode, trace, context } from "@opentelemetry/api";
 import { getActiveSpan, WithSpan } from "../metrics/tracing/tracing-utils";
-import type { DataSource, SessionInfo } from "./types";
+import type { DataSource, SessionInfo, DataSourceSuggestion } from "./types";
 
 
 /**
@@ -10,6 +10,56 @@ import type { DataSource, SessionInfo } from "./types";
 export class ThoughtSpotService {
     constructor(private client: ThoughtSpotRestApi) { }
 
+    @WithSpan('discover-data-sources')
+    async discoverDataSources(query?: string): Promise<DataSource[] | DataSourceSuggestion[] | null> {
+        const span = getActiveSpan();
+        span?.addEvent("discover-data-sources");
+        
+        // If a query is provided, use intelligent data source suggestions
+        if (query) {
+            span?.addEvent("get-data-source-suggestions");
+            return await this.getDataSourceSuggestions(query);
+        }
+        
+        // Otherwise, fallback to getting all data sources
+        const dataSources = await this.getDataSources();
+        return dataSources;
+    }
+
+    /**
+     * Get intelligent data source suggestions based on a query using GraphQL
+     */
+    @WithSpan('get-data-source-suggestions')
+    async getDataSourceSuggestions(query: string): Promise<DataSourceSuggestion[] | null> {
+        const span = getActiveSpan();
+        
+        try {
+            span?.setAttribute("query", query);
+            span?.addEvent("query-get-data-source-suggestions");
+                        
+            const response = await (this.client as any).queryGetDataSourceSuggestions(query);
+            
+            span?.setStatus({ code: SpanStatusCode.OK, message: "Data source suggestions retrieved" });
+            
+            // Check if we have any data sources
+            if (!response.dataSources || response.dataSources.length === 0) {
+                span?.setAttribute("suggestions_count", 0);
+                return null;
+            }
+            
+            span?.setAttribute("suggestions_count", response.dataSources.length);
+            
+            // Return top 2 data sources (or just 1 if only 1 available)
+            const topDataSources = response.dataSources.slice(0, 2);
+            return topDataSources;
+            
+        } catch (error) {
+            span?.setStatus({ code: SpanStatusCode.ERROR, message: (error as Error).message });
+            console.error("Error getting data source suggestions: ", error);
+            throw error;
+        }
+    }
+
     /**
      * Get relevant questions for a given query and data sources
      */
@@ -342,6 +392,7 @@ export class ThoughtSpotService {
             releaseVersion: info.releaseVersion,
             currentOrgId: info.currentOrgId,
             privileges: info.privileges,
+            enableSpotterDataSourceDiscovery: info.enableSpotterDataSourceDiscovery,
         };
     }
 
@@ -439,10 +490,18 @@ export async function getDataSources(
     return service.getDataSources();
 }
 
+export async function getDataSourceSuggestions(
+    query: string,
+    client: ThoughtSpotRestApi,
+): Promise<DataSourceSuggestion[] | null> {
+    const service = new ThoughtSpotService(client);
+    return service.getDataSourceSuggestions(query);
+}
+
 export async function getSessionInfo(client: ThoughtSpotRestApi): Promise<SessionInfo> {
     const service = new ThoughtSpotService(client);
     return service.getSessionInfo();
 }
 
 // Export types
-export type { DataSource, SessionInfo } from "./types";
+export type { DataSource, SessionInfo, DataSourceSuggestion, DataSourceSuggestionResponse } from "./types";

src/thoughtspot/types.ts

@@ -4,6 +4,20 @@ export interface DataSource {
     description: string;
 }
 
+export interface DataSourceSuggestion {
+    confidence: number;
+    header: {
+        description: string;
+        displayName: string;
+        guid: string;
+    };
+    llmReasoning: string;
+}
+
+export interface DataSourceSuggestionResponse {
+    dataSources: DataSourceSuggestion[];
+}
+
 export interface SessionInfo {
     mixpanelToken: string;
     userGUID: string;
@@ -13,4 +27,5 @@ export interface SessionInfo {
     releaseVersion: string;
     currentOrgId: string;
     privileges: any;
+    enableSpotterDataSourceDiscovery?: boolean;
 } 

src/utils.ts

@@ -90,6 +90,39 @@ export class McpServerError extends Error {
     }
 }
 
+/**
+ * Store a value in Cloudflare KV
+ */
+export async function putInKV(key: string, value: any, env?: any): Promise<void> {
+    if (!env?.OAUTH_KV) {
+        return;
+    }
+    try {
+        await env.OAUTH_KV.put(key, JSON.stringify(value), { expirationTtl: 60 * 60 * 3 });
+    } catch (error) {
+        console.error('Error storing in KV:', error);
+    }
+}
+
+/**
+ * Retrieve a value from Cloudflare KV
+ */
+export async function getFromKV(key: string, env?: any): Promise<any> {
+    console.log('[DEBUG] Getting from KV', key);
+    
+    if (!env?.OAUTH_KV) {
+        return undefined;
+    }
+    
+    try {
+        const value = await env.OAUTH_KV.get(key, { type: "json" });
+        return value;
+    } catch (error) {
+        console.error('Error retrieving from KV:', error);
+        return undefined;
+    }
+}
+
 export function instrumentedMCPServer<T extends BaseMCPServer>(MCPServer: new (ctx: Context) => T, config: ResolveConfigFn) {
     const Agent = class extends McpAgent<Env, any, Props> {
         server = new MCPServer(this);