Skip to content

Add Agent Identity #121

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
kalaiyarasiganeshalingam wants to merge 13 commits into
base: main
Choose a base branch
from
Open

Add Agent Identity #121

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
3d2ce95
[Automated] Update the native jar versions
kalaiyarasiganeshalingam Feb 2, 2026
4fb03bf
[Automated] Update the native jar versions
kalaiyarasiganeshalingam Feb 2, 2026
0ccb6b7
[Automated] Update the native jar versions
kalaiyarasiganeshalingam Feb 2, 2026
c4033cd
[Automated] Update the native jar versions
kalaiyarasiganeshalingam Feb 2, 2026
208bd17
Add scope in the tool
kalaiyarasiganeshalingam Feb 6, 2026
857ee8b
Merge branch 'main' into main
kalaiyarasiganeshalingam Feb 19, 2026
1b907fd
[Automated] Update the native jar versions
kalaiyarasiganeshalingam Feb 19, 2026
bf3e1bf
Handle panic
kalaiyarasiganeshalingam Feb 22, 2026
b40f08d
[Automated] Update the native jar versions
kalaiyarasiganeshalingam Feb 23, 2026
8783860
[Automated] Update the native jar versions
kalaiyarasiganeshalingam Feb 23, 2026
e782170
[Automated] Update the native jar versions
kalaiyarasiganeshalingam Feb 23, 2026
d98f8c8
Fix comments
kalaiyarasiganeshalingam Feb 23, 2026
83944cc
Add tracing to agent identity
kalaiyarasiganeshalingam Feb 23, 2026
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
29 changes: 29 additions & 0 deletions ballerina/Dependencies.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,21 +12,26 @@ org = "ballerina"
name = "ai"
version = "1.10.0"
dependencies = [
{org = "ballerina", name = "cache"},
{org = "ballerina", name = "constraint"},
{org = "ballerina", name = "crypto"},
{org = "ballerina", name = "data.jsondata"},
{org = "ballerina", name = "data.xmldata"},
{org = "ballerina", name = "file"},
{org = "ballerina", name = "http"},
{org = "ballerina", name = "io"},
{org = "ballerina", name = "jballerina.java"},
{org = "ballerina", name = "jwt"},
{org = "ballerina", name = "lang.array"},
{org = "ballerina", name = "lang.regexp"},
{org = "ballerina", name = "lang.runtime"},
{org = "ballerina", name = "lang.string"},
{org = "ballerina", name = "log"},
{org = "ballerina", name = "math.vector"},
{org = "ballerina", name = "mcp"},
{org = "ballerina", name = "mime"},
{org = "ballerina", name = "observe"},
{org = "ballerina", name = "random"},
{org = "ballerina", name = "test"},
{org = "ballerina", name = "time"},
{org = "ballerina", name = "url"},
Expand Down Expand Up @@ -61,6 +66,9 @@ dependencies = [
{org = "ballerina", name = "task"},
{org = "ballerina", name = "time"}
]
modules = [
{org = "ballerina", packageName = "cache", moduleName = "cache"}
]

[[package]]
org = "ballerina"
Expand All @@ -81,6 +89,9 @@ dependencies = [
{org = "ballerina", name = "jballerina.java"},
{org = "ballerina", name = "time"}
]
modules = [
{org = "ballerina", packageName = "crypto", moduleName = "crypto"}
]

[[package]]
org = "ballerina"
Expand Down Expand Up @@ -187,6 +198,9 @@ dependencies = [
{org = "ballerina", name = "log"},
{org = "ballerina", name = "time"}
]
modules = [
{org = "ballerina", packageName = "jwt", moduleName = "jwt"}
]

[[package]]
org = "ballerina"
Expand Down Expand Up @@ -271,6 +285,9 @@ dependencies = [
{org = "ballerina", name = "jballerina.java"},
{org = "ballerina", name = "lang.regexp"}
]
modules = [
{org = "ballerina", packageName = "lang.string", moduleName = "lang.string"}
]

[[package]]
org = "ballerina"
Expand Down Expand Up @@ -363,6 +380,18 @@ dependencies = [
{org = "ballerina", name = "jballerina.java"}
]

[[package]]
org = "ballerina"
name = "random"
version = "1.7.0"
dependencies = [
{org = "ballerina", name = "jballerina.java"},
{org = "ballerina", name = "time"}
]
modules = [
{org = "ballerina", packageName = "random", moduleName = "random"}
]

[[package]]
org = "ballerina"
name = "task"
Expand Down
149 changes: 102 additions & 47 deletions ballerina/agent-utils.bal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

import ai.observe;

import ballerina/cache;
import ballerina/io;
import ballerina/log;
import ballerina/time;
Expand Down Expand Up @@ -85,9 +86,11 @@ public type ToolOutput record {|

type BaseAgent distinct isolated object {
ModelProvider model;
ToolStore toolStore;
ToolManager toolManager;
Memory memory;
boolean stateless;
cache:Cache tokenManager;
AuthConfig? auth;

# Parse the llm response and extract the tool to be executed.
#
Expand Down Expand Up @@ -115,6 +118,9 @@ class Executor {
private final BaseAgent agent;
# Contains the current execution progress for the agent and the query
public ExecutionProgress progress;
private string agentId = "";
private boolean & readonly isAuthEnabled = false;


# Initialize the executor with the agent and the query.
#
Expand All @@ -126,6 +132,11 @@ class Executor {
self.sessionId = sessionId;
self.agent = agent;
self.progress = progress;
AuthConfig? auth = agent.auth;
if auth is AuthConfig {
self.isAuthEnabled = true;
self.agentId = auth.agentId;
}
}

# Checks whether agent has more steps to execute.
Expand All @@ -140,6 +151,12 @@ class Executor {
# + return - generated LLM response during the reasoning or an error if the reasoning fails
public isolated function reason() returns json|Error {
if self.isCompleted {
if self.isAuthEnabled {
log:printError("Task is already completed. No more reasoning is needed.",
executionId = self.progress.executionId,
agentId = self.agentId
);
}
return error TaskCompletedError("Task is already completed. No more reasoning is needed.");
}
log:printDebug("LLM reasoning started",
Expand Down Expand Up @@ -181,56 +198,94 @@ class Executor {
if toolCallId is string {
span.addId(toolCallId);
}
string? toolDescription = self.agent.toolStore.getToolDescription(toolName);
ToolManager toolManager = self.agent.toolManager;
string? toolDescription = toolManager.getToolDescription(toolName);
if toolDescription is string {
span.addDescription(toolDescription);

}
span.addType(self.agent.toolStore.isMcpTool(toolName) ? observe:EXTENTION : observe:FUNCTION);
boolean isMcpTool = toolManager.isMcpTool(toolName);
span.addType(isMcpTool ? observe:EXTENTION : observe:FUNCTION);
span.addArguments(parsedOutput.arguments);

ToolOutput|ToolExecutionError|LlmInvalidGenerationError output = self.agent.toolStore.execute(parsedOutput,
self.progress.context);
if output is Error {
if output is ToolNotFoundError {
observation = "Tool is not found. Please check the tool name and retry.";
} else if output is ToolInvalidInputError {
observation = "Tool execution failed due to invalid inputs. Retry with correct inputs.";
} else {
observation = "Tool execution failed. Retry with correct inputs.";
}
observation = string `${observation.toString()} <detail>${output.toString()}</detail>`;
LlmInvalidGenerationError|ToolExecutionError? validateRes = toolManager.
validateTool(parsedOutput, self.agent.auth, self.agent.tokenManager,
self.progress.context, isMcpTool);
if validateRes is Error {
log:printError("Tool validation failed",
executionId = self.progress.executionId,
agentId = self.agentId,
sessionId = self.sessionId,
toolName = toolName,
'error = validateRes
);
observation = "Tool extraction failed due to tool validation";
executionResult = {
llmResponse,
'error: output,
observation: observation.toString()
'error: validateRes,
observation: "Tool extraction failed due to tool validation"
};

log:printDebug("Tool execution resulted in error",
executionId = self.progress.executionId,
observation = observation.toString(),
sessionId = self.sessionId,
toolName = toolName
);

Error toolExecutionError = error Error(observation.toString(), details = {parsedOutput});
span.close(toolExecutionError);
} else {
anydata|error value = output.value;
observation = value is error ? value.toString() : value;
log:printDebug("Tool execution successful",
executionId = self.progress.executionId,
sessionId = self.sessionId,
toolName = toolName,
output = observation
);
executionResult = {
tool: parsedOutput,
observation: value
};

span.addOutput(observation);
span.close();
ToolOutput|ToolExecutionError|LlmInvalidGenerationError output = toolManager.execute(parsedOutput,
self.progress.context);
if output is Error {
if output is ToolNotFoundError {
observation = "Tool is not found. Please check the tool name and retry.";
} else if output is ToolInvalidInputError {
observation = "Tool execution failed due to invalid inputs. Retry with correct inputs.";
} else {
observation = "Tool execution failed. Retry with correct inputs.";
}
observation = string `${observation.toString()} <detail>${output.toString()}</detail>`;
executionResult = {
llmResponse,
'error: output,
observation: observation.toString()
};

if self.isAuthEnabled {
log:printError("Tool execution resulted in error",
executionId = self.progress.executionId,
agentId = self.agentId,
observation = observation.toString(),
sessionId = self.sessionId,
toolName = toolName
);
} else {
log:printDebug("Tool execution resulted in error",
executionId = self.progress.executionId,
observation = observation.toString(),
sessionId = self.sessionId,
toolName = toolName
);
}

Error toolExecutionError = error Error(observation.toString(), details = {parsedOutput});
span.close(toolExecutionError);
} else {
anydata|error value = output.value;
observation = value is error ? value.toString() : value;
if self.isAuthEnabled {
log:printInfo("Tool execution successful",
executionId = self.progress.executionId,
agentId = self.agentId,
sessionId = self.sessionId,
toolName = toolName
);
} else {
log:printDebug("Tool execution successful",
executionId = self.progress.executionId,
sessionId = self.sessionId,
toolName = toolName,
output = observation
);
}
executionResult = {
tool: parsedOutput,
observation: value
};

span.addOutput(observation);
span.close();
}
}
} else {
log:printDebug("Failed to parse LLM response as valid tool or chat",
Expand Down Expand Up @@ -303,9 +358,9 @@ isolated function run(BaseAgent agent, string instruction, string query, int max
executionId = executionId,
sessionId = sessionId,
maxIterations = maxIter,
tools = agent.toolStore.tools.toString(),
tools = agent.toolManager.tools.toString(),
isStateless = agent.stateless
);
);

(ExecutionResult|ExecutionError|Error)[] steps = [];
string? content = ();
Expand Down Expand Up @@ -470,7 +525,7 @@ isolated function getOutputOfIteration(ExecutionResult|LlmChatResponse|Execution
}

isolated function buildCurrentIterationHistory(ExecutionProgress progress,
ChatMessage[] conversationHistoryUpToCurrentUserQuery) returns ChatMessage[] {
ChatMessage[] conversationHistoryUpToCurrentUserQuery) returns ChatMessage[] {
ChatMessage[] messages = createFunctionCallMessages(progress);
messages.unshift(...conversationHistoryUpToCurrentUserQuery);
return messages;
Expand All @@ -497,7 +552,7 @@ isolated function getObservationString(anydata|error observation) returns string
#
# + agent - Agent instance
# + return - Array of tools registered with the agent
public isolated function getTools(Agent agent) returns Tool[] => agent.functionCallAgent.toolStore.tools.toArray();
public isolated function getTools(Agent agent) returns Tool[] => agent.functionCallAgent.toolManager.tools.toArray();

isolated function updateMemory(Memory memory, string sessionId, ChatMessage[] messages) {
error? updationStation = memory.update(sessionId, messages);
Expand Down
Loading
Loading