diff --git a/src/agent/grok-agent.ts b/src/agent/grok-agent.ts index 673cc20..bf6a983 100644 --- a/src/agent/grok-agent.ts +++ b/src/agent/grok-agent.ts @@ -172,6 +172,34 @@ Current working directory: ${process.cwd()}`, return currentModel.toLowerCase().includes("grok"); } + // Heuristic: enable web search only when likely needed + private shouldUseSearchFor(message: string): boolean { + const q = message.toLowerCase(); + const keywords = [ + "today", + "latest", + "news", + "trending", + "breaking", + "current", + "now", + "recent", + "x.com", + "twitter", + "tweet", + "what happened", + "as of", + "update on", + "release notes", + "changelog", + "price", + ]; + if (keywords.some((k) => q.includes(k))) return true; + // crude date pattern (e.g., 2024/2025) may imply recency + if (/(20\d{2})/.test(q)) return true; + return false; + } + async processUserMessage(message: string): Promise { // Add user message to conversation const userEntry: ChatEntry = { @@ -192,7 +220,9 @@ Current working directory: ${process.cwd()}`, this.messages, tools, undefined, - this.isGrokModel() ? { search_parameters: { mode: "auto" } } : undefined + this.isGrokModel() && this.shouldUseSearchFor(message) + ? { search_parameters: { mode: "auto" } } + : { search_parameters: { mode: "off" } } ); // Agent loop - continue until no more tool calls or max rounds reached @@ -286,9 +316,9 @@ Current working directory: ${process.cwd()}`, this.messages, tools, undefined, - this.isGrokModel() + this.isGrokModel() && this.shouldUseSearchFor(message) ? { search_parameters: { mode: "auto" } } - : undefined + : { search_parameters: { mode: "off" } } ); } else { // No more tool calls, add final response @@ -389,6 +419,7 @@ Current working directory: ${process.cwd()}`, const maxToolRounds = this.maxToolRounds; // Prevent infinite loops let toolRounds = 0; let totalOutputTokens = 0; + let lastTokenUpdate = 0; try { // Agent loop - continue until no more tool calls or max rounds reached @@ -409,9 +440,9 @@ Current working directory: ${process.cwd()}`, this.messages, tools, undefined, - this.isGrokModel() + this.isGrokModel() && this.shouldUseSearchFor(message) ? { search_parameters: { mode: "auto" } } - : undefined + : { search_parameters: { mode: "off" } } ); let accumulatedMessage: any = {}; let accumulatedContent = ""; @@ -468,12 +499,16 @@ Current working directory: ${process.cwd()}`, }; // Emit token count update - yield { - type: "token_count", - tokenCount: inputTokens + totalOutputTokens, - }; - } + const now = Date.now(); + if (now - lastTokenUpdate > 250) { + lastTokenUpdate = now; + yield { + type: "token_count", + tokenCount: inputTokens + totalOutputTokens, + }; + } } + } // Add assistant entry to history const assistantEntry: ChatEntry = { @@ -548,6 +583,7 @@ Current working directory: ${process.cwd()}`, inputTokens = this.tokenCounter.countMessageTokens( this.messages as any ); + // Final token update after tools processed yield { type: "token_count", tokenCount: inputTokens + totalOutputTokens, diff --git a/src/grok/client.ts b/src/grok/client.ts index 06e2813..bf10180 100644 --- a/src/grok/client.ts +++ b/src/grok/client.ts @@ -48,6 +48,7 @@ export interface GrokResponse { export class GrokClient { private client: OpenAI; private currentModel: string = "grok-code-fast-1"; + private defaultMaxTokens: number; constructor(apiKey: string, model?: string, baseURL?: string) { this.client = new OpenAI({ @@ -55,6 +56,8 @@ export class GrokClient { baseURL: baseURL || process.env.GROK_BASE_URL || "https://api.x.ai/v1", timeout: 360000, }); + const envMax = Number(process.env.GROK_MAX_TOKENS); + this.defaultMaxTokens = Number.isFinite(envMax) && envMax > 0 ? envMax : 1536; if (model) { this.currentModel = model; } @@ -81,7 +84,7 @@ export class GrokClient { tools: tools || [], tool_choice: tools && tools.length > 0 ? "auto" : undefined, temperature: 0.7, - max_tokens: 4000, + max_tokens: this.defaultMaxTokens, }; // Add search parameters if specified @@ -111,7 +114,7 @@ export class GrokClient { tools: tools || [], tool_choice: tools && tools.length > 0 ? "auto" : undefined, temperature: 0.7, - max_tokens: 4000, + max_tokens: this.defaultMaxTokens, stream: true, }; diff --git a/src/ui/components/chat-interface.tsx b/src/ui/components/chat-interface.tsx index e075886..abf841d 100644 --- a/src/ui/components/chat-interface.tsx +++ b/src/ui/components/chat-interface.tsx @@ -108,10 +108,9 @@ function ChatInterfaceWithAgent({ setChatHistory([]); }, []); - // Process initial message if provided + // Process initial message if provided (streaming for faster feedback) useEffect(() => { if (initialMessage && agent) { - // First, immediately add the user message to chat history const userEntry: ChatEntry = { type: "user", content: initialMessage, @@ -119,12 +118,119 @@ function ChatInterfaceWithAgent({ }; setChatHistory([userEntry]); - // Then process the message asynchronously const processInitialMessage = async () => { setIsProcessing(true); - const entries = await agent.processUserMessage(initialMessage); - setChatHistory(entries); + setIsStreaming(true); + + try { + let streamingEntry: ChatEntry | null = null; + for await (const chunk of agent.processUserMessageStream(initialMessage)) { + switch (chunk.type) { + case "content": + if (chunk.content) { + if (!streamingEntry) { + const newStreamingEntry = { + type: "assistant" as const, + content: chunk.content, + timestamp: new Date(), + isStreaming: true, + }; + setChatHistory((prev) => [...prev, newStreamingEntry]); + streamingEntry = newStreamingEntry; + } else { + setChatHistory((prev) => + prev.map((entry, idx) => + idx === prev.length - 1 && entry.isStreaming + ? { ...entry, content: entry.content + chunk.content } + : entry + ) + ); + } + } + break; + case "token_count": + if (chunk.tokenCount !== undefined) { + setTokenCount(chunk.tokenCount); + } + break; + case "tool_calls": + if (chunk.toolCalls) { + // Stop streaming for the current assistant message + setChatHistory((prev) => + prev.map((entry) => + entry.isStreaming + ? { + ...entry, + isStreaming: false, + toolCalls: chunk.toolCalls, + } + : entry + ) + ); + streamingEntry = null; + + // Add individual tool call entries to show tools are being executed + chunk.toolCalls.forEach((toolCall) => { + const toolCallEntry: ChatEntry = { + type: "tool_call", + content: "Executing...", + timestamp: new Date(), + toolCall: toolCall, + }; + setChatHistory((prev) => [...prev, toolCallEntry]); + }); + } + break; + case "tool_result": + if (chunk.toolCall && chunk.toolResult) { + setChatHistory((prev) => + prev.map((entry) => { + if (entry.isStreaming) { + return { ...entry, isStreaming: false }; + } + if ( + entry.type === "tool_call" && + entry.toolCall?.id === chunk.toolCall?.id + ) { + return { + ...entry, + type: "tool_result", + content: chunk.toolResult.success + ? chunk.toolResult.output || "Success" + : chunk.toolResult.error || "Error occurred", + toolResult: chunk.toolResult, + }; + } + return entry; + }) + ); + streamingEntry = null; + } + break; + case "done": + if (streamingEntry) { + setChatHistory((prev) => + prev.map((entry) => + entry.isStreaming ? { ...entry, isStreaming: false } : entry + ) + ); + } + setIsStreaming(false); + break; + } + } + } catch (error: any) { + const errorEntry: ChatEntry = { + type: "assistant", + content: `Error: ${error.message}`, + timestamp: new Date(), + }; + setChatHistory((prev) => [...prev, errorEntry]); + setIsStreaming(false); + } + setIsProcessing(false); + processingStartTime.current = 0; }; processInitialMessage();