feat: insights chat

frontend/lib/chat/agents/base-agent.ts

@@ -0,0 +1,167 @@
+// Copyright (c) 2025 The Linux Foundation and each contributor.
+// SPDX-License-Identifier: MIT
+import { z } from "zod";
+import { generateText } from "ai";
+
+export abstract class BaseAgent<TInput, TOutput> {
+  abstract readonly name: string;
+  abstract readonly outputSchema: z.ZodSchema<TOutput>;
+  abstract readonly temperature: number;
+  abstract readonly maxSteps: number;
+
+  /**
+   * Generates JSON format instructions based on the Zod schema
+   */
+  protected generateJSONInstructions(): string {
+    try {
+      const schemaShape = this.getSchemaShape(this.outputSchema);
+      return `\n\nReturn your response as a JSON object with this exact structure:\n${JSON.stringify(
+        schemaShape,
+        null,
+        2
+      )}\n\nDo not include any other text in your response, only the JSON object.`;
+    } catch (error) {
+      return "\n\nReturn your response as a valid JSON object.";
+    }
+  }
+
+  /**
+   * Extracts the shape of a Zod schema for documentation
+   */
+  private getSchemaShape(schema: z.ZodSchema<any>): any {
+    if (schema instanceof z.ZodObject) {
+      const shape: any = {};
+      const schemaShape = (schema as any).shape;
+
+      for (const key in schemaShape) {
+        const field = schemaShape[key];
+        if (field instanceof z.ZodString) {
+          shape[key] = field.description || "string";
+        } else if (field instanceof z.ZodNumber) {
+          shape[key] = field.description || "number";
+        } else if (field instanceof z.ZodBoolean) {
+          shape[key] = field.description || "boolean";
+        } else if (field instanceof z.ZodArray) {
+          const elementType = this.getFieldDescription(field._def.type);
+          shape[key] = field.description || `array of ${elementType}`;
+        } else if (field instanceof z.ZodEnum) {
+          const values = (field as any)._def.values;
+          shape[key] = field.description || `one of: ${values.join(" | ")}`;
+        } else if (field instanceof z.ZodOptional || field instanceof z.ZodNullable) {
+          const innerType = (field as any)._def.innerType;
+          shape[key] = `${this.getFieldDescription(innerType)} (optional)`;
+        } else if (field instanceof z.ZodAny) {
+          shape[key] = field.description || "any value";
+        } else {
+          shape[key] = field.description || "value";
+        }
+      }
+
+      return shape;
+    }
+
+    return {};
+  }
+
+  private getFieldDescription(field: any): string {
+    if (field instanceof z.ZodString) return "string";
+    if (field instanceof z.ZodNumber) return "number";
+    if (field instanceof z.ZodBoolean) return "boolean";
+    if (field instanceof z.ZodAny) return "any";
+    return "value";
+  }
+
+  async execute(input: TInput): Promise<TOutput> {
+    try {
+      const systemPrompt = await this.getSystemPrompt(input);
+      const userPrompt = this.getUserPrompt(input);
+      const tools = this.getTools(input);
+
+      // Append JSON format instructions to system prompt
+      const jsonInstructions = this.generateJSONInstructions();
+      const fullSystemPrompt = systemPrompt + jsonInstructions;
+
+      // Check if we have messages in the input
+      const hasMessages = typeof input === 'object' && input !== null && 'messages' in input && Array.isArray((input as any).messages);
+
+      const generateConfig: any = {
+        model: this.getModel(input),
+        system: fullSystemPrompt,
+        tools,
+        maxSteps: this.maxSteps,
+        temperature: this.temperature,
+      };
+
+      // Add any provider-specific options
+      const providerOptions = this.getProviderOptions(input);
+      if (providerOptions) {
+        generateConfig.providerOptions = providerOptions;
+      }
+
+      // Use messages if available, otherwise use prompt
+      if (hasMessages) {
+        generateConfig.messages = (input as any).messages;
+      } else {
+        generateConfig.prompt = userPrompt;
+      }
+
+
+      const response = await generateText(generateConfig);
+
+
+      // Extract and validate JSON from response
+      return this.extractJSON(response.text);
+    } catch (error) {
+      throw this.createError(error);
+    }
+  }
+
+  /**
+   * Extract and validate JSON from the response text
+   */
+  protected extractJSON(text: string): TOutput {
+    // Remove markdown code block wrappers if present
+    let cleanText = text.replaceAll(/```json\s*\n?/g, "");
+    cleanText = cleanText.replaceAll(/\n?```/g, "");
+
+    // Try to find JSON in the response
+    const jsonMatch = cleanText.match(/\{[\s\S]*\}/);
+    if (!jsonMatch) {
+      console.error("No JSON found in the response");
+      console.error(text);
+      throw new Error(`${this.name} agent did not return valid JSON`);
+    }
+
+    const jsonStr = jsonMatch[0];
+    let parsedOutput: any;
+
+    try {
+      parsedOutput = JSON.parse(jsonStr);
+    } catch (error) {
+      throw new Error(`Failed to parse ${this.name} JSON: ${error}`);
+    }
+
+    // Validate against schema
+    try {
+      const validatedOutput = this.outputSchema.parse(parsedOutput);
+      console.log(`${this.name} Agent Output:`, validatedOutput);
+      return validatedOutput;
+    } catch (error) {
+      console.error(`Failed to validate ${this.name} JSON`, error);
+      throw new Error(`Failed to validate ${this.name} JSON: ${error}`);
+    }
+  }
+
+  protected abstract getModel(input: TInput): any;
+  protected abstract getSystemPrompt(input: TInput): string | Promise<string>;
+  protected abstract getUserPrompt(input: TInput): string;
+  protected abstract getTools(input: TInput): Record<string, any>;
+  protected abstract createError(error: unknown): Error;
+
+  /**
+   * Override this method to provide provider-specific options
+   */
+  protected getProviderOptions(_input: TInput): any {
+    return undefined;
+  }
+}

