Skip to content

Implement Search Mode Selector and Agent Personas #328

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Implement Search Mode Selector and Agent Personas #328

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion app/actions.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,8 @@ async function submit(formData?: FormData, skip?: boolean) {
: ((formData?.get('related_query') as string) ||
(formData?.get('input') as string))

const searchMode = formData?.get('searchMode') as string || 'Standard'

Comment on lines +136 to +137
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You’re trusting the client-provided searchMode string without validation. This allows unexpected values to flow into your agent logic and could lead to inconsistent prompts/tool availability. Validate and normalize it to a known set of modes before use.

Suggestion

Validate the value before using it:

const rawSearchMode = (formData?.get('searchMode') as string) ?? 'Standard'
const allowedModes = new Set(['Standard', 'Geospatial', 'Web Search'])
const searchMode = allowedModes.has(rawSearchMode) ? rawSearchMode : 'Standard'

Reply with "@CharlieHelps yes please" if you'd like me to add a commit with this suggestion.

if (userInput.toLowerCase().trim() === 'what is a planet computer?' || userInput.toLowerCase().trim() === 'what is qcx-terra?') {
const definition = userInput.toLowerCase().trim() === 'what is a planet computer?'
? `A planet computer is a proprietary environment aware system that interoperates weather forecasting, mapping and scheduling using cutting edge multi-agents to streamline automation and exploration on a planet. Available for our Pro and Enterprise customers. [QCX Pricing](https://www.queue.cx/#pricing)`
Expand Down Expand Up @@ -324,7 +326,8 @@ async function submit(formData?: FormData, skip?: boolean) {
uiStream,
streamText,
messages,
useSpecificAPI
useSpecificAPI,
searchMode
)
answer = fullResponse
toolOutputs = toolResponses
Expand Down
45 changes: 42 additions & 3 deletions components/chat-panel.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,13 @@
import { useEffect, useState, useRef, ChangeEvent, forwardRef, useImperativeHandle } from 'react'
import type { AI, UIState } from '@/app/actions'
import { useUIState, useActions } from 'ai/rsc'
// Removed import of useGeospatialToolMcp as it's no longer used/available
import { cn } from '@/lib/utils'
import { UserMessage } from './user-message'
import { Button } from './ui/button'
import { ArrowRight, Plus, Paperclip, X } from 'lucide-react'
import { ArrowRight, Plus, Paperclip, X, Search, MapPin, Globe } from 'lucide-react'
import Textarea from 'react-textarea-autosize'
import { nanoid } from 'nanoid'
import { RadioGroup, RadioGroupItem } from '@/components/ui/radio-group'

interface ChatPanelProps {
messages: UIState
Expand All @@ -27,6 +27,7 @@ export const ChatPanel = forwardRef<ChatPanelRef, ChatPanelProps>(({ messages, i
// Removed mcp instance as it's no longer passed to submit
const [isMobile, setIsMobile] = useState(false)
const [selectedFile, setSelectedFile] = useState<File | null>(null)
const [searchMode, setSearchMode] = useState('Standard')
const inputRef = useRef<HTMLTextAreaElement>(null)
const formRef = useRef<HTMLFormElement>(null)
const fileInputRef = useRef<HTMLInputElement>(null)
Expand Down Expand Up @@ -169,15 +170,53 @@ export const ChatPanel = forwardRef<ChatPanelRef, ChatPanelProps>(({ messages, i
isMobile ? 'px-2 pb-2 pt-1 h-full flex flex-col justify-center' : ''
)}
>
<div className="flex justify-center mb-2">
<RadioGroup
defaultValue="Standard"
className="flex bg-muted rounded-full p-1"
onValueChange={setSearchMode}
>
<RadioGroupItem value="Standard" id="Standard" className="sr-only" />
<label
htmlFor="Standard"
className={`cursor-pointer rounded-full px-3 py-1 text-sm ${
searchMode === 'Standard' ? 'bg-background text-foreground' : 'text-muted-foreground'
}`}
>
<Search className="h-4 w-4" />
</label>
<RadioGroupItem value="Geospatial" id="Geospatial" className="sr-only" />
<label
htmlFor="Geospatial"
className={`cursor-pointer rounded-full px-3 py-1 text-sm ${
searchMode === 'Geospatial' ? 'bg-background text-foreground' : 'text-muted-foreground'
}`}
>
<MapPin className="h-4 w-4" />
</label>
<RadioGroupItem value="Web Search" id="Web Search" className="sr-only" />
<label
htmlFor="Web Search"
className={`cursor-pointer rounded-full px-3 py-1 text-sm ${
searchMode === 'Web Search' ? 'bg-background text-foreground' : 'text-muted-foreground'
}`}
>
<Globe className="h-4 w-4" />
</label>
Comment on lines +197 to +205
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

id attributes must not contain spaces. Using id="Web Search" can break the label–input association and CSS selectors, leading to the labels not toggling the radio inputs reliably. This is a correctness bug in the new UI.

Suggestion

Use space-free ids and keep the display values as-is. Also bind value to state for consistency (see next comment). For example:

<RadioGroup
  value={searchMode}
  onValueChange={setSearchMode}
  defaultValue="Standard"
  className="flex bg-muted rounded-full p-1"
>
  <RadioGroupItem value="Standard" id="standard" className="sr-only" />
  <label htmlFor="standard" className={`cursor-pointer rounded-full px-3 py-1 text-sm ${searchMode === 'Standard' ? 'bg-background text-foreground' : 'text-muted-foreground'}`}>
    <Search className="h-4 w-4" />
  </label>

  <RadioGroupItem value="Geospatial" id="geospatial" className="sr-only" />
  <label htmlFor="geospatial" className={`cursor-pointer rounded-full px-3 py-1 text-sm ${searchMode === 'Geospatial' ? 'bg-background text-foreground' : 'text-muted-foreground'}`}>
    <MapPin className="h-4 w-4" />
  </label>

  <RadioGroupItem value="Web Search" id="web-search" className="sr-only" />
  <label htmlFor="web-search" className={`cursor-pointer rounded-full px-3 py-1 text-sm ${searchMode === 'Web Search' ? 'bg-background text-foreground' : 'text-muted-foreground'}`}>
    <Globe className="h-4 w-4" />
  </label>
</RadioGroup>

Reply with "@CharlieHelps yes please" if you'd like me to add a commit with this suggestion

Comment on lines +180 to +205
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The icon-only labels lack accessible names. Screen-reader users won’t know what each control does. Provide accessible labels (e.g., aria-label) or include visually hidden text.

Suggestion

Add aria-label to each label or include an sr-only span:

<label htmlFor="standard" aria-label="Standard mode" className={...}>
  <Search className="h-4 w-4" />
</label>
<label htmlFor="geospatial" aria-label="Geospatial mode" className={...}>
  <MapPin className="h-4 w-4" />
</label>
<label htmlFor="web-search" aria-label="Web Search mode" className={...}>
  <Globe className="h-4 w-4" />
</label>

Reply with "@CharlieHelps yes please" if you'd like me to add a commit with this suggestion

</RadioGroup>
</div>
<div
className={cn(
'relative flex items-start w-full',
isMobile && 'mobile-chat-input' // Apply mobile chat input styling
)}
>
<input type="hidden" name="searchMode" value={searchMode} />
<input
type="file"
ref={fileInputRef}
ref={fileInputNode => {
fileInputRef.current = fileInputNode
}}
onChange={handleFileChange}
className="hidden"
accept="text/plain,image/png,image/jpeg,image/webp"
Expand Down
45 changes: 45 additions & 0 deletions components/map/location-links.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
'use client';

import { useMapData } from './map-data-context';
import { Button } from '@/components/ui/button';
import { MapPin } from 'lucide-react';

interface Location {
latitude: number;
longitude: number;
place_name: string;
}

interface LocationLinksProps {
locations: Location[];
}

export const LocationLinks: React.FC<LocationLinksProps> = ({ locations }) => {
const { setMapData } = useMapData();

const handleFlyTo = (location: Location) => {
setMapData(prevData => ({
...prevData,
targetPosition: [location.longitude, location.latitude],
mapFeature: {
place_name: location.place_name,
}
}));
};

return (
<div className="flex flex-col space-y-2">
{locations.map((location, index) => (
<Button
key={index}
variant="outline"
className="justify-start"
onClick={() => handleFlyTo(location)}
>
<MapPin className="mr-2 h-4 w-4" />
{location.place_name}
</Button>
Comment on lines +33 to +41
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using the array index as a React key can cause unnecessary re-renders or state mismatches if the list changes. Even if the list is static per response, a stable key is preferable.

Suggestion

Use a stable key derived from the location itself, e.g., coordinates or place_name:

{locations.map((location) => (
  <Button
    key={`${location.longitude},${location.latitude}`}
    variant="outline"
    className="justify-start"
    onClick={() => handleFlyTo(location)}
  >
    <MapPin className="mr-2 h-4 w-4" />
    {location.place_name}
  </Button>
))}

Reply with "@CharlieHelps yes please" if you'd like me to add a commit with this suggestion

))}
</div>
);
};
85 changes: 31 additions & 54 deletions components/map/map-query-handler.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,83 +1,60 @@
'use client';

import { useEffect } from 'react';
// Removed useMCPMapClient as we'll use data passed via props
import { useMapData } from './map-data-context';
import { useMapData } from './map-data-context';
import { LocationLinks } from './location-links';

interface Location {
latitude: number;
longitude: number;
place_name: string;
}

// Define the expected structure of the mcp_response from geospatialTool
interface McpResponseData {
location: {
latitude?: number;
longitude?: number;
place_name?: string;
address?: string;
};
locations?: Location[]; // Support multiple locations
location?: Location; // Support single location
mapUrl?: string;
}

interface GeospatialToolOutput {
type: string; // e.g., "MAP_QUERY_TRIGGER"
type: string;
originalUserInput: string;
timestamp: string;
mcp_response: McpResponseData | null;
}

interface MapQueryHandlerProps {
// originalUserInput: string; // Kept for now, but primary data will come from toolOutput
toolOutput?: GeospatialToolOutput | null; // The direct output from geospatialTool
toolOutput?: GeospatialToolOutput | null;
}

export const MapQueryHandler: React.FC<MapQueryHandlerProps> = ({ toolOutput }) => {
const { setMapData } = useMapData();

useEffect(() => {
if (toolOutput && toolOutput.mcp_response && toolOutput.mcp_response.location) {
const { latitude, longitude, place_name } = toolOutput.mcp_response.location;
if (toolOutput && toolOutput.mcp_response) {
const { locations, location } = toolOutput.mcp_response;

if (typeof latitude === 'number' && typeof longitude === 'number') {
console.log(`MapQueryHandler: Received data from geospatialTool. Place: ${place_name}, Lat: ${latitude}, Lng: ${longitude}`);
setMapData(prevData => ({
...prevData,
// Ensure coordinates are in [lng, lat] format for MapboxGL
targetPosition: [longitude, latitude],
// Optionally store more info from mcp_response if needed by MapboxMap component later
mapFeature: {
place_name,
// Potentially add mapUrl or other details from toolOutput.mcp_response
mapUrl: toolOutput.mcp_response?.mapUrl
}
}));
if (locations && locations.length > 1) {
// Multiple locations: handled by LocationLinks, no automatic fly-to
} else {
console.warn("MapQueryHandler: Invalid latitude/longitude in toolOutput.mcp_response:", toolOutput.mcp_response.location);
// Clear target position if data is invalid
setMapData(prevData => ({
...prevData,
targetPosition: null,
mapFeature: null
}));
const singleLocation = locations && locations.length === 1 ? locations[0] : location;
if (singleLocation && typeof singleLocation.latitude === 'number' && typeof singleLocation.longitude === 'number') {
setMapData(prevData => ({
...prevData,
targetPosition: [singleLocation.longitude, singleLocation.latitude],
mapFeature: {
place_name: singleLocation.place_name,
mapUrl: toolOutput.mcp_response?.mapUrl
}
}));
}
}
Comment on lines +34 to 51
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When the tool output is invalid (no valid coordinates), the previous implementation cleared targetPosition/mapFeature. The new version silently does nothing, which can leave stale map state visible. Consider resetting the map state when no valid single location is present and there aren’t multiple locations to show.

Suggestion

Add a fallback reset when no valid single location is available and there aren’t multiple locations:

if (locations && locations.length > 1) {
  // handled by LocationLinks
} else {
  const singleLocation = locations && locations.length === 1 ? locations[0] : location;
  if (singleLocation && typeof singleLocation.latitude === 'number' && typeof singleLocation.longitude === 'number') {
    setMapData(prevData => ({
      ...prevData,
      targetPosition: [singleLocation.longitude, singleLocation.latitude],
      mapFeature: {
        place_name: singleLocation.place_name,
        mapUrl: toolOutput.mcp_response?.mapUrl
      }
    }));
  } else {
    setMapData(prevData => ({ ...prevData, targetPosition: null, mapFeature: null }));
  }
}

Reply with "@CharlieHelps yes please" if you'd like me to add a commit with this suggestion

} else {
// This case handles when toolOutput or its critical parts are missing.
// Depending on requirements, could fall back to originalUserInput and useMCPMapClient,
// or simply log that no valid data was provided from the tool.
// For this subtask, we primarily focus on using the new toolOutput.
if (toolOutput) { // It exists, but data is not as expected
console.warn("MapQueryHandler: toolOutput provided, but mcp_response or location data is missing.", toolOutput);
}
// If toolOutput is null/undefined, this component might not need to do anything,
// or it's an indication that it shouldn't have been rendered/triggered.
// For now, if no valid toolOutput, we clear map data or leave it as is.
// setMapData(prevData => ({ ...prevData, targetPosition: null, mapFeature: null }));
}
// The dependencies for this useEffect should be based on the props that trigger its logic.
// If originalUserInput and the old MCP client were still used as a fallback, they'd be dependencies.
}, [toolOutput, setMapData]);

// This component is a handler and does not render any visible UI itself.
// Its purpose is to trigger map data updates based on AI tool results.
// If it were to use the old useMCPMapClient, mcpLoading and mcpError would be relevant.
// It could return a small status indicator or debug info if needed for development.
if (toolOutput?.mcp_response?.locations && toolOutput.mcp_response.locations.length > 1) {
return <LocationLinks locations={toolOutput.mcp_response.locations} />;
}

return null;
// Example for debugging with previous client:
// return <div data-map-query-processed={originalUserInput} data-mcp-loading={mcpLoading} data-mcp-error={mcpError} style={{display: 'none'}} />;
};
98 changes: 44 additions & 54 deletions lib/agents/researcher.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,12 @@ import { getTools } from './tools'
import { getModel } from '../utils'

export async function researcher(
dynamicSystemPrompt: string, // New parameter
dynamicSystemPrompt: string,
uiStream: ReturnType<typeof createStreamableUI>,
streamText: ReturnType<typeof createStreamableValue<string>>,
messages: CoreMessage[],
// mcp: any, // Removed mcp parameter
useSpecificModel?: boolean
useSpecificModel?: boolean,
searchMode?: string
) {
let fullResponse = ''
let hasError = false
Expand All @@ -28,59 +28,49 @@ export async function researcher(
)

const currentDate = new Date().toLocaleString()
// Default system prompt, used if dynamicSystemPrompt is not provided
const default_system_prompt = `As a comprehensive AI assistant, you can search the web, retrieve information from URLs except from maps -here use the Geospatial tools provided, and understand geospatial queries to assist the user and display information on a map.
Current date and time: ${currentDate}. When tools are not needed, provide direct, helpful answers based on your knowledge.Match the language of your response to the user's language.
Always aim to directly address the user's question. If using information from a tool (like web search), cite the source URL.

There are also some proconfigured example queires.
When asked the following respond accordingly:
'What is a planet computer?' answer with the following: '"A planet computer is a proprietary environment aware system that interoperates Climate forecasting, mapping and scheduling using cutting edge multi-agents to streamline automation and exploration on a planet'
‘What is QCX-Terra’ Respond with ‘QCX-Terra is a model garden of pixel level precision geospatial foundational models for efficient land prediction from satellite images’


Tool Usage Guide:

- For general web searches for factual information: Use the 'search' tool.
- For retrieving content from specific URLs provided by the user: Use the 'retrieve' tool. (Do not use this for URLs found in search results).

- For any questions involving locations, places, addresses, geographical features, finding businesses or points of interest, distances between locations, or directions: You MUST use the 'geospatialQueryTool'. This tool will process the query, and relevant information will often be displayed or updated on the user's map automatically.**
Examples of queries for 'geospatialQueryTool':
Location Discovery
"Find coffee shops within walking distance of the Empire State Building"
"Show me gas stations along the route from Boston to New York"
"What restaurants are near Times Square?"
Navigation & Travel
"Get driving directions from LAX to Hollywood with current traffic"
"How long would it take to walk from Central Park to Times Square?"
"Calculate travel time from my hotel (Four Seasons) to JFK Airport by taxi during rush hour"
Visualization & Maps
"Create a map image showing the route from Golden Gate Bridge to Fisherman's Wharf with markers at both locations"
"Show me a satellite view of Manhattan with key landmarks marked"
"Generate a map highlighting all Starbucks locations within a mile of downtown Seattle"
Analysis & Planning
"Show me areas reachable within 30 minutes of downtown Portland by car"
"Calculate a travel time matrix between these 3 hotel locations (Marriott, Sheraton and Hilton) and the convention center in Denver"
"Find the optimal route visiting these 3 tourist attractions (Golden Gate, Musical Stairs and Fisherman's Wharf) in San Francisco"

When you use 'geospatialQueryTool', you don't need to describe how the map will change; simply provide your textual answer based on the query, and trust the map will update appropriately.
`;

const systemToUse = dynamicSystemPrompt && dynamicSystemPrompt.trim() !== '' ? dynamicSystemPrompt : default_system_prompt;

const result = await nonexperimental_streamText({
model: getModel() as LanguageModel,
maxTokens: 2500,
system: systemToUse, // Use the dynamic or default system prompt
messages,
tools: getTools({
uiStream,
fullResponse,
// mcp // mcp parameter is no longer passed to getTools
})

let systemPrompt = `Current date and time: ${currentDate}. Match the language of your response to the user's language. Always aim to directly address the user's question. If using information from a tool (like web search), cite the source URL.`;

const standardPrompt = `As a comprehensive AI assistant, you can search the web, retrieve information from URLs, and understand geospatial queries to assist the user and display information on a map. When tools are not needed, provide direct, helpful answers based on your knowledge.

Tool Usage Guide:
- For general web searches: Use the 'search' tool.
- For retrieving content from specific URLs: Use the 'retrieve' tool.
- For any questions involving locations, places, or directions: You MUST use the 'geospatialQueryTool'.`;

const geospatialPrompt = `You are a specialized Geospatial AI assistant. Your primary function is to understand and respond to geospatial queries. You MUST prioritize using the 'geospatialQueryTool' for any questions involving locations, places, addresses, geographical features, businesses, points of interest, distances, or directions. Only use other tools if geospatial queries are not applicable.`;

const webSearchPrompt = `You are a specialized Web Search AI assistant. Your primary function is to search the web and retrieve information from URLs to answer user questions. You MUST prioritize using the 'search' and 'retrieve' tools. Only use other tools if web searches are not applicable.`;

switch (searchMode) {
case 'Geospatial':
systemPrompt = `${geospatialPrompt}\n${systemPrompt}`;
break;
case 'Web Search':
systemPrompt = `${webSearchPrompt}\n${systemPrompt}`;
break;
default:
systemPrompt = `${standardPrompt}\n${systemPrompt}`;
break;
}
Comment on lines +32 to +55
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

dynamicSystemPrompt is now ignored. This is a functional regression from the previous implementation where a caller-provided system prompt could shape the agent’s behavior. The new persona logic should augment or allow override by dynamicSystemPrompt, not discard it.

Suggestion

Combine the persona prompt with dynamicSystemPrompt, giving the dynamic prompt precedence if present:

let systemPrompt = `Current date and time: ${currentDate}. Match the language of your response to the user's language. Always aim to directly address the user's question. If using information from a tool (like web search), cite the source URL.`;

// Build personaPrompt based on searchMode as you do now...

// If a dynamic system prompt is provided, prepend or override
if (dynamicSystemPrompt && dynamicSystemPrompt.trim() !== '') {
  systemPrompt = `${dynamicSystemPrompt.trim()}\n\n${systemPrompt}`;
}

Reply with "@CharlieHelps yes please" if you'd like me to add a commit with this suggestion


const allTools = getTools({ uiStream, fullResponse });
let availableTools: any = allTools;

if (searchMode === 'Geospatial') {
availableTools = { geospatialQueryTool: allTools.geospatialQueryTool };
} else if (searchMode === 'Web Search') {
availableTools = { search: allTools.search, retrieve: allTools.retrieve };
}
Comment on lines +60 to +64
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Over-restricting tools by removing them entirely can degrade fallback behavior. For example, the Geospatial persona might still benefit from retrieve to cite sources. Consider keeping complementary tools available and relying on the persona prompt to prioritize usage, rather than hard removal.

Suggestion

Keep complementary tools available for personas while instructing prioritization via the prompt. For example, allow retrieve with Geospatial:

if (searchMode === 'Geospatial') {
  availableTools = {
    geospatialQueryTool: allTools.geospatialQueryTool,
    retrieve: allTools.retrieve
  };
}

Alternatively, leave all tools available and rely solely on the persona prompt to guide selection.

Reply with "@CharlieHelps yes please" if you'd like me to add a commit with this suggestion


Comment on lines +57 to +65
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Avoid any for availableTools. It weakens type safety and risks passing malformed tool definitions to the model at runtime. You can keep the correct shape with a Partial<typeof allTools> and narrow explicitly.

Suggestion

Preserve types by using Partial<typeof allTools> and explicit construction:

const allTools = getTools({ uiStream, fullResponse });
let availableTools: Partial<typeof allTools> = allTools;

if (searchMode === 'Geospatial') {
  availableTools = { geospatialQueryTool: allTools.geospatialQueryTool };
} else if (searchMode === 'Web Search') {
  availableTools = { search: allTools.search, retrieve: allTools.retrieve };
}

Reply with "@CharlieHelps yes please" if you'd like me to add a commit with this suggestion

const result = await nonexperimental_streamText({
model: getModel() as LanguageModel,
maxTokens: 2500,
system: systemPrompt,
messages,
tools: availableTools
})

// Remove the spinner
uiStream.update(null)

// Process the response
Expand Down