feat(widgets): use new classifier, implement new widget executor, delete registry

This commit is contained in:
ItzCrazyKns
2025-12-02 11:52:40 +05:30
parent 6899b49ca0
commit 046f159528
10 changed files with 204 additions and 222 deletions

View File

@@ -1,26 +1,24 @@
import { ResearcherOutput, SearchAgentInput } from './types';
import SessionManager from '@/lib/session';
import Classifier from './classifier';
import { WidgetRegistry } from './widgets';
import { classify } from './classifier';
import Researcher from './researcher';
import { getWriterPrompt } from '@/lib/prompts/search/writer';
import fs from 'fs';
import { WidgetExecutor } from './widgets';
class SearchAgent {
async searchAsync(session: SessionManager, input: SearchAgentInput) {
const classifier = new Classifier();
const classification = await classifier.classify({
const classification = await classify({
chatHistory: input.chatHistory,
enabledSources: input.config.sources,
query: input.followUp,
llm: input.config.llm,
});
const widgetPromise = WidgetRegistry.executeAll(classification.widgets, {
const widgetPromise = WidgetExecutor.executeAll({
classification,
chatHistory: input.chatHistory,
followUp: input.followUp,
llm: input.config.llm,
embedding: input.config.embedding,
session: session,
}).then((widgetOutputs) => {
widgetOutputs.forEach((o) => {
session.emitBlock({
@@ -37,7 +35,7 @@ class SearchAgent {
let searchPromise: Promise<ResearcherOutput> | null = null;
if (!classification.skipSearch) {
if (!classification.classification.skipSearch) {
const researcher = new Researcher();
searchPromise = researcher.research(session, {
chatHistory: input.chatHistory,

View File

@@ -26,7 +26,7 @@ const webSearchAction: ResearchAction<typeof actionSchema> = {
name: 'web_search',
description: actionDescription,
schema: actionSchema,
enabled: (config) => config.classification.intents.includes('web_search'),
enabled: (config) => true,
execute: async (input, _) => {
let results: Chunk[] = [];

View File

@@ -61,9 +61,7 @@ class Researcher {
maxIteration,
);
const actionStream = input.config.llm.streamObject<
z.infer<typeof schema>
>({
const actionStream = input.config.llm.streamObject<typeof schema>({
messages: [
{
role: 'system',

View File

@@ -19,26 +19,17 @@ export type SearchAgentInput = {
config: SearchAgentConfig;
};
export interface Intent {
name: string;
description: string;
requiresSearch: boolean;
enabled: (config: { sources: SearchSources[] }) => boolean;
}
export type Widget<TSchema extends z.ZodObject<any> = z.ZodObject<any>> = {
name: string;
description: string;
schema: TSchema;
execute: (
params: z.infer<TSchema>,
additionalConfig: AdditionalConfig,
) => Promise<WidgetOutput>;
export type WidgetInput = {
chatHistory: ChatTurnMessage[];
followUp: string;
classification: ClassifierOutput;
llm: BaseLLM<any>;
};
export type WidgetConfig = {
export type Widget = {
type: string;
params: Record<string, any>;
shouldExecute: (classification: ClassifierOutput) => boolean;
execute: (input: WidgetInput) => Promise<WidgetOutput | void>;
};
export type WidgetOutput = {

View File

@@ -1,66 +1,66 @@
import z from 'zod';
import { Widget } from '../types';
import { evaluate as mathEval } from 'mathjs';
import formatChatHistoryAsString from '@/lib/utils/formatHistory';
import { exp, evaluate as mathEval } from 'mathjs';
const schema = z.object({
type: z.literal('calculation'),
expression: z
.string()
.describe(
"A valid mathematical expression to be evaluated (e.g., '2 + 2', '3 * (4 + 5)').",
),
.describe('Mathematical expression to calculate or evaluate.'),
notPresent: z
.boolean()
.describe('Whether there is any need for the calculation widget.'),
});
const calculationWidget: Widget<typeof schema> = {
name: 'calculation',
description: `Performs mathematical calculations and evaluates mathematical expressions. Supports arithmetic operations, algebraic equations, functions, and complex mathematical computations.
const system = `
<role>
Assistant is a calculation expression extractor. You will recieve a user follow up and a conversation history.
Your task is to determine if there is a mathematical expression that needs to be calculated or evaluated. If there is, extract the expression and return it. If there is no need for any calculation, set notPresent to true.
</role>
**What it provides:**
- Evaluates mathematical expressions and returns computed results
- Handles basic arithmetic (+, -, *, /)
- Supports functions (sqrt, sin, cos, log, etc.)
- Can process complex expressions with parentheses and order of operations
<instructions>
Make sure that the extracted expression is valid and can be used to calculate the result with Math JS library (https://mathjs.org/). If the expression is not valid, set notPresent to true.
If you feel like you cannot extract a valid expression, set notPresent to true.
</instructions>
**When to use:**
- User asks to calculate, compute, or evaluate a mathematical expression
- Questions like "what is X", "calculate Y", "how much is Z" where X/Y/Z are math expressions
- Any request involving numbers and mathematical operations
**Example call:**
<output_format>
You must respond in the following JSON format without any extra text, explanations or filler sentences:
{
"type": "calculation",
"expression": "25% of 480"
"expression": string,
"notPresent": boolean
}
</output_format>
`;
{
"type": "calculation",
"expression": "sqrt(144) + 5 * 2"
}
const calculationWidget: Widget = {
type: 'calculationWidget',
shouldExecute: (classification) =>
classification.classification.showCalculationWidget,
execute: async (input) => {
const output = await input.llm.generateObject<typeof schema>({
messages: [
{
role: 'system',
content: system,
},
{
role: 'user',
content: `<conversation_history>\n${formatChatHistoryAsString(input.chatHistory)}\n</conversation_history>\n<user_follow_up>\n${input.followUp}\n</user_follow_up>`,
},
],
schema,
});
**Important:** The expression must be valid mathematical syntax that can be evaluated by mathjs. Format percentages as "0.25 * 480" or "25% of 480". Do not include currency symbols, units, or non-mathematical text in the expression.`,
schema: schema,
execute: async (params, _) => {
try {
const result = mathEval(params.expression);
const result = mathEval(output.expression);
return {
type: 'calculation_result',
llmContext: `The result of the expression "${params.expression}" is ${result}.`,
llmContext: `The result of the calculation for the expression "${output.expression}" is: ${result}`,
data: {
expression: params.expression,
result: result,
expression: output.expression,
result,
},
};
} catch (error) {
return {
type: 'calculation_result',
llmContext: 'Failed to evaluate mathematical expression.',
data: {
expression: params.expression,
result: `Error evaluating expression: ${error}`,
},
};
}
},
};

View File

@@ -0,0 +1,36 @@
import { Widget, WidgetInput, WidgetOutput } from '../types';
class WidgetExecutor {
static widgets = new Map<string, Widget>();
static register(widget: Widget) {
this.widgets.set(widget.type, widget);
}
static getWidget(type: string): Widget | undefined {
return this.widgets.get(type);
}
static async executeAll(input: WidgetInput): Promise<WidgetOutput[]> {
const results: WidgetOutput[] = [];
await Promise.all(
Array.from(this.widgets.values()).map(async (widget) => {
try {
if (widget.shouldExecute(input.classification)) {
const output = await widget.execute(input);
if (output) {
results.push(output);
}
}
} catch (e) {
console.log(`Error executing widget ${widget.type}:`, e);
}
}),
);
return results;
}
}
export default WidgetExecutor;

View File

@@ -1,10 +1,10 @@
import calculationWidget from './calculationWidget';
import WidgetRegistry from './registry';
import WidgetExecutor from './executor';
import weatherWidget from './weatherWidget';
import stockWidget from './stockWidget';
WidgetRegistry.register(weatherWidget);
WidgetRegistry.register(calculationWidget);
WidgetRegistry.register(stockWidget);
WidgetExecutor.register(weatherWidget);
WidgetExecutor.register(calculationWidget);
WidgetExecutor.register(stockWidget);
export { WidgetRegistry };
export { WidgetExecutor };

View File

@@ -1,65 +0,0 @@
import {
AdditionalConfig,
SearchAgentConfig,
Widget,
WidgetConfig,
WidgetOutput,
} from '../types';
class WidgetRegistry {
private static widgets = new Map<string, Widget>();
static register(widget: Widget<any>) {
this.widgets.set(widget.name, widget);
}
static get(name: string): Widget | undefined {
return this.widgets.get(name);
}
static getAll(): Widget[] {
return Array.from(this.widgets.values());
}
static getDescriptions(): string {
return Array.from(this.widgets.values())
.map((widget) => `${widget.name}: ${widget.description}`)
.join('\n\n');
}
static async execute(
name: string,
params: any,
config: AdditionalConfig,
): Promise<WidgetOutput> {
const widget = this.get(name);
if (!widget) {
throw new Error(`Widget with name ${name} not found`);
}
return widget.execute(params, config);
}
static async executeAll(
widgets: WidgetConfig[],
additionalConfig: AdditionalConfig,
): Promise<WidgetOutput[]> {
const results: WidgetOutput[] = [];
await Promise.all(
widgets.map(async (widgetConfig) => {
const output = await this.execute(
widgetConfig.type,
widgetConfig.params,
additionalConfig,
);
results.push(output);
}),
);
return results;
}
}
export default WidgetRegistry;

View File

@@ -1,13 +1,13 @@
import z from 'zod';
import { Widget } from '../types';
import YahooFinance from 'yahoo-finance2';
import formatChatHistoryAsString from '@/lib/utils/formatHistory';
const yf = new YahooFinance({
suppressNotices: ['yahooSurvey'],
});
const schema = z.object({
type: z.literal('stock'),
name: z
.string()
.describe(
@@ -19,60 +19,59 @@ const schema = z.object({
.describe(
"Optional array of up to 3 stock names to compare against the base name (e.g., ['Microsoft', 'GOOGL', 'Meta']). Charts will show percentage change comparison.",
),
notPresent: z
.boolean()
.describe('Whether there is no need for the stock widget.'),
});
const stockWidget: Widget<typeof schema> = {
name: 'stock',
description: `Provides comprehensive real-time stock market data and financial information for any publicly traded company. Returns detailed quote data, market status, trading metrics, and company fundamentals.
const systemPrompt = `
<role>
You are a stock ticker/name extractor. You will receive a user follow up and a conversation history.
Your task is to determine if the user is asking about stock information and extract the stock name(s) they want data for.
</role>
You can set skipSearch to true if the stock widget can fully answer the user's query without needing additional web search.
<instructions>
- If the user is asking about a stock, extract the primary stock name or ticker.
- If the user wants to compare stocks, extract up to 3 comparison stock names in comparisonNames.
- You can use either stock names (e.g., "Nvidia", "Apple") or tickers (e.g., "NVDA", "AAPL").
- If you cannot determine a valid stock or the query is not stock-related, set notPresent to true.
- If no comparison is needed, set comparisonNames to an empty array.
</instructions>
**What it provides:**
- **Real-time Price Data**: Current price, previous close, open price, day's range (high/low)
- **Market Status**: Whether market is currently open or closed, trading sessions
- **Trading Metrics**: Volume, average volume, bid/ask prices and sizes
- **Performance**: Price changes (absolute and percentage), 52-week high/low range
- **Valuation**: Market capitalization, P/E ratio, earnings per share (EPS)
- **Dividends**: Dividend rate, dividend yield, ex-dividend date
- **Company Info**: Full company name, exchange, currency, sector/industry (when available)
- **Advanced Metrics**: Beta, trailing/forward P/E, book value, price-to-book ratio
- **Charts Data**: Historical price movements for visualization
- **Comparison**: Compare up to 3 stocks side-by-side with percentage-based performance visualization
**When to use:**
- User asks about a stock price ("What's AAPL stock price?", "How is Tesla doing?")
- Questions about company market performance ("Is Microsoft up or down today?")
- Requests for stock market data, trading info, or company valuation
- Queries about dividends, P/E ratio, market cap, or other financial metrics
- Any stock/equity-related question for a specific company
- Stock comparisons ("Compare AAPL vs MSFT", "How is TSLA doing vs RIVN and LCID?")
**Example calls:**
<output_format>
You must respond in the following JSON format without any extra text, explanations or filler sentences:
{
"type": "stock",
"name": "AAPL"
"name": string,
"comparisonNames": string[],
"notPresent": boolean
}
</output_format>
`;
{
"type": "stock",
"name": "TSLA",
"comparisonNames": ["RIVN", "LCID"]
}
const stockWidget: Widget = {
type: 'stockWidget',
shouldExecute: (classification) =>
classification.classification.showStockWidget,
execute: async (input) => {
const output = await input.llm.generateObject<typeof schema>({
messages: [
{
role: 'system',
content: systemPrompt,
},
{
role: 'user',
content: `<conversation_history>\n${formatChatHistoryAsString(input.chatHistory)}\n</conversation_history>\n<user_follow_up>\n${input.followUp}\n</user_follow_up>`,
},
],
schema,
});
{
"type": "stock",
"name": "Google",
"comparisonNames": ["Microsoft", "Meta", "Amazon"]
}
if (output.notPresent) {
return;
}
**Important:**
- You can use both tickers and names (prefer name when you're not aware of the ticker).
- For companies with multiple share classes, use the most common one.
- The widget works for stocks listed on major exchanges (NYSE, NASDAQ, etc.)
- Returns comprehensive data; the UI will display relevant metrics based on availability
- Market data may be delayed by 15-20 minutes for free data sources during trading hours`,
schema: schema,
execute: async (params, _) => {
const params = output;
try {
const name = params.name;

View File

@@ -1,8 +1,8 @@
import z from 'zod';
import { Widget } from '../types';
import formatChatHistoryAsString from '@/lib/utils/formatHistory';
const WeatherWidgetSchema = z.object({
type: z.literal('weather'),
const schema = z.object({
location: z
.string()
.describe(
@@ -18,38 +18,63 @@ const WeatherWidgetSchema = z.object({
.describe(
'Longitude coordinate in decimal degrees (e.g., -74.0060). Only use when location name is empty.',
),
notPresent: z
.boolean()
.describe('Whether there is no need for the weather widget.'),
});
const weatherWidget: Widget<typeof WeatherWidgetSchema> = {
name: 'weather',
description: `Provides comprehensive current weather information and forecasts for any location worldwide. Returns real-time weather data including temperature, conditions, humidity, wind, and multi-day forecasts.
const systemPrompt = `
<role>
You are a location extractor for weather queries. You will receive a user follow up and a conversation history.
Your task is to determine if the user is asking about weather and extract the location they want weather for.
</role>
You can set skipSearch to true if the weather widget can fully answer the user's query without needing additional web search.
<instructions>
- If the user is asking about weather, extract the location name OR coordinates (never both).
- If using location name, set lat and lon to 0.
- If using coordinates, set location to empty string.
- If you cannot determine a valid location or the query is not weather-related, set notPresent to true.
- Location should be specific (city, state/region, country) for best results.
- You have to give the location so that it can be used to fetch weather data, it cannot be left empty unless notPresent is true.
- Make sure to infer short forms of location names (e.g., "NYC" -> "New York City", "LA" -> "Los Angeles").
</instructions>
**What it provides:**
- Current weather conditions (temperature, feels-like, humidity, precipitation)
- Wind speed, direction, and gusts
- Weather codes/conditions (clear, cloudy, rainy, etc.)
- Hourly forecast for next 24 hours
- Daily forecast for next 7 days (high/low temps, precipitation probability)
- Timezone information
**When to use:**
- User asks about weather in a location ("weather in X", "is it raining in Y")
- Questions about temperature, conditions, or forecast
- Any weather-related query for a specific place
**Example call:**
<output_format>
You must respond in the following JSON format without any extra text, explanations or filler sentences:
{
"type": "weather",
"location": "San Francisco, CA, USA",
"lat": 0,
"lon": 0
"location": string,
"lat": number,
"lon": number,
"notPresent": boolean
}
</output_format>
`;
const weatherWidget: Widget = {
type: 'weatherWidget',
shouldExecute: (classification) =>
classification.classification.showWeatherWidget,
execute: async (input) => {
const output = await input.llm.generateObject<typeof schema>({
messages: [
{
role: 'system',
content: systemPrompt,
},
{
role: 'user',
content: `<conversation_history>\n${formatChatHistoryAsString(input.chatHistory)}\n</conversation_history>\n<user_follow_up>\n${input.followUp}\n</user_follow_up>`,
},
],
schema,
});
if (output.notPresent) {
return;
}
const params = output;
**Important:** Provide EITHER a location name OR latitude/longitude coordinates, never both. If using location name, set lat/lon to 0. Location should be specific (city, state/region, country) for best results.`,
schema: WeatherWidgetSchema,
execute: async (params, _) => {
try {
if (
params.location === '' &&