frontend/lib/chat/agents/pipe.ts

@@ -0,0 +1,72 @@
+// Copyright (c) 2025 The Linux Foundation and each contributor.
+// SPDX-License-Identifier: MIT
+import { z } from 'zod';
+import { pipePrompt } from '../prompts/pipe';
+import { BaseAgent } from './base-agent';
+
+// Output schema for pipe agent
+export const pipeOutputSchema = z.object({
+  tools: z.array(z.string()).describe("Ordered array of tools used to fetch the data"),
+  explanation: z.string().describe("Brief explanation of why these tools answer the question"),
+  data: z.any().describe("The actual data returned from executing the tools")
+});
+
+export type PipeOutput = z.infer<typeof pipeOutputSchema>;
+
+interface PipeAgentInput {
+  model: any; // Bedrock model instance
+  messages: any[];
+  tools: Record<string, any>; // Filtered pipe tools based on router decision
+  date: string;
+  projectName: string;
+  pipe: string;
+  parametersString: string;
+  segmentId: string | null;
+  reformulatedQuestion: string;
+  toolNames: string[]; // Array of tool names from router
+}
+
+export class PipeAgent extends BaseAgent<PipeAgentInput, PipeOutput> {
+  readonly name = "Pipe";
+  readonly outputSchema = pipeOutputSchema;
+  readonly temperature = 0;
+  readonly maxSteps = 10;
+
+  protected getModel(input: PipeAgentInput): any {
+    return input.model;
+  }
+
+  protected getSystemPrompt(input: PipeAgentInput): string {
+    return pipePrompt(
+      input.date,
+      input.projectName,
+      input.pipe,
+      input.parametersString,
+      input.segmentId,
+      input.reformulatedQuestion,
+      input.toolNames
+    );
+  }
+
+  protected getUserPrompt(_input: PipeAgentInput): string {
+    // Not used when messages are provided, but required by base class
+    return _input.reformulatedQuestion;
+  }
+
+  protected getTools(input: PipeAgentInput): Record<string, any> {
+    return input.tools;
+  }
+
+  protected createError(error: unknown): Error {
+    if (error instanceof Error) {
+      return new Error(`Pipe agent error: ${error.message}`);
+    }
+    return new Error(`Pipe agent error: ${String(error)}`);
+  }
+}
+
+// Convenience function to maintain backward compatibility
+export async function runPipeAgent(params: PipeAgentInput): Promise<PipeOutput> {
+  const agent = new PipeAgent();
+  return agent.execute(params);
+}

frontend/lib/chat/agents/router.ts

