Add library search and get tools for BI Copilot Agent #1400
Conversation
WalkthroughThe PR replaces a single library-provider flow with a two-step library discovery (keyword search → get details), adds toolCallId propagation for tool correlation, extends RPC types and client APIs for library search, updates prompts and tools to the new workflow, and normalizes function/library handling and UI message flows. Changes
Sequence DiagramsequenceDiagram
participant Agent as AI Agent
participant Executor as Tool Executor
participant SearchTool as Library Search Tool
participant LangClient as Lang Client
participant API as Search API
participant GetTool as Library Get Tool
Agent->>Executor: Initialize tools (search, get)
Agent->>SearchTool: tool_call(keywords, toolCallId)
activate SearchTool
SearchTool->>LangClient: getCopilotLibrariesBySearch(keywords)
LangClient->>API: POST /copilotLibraryManager/getLibrariesBySearch
API-->>LangClient: MinifiedLibrary[]
SearchTool->>Agent: tool_result(candidates, toolCallId)
deactivate SearchTool
Agent->>GetTool: tool_call(libraryId, toolCallId)
activate GetTool
GetTool->>GetTool: Fetch full library details
GetTool->>Agent: tool_result(fullLibrary, toolCallId)
deactivate GetTool
Agent->>Executor: Continue processing with library details
Estimated Code Review Effort🎯 4 (Complex) | ⏱️ ~45 minutes
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
RNViththagan
force-pushed
the
migration
branch
2 times, most recently
from
9b46c9a to
95aa2b1
Compare
RNViththagan
force-pushed
the
migration
branch
2 times, most recently
from
f7fb541 to
4fe6bcd
Compare
workspaces/ballerina/ballerina-extension/src/features/ai/agent/tools/connector-generator.ts
Outdated
Show resolved
Hide resolved
RNViththagan
force-pushed
the
migration
branch
from
547f098 to
5ed341b
Compare
workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/index.tsx
Outdated
Show resolved
Hide resolved
RNViththagan
force-pushed
the
migration
branch
from
79bff74 to
c57d1b7
Compare
RNViththagan
force-pushed
the
migration
branch
from
c57d1b7 to
9ce90c7
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
workspaces/ballerina/ballerina-extension/src/features/ai/utils/libs/function-registry.ts (1)
377-385:⚠️ Potential issue | 🟡 Minor
nullassigned to optional fields typed asT[] | undefined.
functionsandserviceson theLibraryinterface are optional (?), typed asRemoteFunction[] | undefinedandService[] | undefinedrespectively. Assigningnullviolates this type contract. While optional chaining (?.) works with bothnullandundefined, code expecting the type contract (e.g.,if (lib.functions === undefined)) would fail.🔧 Suggested fix
const maximizedLib: Library = { name: funcResponse.name, description: originalLib.description, clients: filteredClients, - functions: filteredFunctions ? filteredFunctions : null, + functions: filteredFunctions ?? undefined, // Get only the type definitions that are actually used by the selected functions and clients typeDefs: getOwnTypeDefsForLib(filteredClients, filteredFunctions, originalLib.typeDefs), - services: originalLib.services ? originalLib.services : null, + services: originalLib.services ?? undefined, };
🤖 Fix all issues with AI agents
In `@workspaces/ballerina/ballerina-extension/src/features/ai/agent/prompts.ts`:
- Around line 118-119: Edit mode's "Step 2: Identify necessary libraries" prompt
is missing the fallback to call ${CONNECTOR_GENERATOR_TOOL} when no suitable
library is found (present in Plan mode). Modify the Edit mode Step 2 prompt in
prompts.ts to mirror Plan mode: after instructing to use ${LIBRARY_SEARCH_TOOL}
and ${LIBRARY_GET_TOOL}, add a clear fallback sentence such as "If NO suitable
library is found, call ${CONNECTOR_GENERATOR_TOOL}." Ensure you update the Edit
mode prompt text where Step 2 is defined so it matches the Plan mode behavior
and references ${CONNECTOR_GENERATOR_TOOL}.
In
`@workspaces/ballerina/ballerina-extension/src/features/ai/agent/tools/library-search.ts`:
- Around line 152-189: The tool documentation in LibrarySearchTool refers to a
non-existent "LibraryProviderTool" in its description text and examples; update
all occurrences to the actual tool name "LibraryGetTool" inside
workspaces/ballerina/ballerina-extension/src/features/ai/agent/tools/library-search.ts
so the agent calls the correct tool (search for the symbol/identifier
LibrarySearchTool and replace every "LibraryProviderTool" string in its
description, examples, and any related comments with "LibraryGetTool").
- Around line 86-91: The early return when keywords is empty (after slicing via
MAX_SEARCH_KEYWORDS) bypasses emitting the corresponding tool_result, leaving
the UI stuck; instead of returning immediately in the keywords.length === 0
branch, emit the same tool_result event used elsewhere (the tool_result emission
that should pair with the earlier tool_call) with { libraries: [] } and then
return that value — i.e., locate the keywords check around
params.keywords.slice(0, MAX_SEARCH_KEYWORDS) and replace the immediate return
with a call to the existing tool_result emission helper (use the same
event-emission function/mechanism used after tool_call) followed by returning {
libraries: [] } so the UI receives a tool_result for the empty-keyword case.
In
`@workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/index.tsx`:
- Around line 424-444: The current replacement for LibrarySearchTool
reconstructs originalMessage from toolOutput?.searchDescription and then does
content.replace(...) which can silently fail if toolInput?.searchDescription and
toolOutput differ; instead, locate the last message by toolCallId and replace
the entire <toolcall id="...">...</toolcall> block with a <toolresult
id="...">...</toolresult> irrespective of the inner text. Update the logic in
the LibrarySearchTool branch (look for response.toolName ===
"LibrarySearchTool", toolCallId, toolOutput, and updateLastMessage) to match and
replace by the unique toolCallId (as you do for LibraryGetTool) so the
spinner/tag is always removed even when descriptions diverge.
🧹 Nitpick comments (5)
workspaces/ballerina/ballerina-extension/src/features/ai/agent/tools/connector-generator.ts (1)
51-51: Address the TODO before merging or convert to a tracked issue.This TODO about verifying alignment with the updated tool workflow should be resolved to avoid stale notes accumulating in the codebase.
Would you like me to open an issue to track this verification task?
workspaces/ballerina/ballerina-extension/src/features/ai/utils/libs/libraries.ts (1)
22-26: Remove commented-out dead code.This block has been commented out and serves no purpose. Consider removing it to keep the file clean.
workspaces/ballerina/ballerina-extension/src/features/ai/agent/tools/library-get.ts (1)
132-141: User prompt logged in full — verify this is acceptable for your privacy posture.Line 139 logs
input.userPromptverbatim. While this is a local extension console log (not sent externally), user prompts can contain sensitive information. Confirm this level of logging is acceptable or consider truncating.workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/index.tsx (1)
384-384: Inconsistent equality operators:==vs===.Lines 384, 457, and 466 use loose equality (
==) forresponse.toolNamecomparisons while the newly added branches (lines 367, 378, 424, 445) use strict equality (===). These are pre-existing but now visually inconsistent. Consider unifying to===for consistency.Also applies to: 457-457, 466-466
workspaces/ballerina/ballerina-extension/src/features/ai/agent/tools/healthcare-library.ts (1)
58-95:HealthcareLibraryProviderToollackstoolCallIdpropagation unlike the new library tools.The
LibrarySearchToolandLibraryGetToolboth propagatetoolCallIdthroughtool_call/tool_resultevents for UI correlation. This tool doesn't, which means if multiple healthcare tool calls ever run concurrently, the UI won't be able to correlate them correctly. Currently this may not be an issue if only one healthcare tool call happens at a time, but it's worth noting for future-proofing.
workspaces/ballerina/ballerina-extension/src/features/ai/agent/prompts.ts
Show resolved
Hide resolved
| // Validate keyword count | ||
| const keywords = params.keywords.slice(0, MAX_SEARCH_KEYWORDS); | ||
| if (keywords.length === 0) { | ||
| console.warn(`[LibrarySearchTool] No keywords provided`); | ||
| return { libraries: [] }; | ||
| } |
There was a problem hiding this comment.
Early return on empty keywords skips tool_result emission, leaving the UI stuck in "searching" state.
The tool_call event is already emitted (lines 74-82) before this check. If keywords are empty after slicing, the function returns { libraries: [] } without emitting a corresponding tool_result event. The UI <toolcall> tag will never be replaced with a <toolresult>.
🐛 Proposed fix
// Validate keyword count
const keywords = params.keywords.slice(0, MAX_SEARCH_KEYWORDS);
if (keywords.length === 0) {
console.warn(`[LibrarySearchTool] No keywords provided`);
+ emitLibrarySearchResult(eventHandler, LIBRARY_SEARCH_TOOL, [], toolCallId, params.searchDescription);
return { libraries: [] };
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // Validate keyword count | |
| const keywords = params.keywords.slice(0, MAX_SEARCH_KEYWORDS); | |
| if (keywords.length === 0) { | |
| console.warn(`[LibrarySearchTool] No keywords provided`); | |
| return { libraries: [] }; | |
| } | |
| // Validate keyword count | |
| const keywords = params.keywords.slice(0, MAX_SEARCH_KEYWORDS); | |
| if (keywords.length === 0) { | |
| console.warn(`[LibrarySearchTool] No keywords provided`); | |
| emitLibrarySearchResult(eventHandler, LIBRARY_SEARCH_TOOL, [], toolCallId, params.searchDescription); | |
| return { libraries: [] }; | |
| } |
🤖 Prompt for AI Agents
In
`@workspaces/ballerina/ballerina-extension/src/features/ai/agent/tools/library-search.ts`
around lines 86 - 91, The early return when keywords is empty (after slicing via
MAX_SEARCH_KEYWORDS) bypasses emitting the corresponding tool_result, leaving
the UI stuck; instead of returning immediately in the keywords.length === 0
branch, emit the same tool_result event used elsewhere (the tool_result emission
that should pair with the earlier tool_call) with { libraries: [] } and then
return that value — i.e., locate the keywords check around
params.keywords.slice(0, MAX_SEARCH_KEYWORDS) and replace the immediate return
with a call to the existing tool_result emission helper (use the same
event-emission function/mechanism used after tool_call) followed by returning {
libraries: [] } so the UI receives a tool_result for the empty-keyword case.
| - Before calling LibraryProviderTool to retrieve full library details | ||
| - When the user query mentions integrations, services, connectors, or specific functionality | ||
| - Whenever you need to find relevant libraries but don't know the exact library names | ||
|
|
||
| **Important - Two-Step Workflow:** | ||
| 1. First, call THIS tool (LibrarySearchTool) with weighted keywords to discover relevant libraries | ||
| 2. Review the returned library names and descriptions | ||
| 3. Select the most appropriate libraries (typically 1-5 libraries) | ||
| 4. Then, call LibraryProviderTool with the selected library names to get detailed API documentation (functions, types, clients, services, etc.) | ||
|
|
||
| **Example Workflows:** | ||
|
|
||
| Example 1 - Stripe Integration: | ||
| User query: "I need to integrate with Stripe payment gateway" | ||
| Keywords: ["Stripe", "payment", "gateway"] // "Stripe" has highest weight | ||
| Call LibrarySearchTool with keywords: ["Stripe", "payment", "gateway"] | ||
| → Returns: [ | ||
| { name: "ballerinax/stripe", description: "Connects to Stripe API for payment processing" } | ||
| ] | ||
| Then call LibraryProviderTool with libraryNames: ["ballerinax/stripe"] | ||
|
|
||
| Example 2 - GitHub API: | ||
| User query: "Create a GitHub integration to list issues" | ||
| Keywords: ["GitHub", "API", "issues"] // "GitHub" has highest weight | ||
| Call LibrarySearchTool with keywords: ["GitHub", "API", "issues"] | ||
| → Returns: [ | ||
| { name: "ballerinax/github", description: "GitHub API connector for repository management" } | ||
| ] | ||
| Then call LibraryProviderTool with libraryNames: ["ballerinax/github"] | ||
|
|
||
| Example 3 - HTTP Service: | ||
| User query: "Create a REST API" | ||
| Keywords: ["HTTP", "REST", "API"] // "HTTP" has highest weight | ||
| Call LibrarySearchTool with keywords: ["HTTP", "REST", "API"] | ||
| → Returns: [ | ||
| { name: "ballerina/http", description: "HTTP client and server implementation" } | ||
| ] | ||
| Then call LibraryProviderTool with libraryNames: ["ballerina/http"] |
There was a problem hiding this comment.
Tool description references non-existent LibraryProviderTool — LLM will attempt to call a tool that doesn't exist.
The description tells the LLM to "call LibraryProviderTool" in five places (lines 152, 160, 171, 180, 189), but the actual tool is named LibraryGetTool. Since the LLM uses these descriptions to decide which tools to invoke, this will cause the agent to fail when it tries to call LibraryProviderTool.
🐛 Fix: Replace all references
**When to use this tool:**
- To discover which libraries are available for a specific use case or integration
-- Before calling LibraryProviderTool to retrieve full library details
+- Before calling LibraryGetTool to retrieve full library details
- When the user query mentions integrations, services, connectors, or specific functionality
...
-4. Then, call LibraryProviderTool with the selected library names to get detailed API documentation (functions, types, clients, services, etc.)
+4. Then, call LibraryGetTool with the selected library names to get detailed API documentation (functions, types, clients, services, etc.)
...
-Then call LibraryProviderTool with libraryNames: ["ballerinax/stripe"]
+Then call LibraryGetTool with libraryNames: ["ballerinax/stripe"]
...
-Then call LibraryProviderTool with libraryNames: ["ballerinax/github"]
+Then call LibraryGetTool with libraryNames: ["ballerinax/github"]
...
-Then call LibraryProviderTool with libraryNames: ["ballerina/http"]
+Then call LibraryGetTool with libraryNames: ["ballerina/http"]🤖 Prompt for AI Agents
In
`@workspaces/ballerina/ballerina-extension/src/features/ai/agent/tools/library-search.ts`
around lines 152 - 189, The tool documentation in LibrarySearchTool refers to a
non-existent "LibraryProviderTool" in its description text and examples; update
all occurrences to the actual tool name "LibraryGetTool" inside
workspaces/ballerina/ballerina-extension/src/features/ai/agent/tools/library-search.ts
so the agent calls the correct tool (search for the symbol/identifier
LibrarySearchTool and replace every "LibraryProviderTool" string in its
description, examples, and any related comments with "LibraryGetTool").
workspaces/ballerina/ballerina-visualizer/src/views/AIPanel/components/AIChat/index.tsx
Show resolved
Hide resolved
RNViththagan
changed the title
Introduces a mode parameter to copilot library requests, enabling filtering by core, healthcare, or all libraries. Unifies search and retrieval endpoints to support mode-based queries, enhancing library selection flexibility.
RNViththagan
force-pushed
the
migration
branch
from
9ce90c7 to
8182a7b
Compare
- Split library discovery into LibrarySearchTool and LibraryGetTool - Add toolCallId support for parallel tool call tracking - Add optional searchDescription parameter for contextual UI messages - Update UI to display tool-specific status with unique IDs - Disable loader auto turn off for tool in segment component
Introduced updateLastMessage helper to eliminate boilerplate code for message updates. Refactored LibrarySearchTool and LibraryGetTool handlers to use the new helper, reducing code duplication.
…alize library objects - Add getLibraryModeFromGenerationType() helper to properly map GenerationType enum to LibraryMode type - Normalize library objects returned from getCopilotFilteredLibraries with default empty arrays for missing fields - Replace type casting with explicit conversion for better type safety
RNViththagan
force-pushed
the
migration
branch
from
8182a7b to
632b381
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In
`@workspaces/ballerina/ballerina-extension/src/features/ai/agent/tools/library-search.ts`:
- Around line 200-208: The execute callback signature should be updated to match
Vercel AI SDK options: rename the second parameter to options and type it as {
toolCallId: string; messages: ModelMessage[]; abortSignal?: AbortSignal;
experimental_context?: unknown }, remove the unnecessary fallback toolCallId
logic, and use options.toolCallId (required) instead of context?.toolCallId;
also surface options.messages and options.abortSignal to the LibrarySearchTool
invocation (or pass them into the call/site that needs them) so execute forwards
the SDK-provided metadata—update the execute function, its parameter name, and
all references to use options.toolCallId, options.messages, and
options.abortSignal and adjust any type imports/usages (e.g., ModelMessage)
accordingly.
🧹 Nitpick comments (1)
workspaces/ballerina/ballerina-extension/src/features/ai/agent/tools/library-get.ts (1)
17-19: Consolidate duplicate"ai"imports.
tool(line 17) andjsonSchema(line 19) are imported from the same"ai"module in separate statements.♻️ Suggested fix
-import { tool } from "ai"; -import { GenerationType } from "../../utils/libs/libraries"; -import { jsonSchema } from "ai"; +import { tool, jsonSchema } from "ai"; +import { GenerationType } from "../../utils/libs/libraries";
| execute: async (input: { keywords: string[]; searchDescription?: string }, context?: { toolCallId?: string }) => { | ||
| // Extract toolCallId from AI SDK context | ||
| const toolCallId = context?.toolCallId || `fallback-${Date.now()}`; | ||
|
|
||
| console.log( | ||
| `[LibrarySearchTool] Called with keywords: [${input.keywords.join(", ")}] [toolCallId: ${toolCallId}] [description: ${input.searchDescription || 'none'}]` | ||
| ); | ||
| return await LibrarySearchTool(input, eventHandler, toolCallId); | ||
| }, |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cat -n workspaces/ballerina/ballerina-extension/src/features/ai/agent/tools/library-search.ts | sed -n '190,220p'Repository: wso2/vscode-extensions
Length of output: 1220
🌐 Web query:
Vercel AI SDK tool execute callback context parameter signature toolCallId
💡 Result:
In the Vercel AI SDK, a tool’s execute gets two parameters: the validated tool input, and a context/options object.
Signature (TypeScript shape):
execute: async (
input: INPUT,
options: {
toolCallId: string;
messages: ModelMessage[];
abortSignal?: AbortSignal;
experimental_context?: unknown;
}
) => RESULT;You can destructure it like:
execute: async (input, { toolCallId, messages, abortSignal, experimental_context }) => { /* ... */ }toolCallId is the unique id for that specific tool invocation (useful for tracking/annotations), and it’s provided via that second “options” parameter. [1], [2]
Fix the execute callback signature to match Vercel AI SDK's actual options parameter structure.
The second parameter should be options: { toolCallId: string; messages: ModelMessage[]; abortSignal?: AbortSignal; experimental_context?: unknown; }, not context?: { toolCallId?: string }. The current signature makes toolCallId optional when the SDK always provides it as required, which is why the fallback logic exists unnecessarily. Additionally, the code ignores messages and abortSignal which are available from the SDK context. Update the parameter name to options and include all SDK properties with correct optionality for proper type safety and to access all available tool invocation metadata.
🤖 Prompt for AI Agents
In
`@workspaces/ballerina/ballerina-extension/src/features/ai/agent/tools/library-search.ts`
around lines 200 - 208, The execute callback signature should be updated to
match Vercel AI SDK options: rename the second parameter to options and type it
as { toolCallId: string; messages: ModelMessage[]; abortSignal?: AbortSignal;
experimental_context?: unknown }, remove the unnecessary fallback toolCallId
logic, and use options.toolCallId (required) instead of context?.toolCallId;
also surface options.messages and options.abortSignal to the LibrarySearchTool
invocation (or pass them into the call/site that needs them) so execute forwards
the SDK-provided metadata—update the execute function, its parameter name, and
all references to use options.toolCallId, options.messages, and
options.abortSignal and adjust any type imports/usages (e.g., ModelMessage)
accordingly.
4 participants
Summary
Implements a two-tool library discovery system that splits library search and retrieval into separate tools with support for parallel execution tracking.
Changes
Resolves: wso2/product-ballerina-integrator#2386
fixes: wso2/product-ballerina-integrator#2401
Handles this from extension side: wso2/product-ballerina-integrator#2416
Summary by CodeRabbit
New Features
Improvements