@@ -0,0 +1,70 @@
+// Copyright (c) 2025 The Linux Foundation and each contributor.
+// SPDX-License-Identifier: MIT
+import { z } from 'zod';
+import { routerPrompt } from '../prompts/router';
+import { BaseAgent } from './base-agent';
+
+// Output schema for router decisions
+export const routerOutputSchema = z.object({
+  next_action: z.enum(["stop", "create_query", "pipes"]),
+  reasoning: z.string().describe("Maximum 2 sentences explaining the decision"),
+  reformulated_question: z.string().describe("Enhanced query with all parameters"),
+  tools: z.array(z.string()).describe("Tools needed for next agent")
+});
+
+export type RouterOutput = z.infer<typeof routerOutputSchema>;
+
+interface RouterAgentInput {
+  model: any; // Bedrock model instance
+  messages: any[];
+  tools: Record<string, any>; // Only list_datasources
+  date: string;
+  projectName: string;
+  pipe: string;
+  parametersString: string;
+  segmentId: string | null;
+}
+
+export class RouterAgent extends BaseAgent<RouterAgentInput, RouterOutput> {
+  readonly name = "Router";
+  readonly outputSchema = routerOutputSchema;
+  readonly temperature = 0;
+  readonly maxSteps = 3; // Allow router to use list_datasources if needed
+
+  protected getModel(input: RouterAgentInput): any {
+    return input.model;
+  }
+
+  protected getSystemPrompt(input: RouterAgentInput): string {
+    return routerPrompt(
+      input.date,
+      input.projectName,
+      input.pipe,
+      input.parametersString,
+      input.segmentId
+    );
+  }
+
+  protected getUserPrompt(_input: RouterAgentInput): string {
+    // Not used when messages are provided, but required by base class
+    return "";
+  }
+
+  protected getTools(input: RouterAgentInput): Record<string, any> {
+    return input.tools;
+  }
+
+  protected createError(error: unknown): Error {
+    if (error instanceof Error) {
+      return new Error(`Router agent error: ${error.message}`);
+    }
+    return new Error(`Router agent error: ${String(error)}`);
+  }
+
+}
+
+// Convenience function to maintain backward compatibility
+export async function runRouterAgent(params: RouterAgentInput): Promise<RouterOutput> {
+  const agent = new RouterAgent();
+  return agent.execute(params);
+}

frontend/lib/chat/agents/text-to-sql.ts

@@ -0,0 +1,78 @@
+// Copyright (c) 2025 The Linux Foundation and each contributor.
+// SPDX-License-Identifier: MIT
+import { z } from 'zod';
+import { textToSqlPrompt } from '../prompts/text-to-sql';
+import { BaseAgent } from './base-agent';
+
+// Output schema for SQL agent
+export const sqlOutputSchema = z.object({
+  sql: z.string().describe("The SQL query used to fetch the data"),
+  explanation: z.string().describe("Brief explanation of why this query answers the question"),
+  data: z.any().describe("The actual data returned from executing the query")
+});
+
+export type SqlOutput = z.infer<typeof sqlOutputSchema>;
+
+interface TextToSqlAgentInput {
+  model: any; // Bedrock model instance
+  messages: any[];
+  tools: Record<string, any>; // list_datasources and execute_query
+  date: string;
+  projectName: string;
+  pipe: string;
+  parametersString: string;
+  segmentId: string | null;
+  reformulatedQuestion: string;
+}
+
+export class TextToSqlAgent extends BaseAgent<TextToSqlAgentInput, SqlOutput> {
+  readonly name = "SQL";
+  readonly outputSchema = sqlOutputSchema;
+  readonly temperature = 0;
+  readonly maxSteps = 10; // Allow multiple steps for SQL generation and execution
+
+  protected getModel(input: TextToSqlAgentInput): any {
+    return input.model;
+  }
+
+  protected getSystemPrompt(input: TextToSqlAgentInput): string {
+    return textToSqlPrompt(
+      input.date,
+      input.projectName,
+      input.pipe,
+      input.parametersString,
+      input.segmentId,
+      input.reformulatedQuestion
+    );
+  }
+
+  protected getUserPrompt(input: TextToSqlAgentInput): string {
+    // Not used when messages are provided, but required by base class
+    return input.reformulatedQuestion;
+  }
+
+  protected getTools(input: TextToSqlAgentInput): Record<string, any> {
+    return input.tools;
+  }
+
+  protected createError(error: unknown): Error {
+    if (error instanceof Error) {
+      return new Error(`SQL agent error: ${error.message}`);
+    }
+    return new Error(`SQL agent error: ${String(error)}`);
+  }
+
+  protected getProviderOptions(_input: TextToSqlAgentInput): any {
+    return {
+      bedrock: {
+        reasoningConfig: { type: "enabled", budgetTokens: 3000 },
+      },
+    };
+  }
+}
+
+// Convenience function to maintain backward compatibility
+export async function runTextToSqlAgent(params: TextToSqlAgentInput): Promise<SqlOutput> {
+  const agent = new TextToSqlAgent();
+  return agent.execute(params);
+}