feat: tool registration — brain tools for agent sessions (P2-004) #76
@@ -30,6 +30,7 @@
|
|||||||
"@opentelemetry/sdk-metrics": "^2.6.0",
|
"@opentelemetry/sdk-metrics": "^2.6.0",
|
||||||
"@opentelemetry/sdk-node": "^0.213.0",
|
"@opentelemetry/sdk-node": "^0.213.0",
|
||||||
"@opentelemetry/semantic-conventions": "^1.40.0",
|
"@opentelemetry/semantic-conventions": "^1.40.0",
|
||||||
|
"@sinclair/typebox": "^0.34.48",
|
||||||
"better-auth": "^1.5.5",
|
"better-auth": "^1.5.5",
|
||||||
"fastify": "^5.0.0",
|
"fastify": "^5.0.0",
|
||||||
"reflect-metadata": "^0.2.0",
|
"reflect-metadata": "^0.2.0",
|
||||||
|
|||||||
@@ -1,11 +1,15 @@
|
|||||||
import { Injectable, Logger, type OnModuleDestroy } from '@nestjs/common';
|
import { Inject, Injectable, Logger, type OnModuleDestroy } from '@nestjs/common';
|
||||||
import {
|
import {
|
||||||
createAgentSession,
|
createAgentSession,
|
||||||
SessionManager,
|
SessionManager,
|
||||||
type AgentSession as PiAgentSession,
|
type AgentSession as PiAgentSession,
|
||||||
type AgentSessionEvent,
|
type AgentSessionEvent,
|
||||||
|
type ToolDefinition,
|
||||||
} from '@mariozechner/pi-coding-agent';
|
} from '@mariozechner/pi-coding-agent';
|
||||||
|
import type { Brain } from '@mosaic/brain';
|
||||||
|
import { BRAIN } from '../brain/brain.tokens.js';
|
||||||
import { ProviderService } from './provider.service.js';
|
import { ProviderService } from './provider.service.js';
|
||||||
|
import { createBrainTools } from './tools/brain-tools.js';
|
||||||
|
|
||||||
export interface AgentSessionOptions {
|
export interface AgentSessionOptions {
|
||||||
provider?: string;
|
provider?: string;
|
||||||
@@ -27,7 +31,15 @@ export class AgentService implements OnModuleDestroy {
|
|||||||
private readonly sessions = new Map<string, AgentSession>();
|
private readonly sessions = new Map<string, AgentSession>();
|
||||||
private readonly creating = new Map<string, Promise<AgentSession>>();
|
private readonly creating = new Map<string, Promise<AgentSession>>();
|
||||||
|
|
||||||
constructor(private readonly providerService: ProviderService) {}
|
private readonly customTools: ToolDefinition[];
|
||||||
|
|
||||||
|
constructor(
|
||||||
|
private readonly providerService: ProviderService,
|
||||||
|
@Inject(BRAIN) private readonly brain: Brain,
|
||||||
|
) {
|
||||||
|
this.customTools = createBrainTools(brain);
|
||||||
|
this.logger.log(`Registered ${this.customTools.length} custom tools`);
|
||||||
|
}
|
||||||
|
|
||||||
async createSession(sessionId: string, options?: AgentSessionOptions): Promise<AgentSession> {
|
async createSession(sessionId: string, options?: AgentSessionOptions): Promise<AgentSession> {
|
||||||
const existing = this.sessions.get(sessionId);
|
const existing = this.sessions.get(sessionId);
|
||||||
@@ -62,6 +74,7 @@ export class AgentService implements OnModuleDestroy {
|
|||||||
modelRegistry: this.providerService.getRegistry(),
|
modelRegistry: this.providerService.getRegistry(),
|
||||||
model: model ?? undefined,
|
model: model ?? undefined,
|
||||||
tools: [],
|
tools: [],
|
||||||
|
customTools: this.customTools,
|
||||||
});
|
});
|
||||||
piSession = result.session;
|
piSession = result.session;
|
||||||
} catch (err) {
|
} catch (err) {
|
||||||
|
|||||||
184
apps/gateway/src/agent/tools/brain-tools.ts
Normal file
184
apps/gateway/src/agent/tools/brain-tools.ts
Normal file
@@ -0,0 +1,184 @@
|
|||||||
|
import { Type } from '@sinclair/typebox';
|
||||||
|
import type { ToolDefinition } from '@mariozechner/pi-coding-agent';
|
||||||
|
import type { Brain } from '@mosaic/brain';
|
||||||
|
|
||||||
|
export function createBrainTools(brain: Brain): ToolDefinition[] {
|
||||||
|
const listProjects: ToolDefinition = {
|
||||||
|
name: 'brain_list_projects',
|
||||||
|
label: 'List Projects',
|
||||||
|
description: 'List all projects in the brain.',
|
||||||
|
parameters: Type.Object({}),
|
||||||
|
async execute() {
|
||||||
|
const projects = await brain.projects.findAll();
|
||||||
|
return {
|
||||||
|
content: [{ type: 'text' as const, text: JSON.stringify(projects, null, 2) }],
|
||||||
|
details: undefined,
|
||||||
|
};
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
const getProject: ToolDefinition = {
|
||||||
|
name: 'brain_get_project',
|
||||||
|
label: 'Get Project',
|
||||||
|
description: 'Get a project by ID.',
|
||||||
|
parameters: Type.Object({
|
||||||
|
id: Type.String({ description: 'Project ID (UUID)' }),
|
||||||
|
}),
|
||||||
|
async execute(_toolCallId, params) {
|
||||||
|
const { id } = params as { id: string };
|
||||||
|
const project = await brain.projects.findById(id);
|
||||||
|
return {
|
||||||
|
content: [
|
||||||
|
{
|
||||||
|
type: 'text' as const,
|
||||||
|
text: project ? JSON.stringify(project, null, 2) : `Project not found: ${id}`,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
details: undefined,
|
||||||
|
};
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
const listTasks: ToolDefinition = {
|
||||||
|
name: 'brain_list_tasks',
|
||||||
|
label: 'List Tasks',
|
||||||
|
description: 'List tasks, optionally filtered by project, mission, or status.',
|
||||||
|
parameters: Type.Object({
|
||||||
|
projectId: Type.Optional(Type.String({ description: 'Filter by project ID' })),
|
||||||
|
missionId: Type.Optional(Type.String({ description: 'Filter by mission ID' })),
|
||||||
|
status: Type.Optional(Type.String({ description: 'Filter by status' })),
|
||||||
|
}),
|
||||||
|
async execute(_toolCallId, params) {
|
||||||
|
const p = params as { projectId?: string; missionId?: string; status?: string };
|
||||||
|
type TaskStatus = 'not-started' | 'in-progress' | 'blocked' | 'done' | 'cancelled';
|
||||||
|
let tasks;
|
||||||
|
if (p.projectId) tasks = await brain.tasks.findByProject(p.projectId);
|
||||||
|
else if (p.missionId) tasks = await brain.tasks.findByMission(p.missionId);
|
||||||
|
else if (p.status) tasks = await brain.tasks.findByStatus(p.status as TaskStatus);
|
||||||
|
else tasks = await brain.tasks.findAll();
|
||||||
|
return {
|
||||||
|
content: [{ type: 'text' as const, text: JSON.stringify(tasks, null, 2) }],
|
||||||
|
details: undefined,
|
||||||
|
};
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
const createTask: ToolDefinition = {
|
||||||
|
name: 'brain_create_task',
|
||||||
|
label: 'Create Task',
|
||||||
|
description: 'Create a new task in the brain.',
|
||||||
|
parameters: Type.Object({
|
||||||
|
title: Type.String({ description: 'Task title' }),
|
||||||
|
description: Type.Optional(Type.String({ description: 'Task description' })),
|
||||||
|
projectId: Type.Optional(Type.String({ description: 'Project ID' })),
|
||||||
|
missionId: Type.Optional(Type.String({ description: 'Mission ID' })),
|
||||||
|
priority: Type.Optional(
|
||||||
|
Type.String({ description: 'Priority: low, medium, high, critical' }),
|
||||||
|
),
|
||||||
|
}),
|
||||||
|
async execute(_toolCallId, params) {
|
||||||
|
const p = params as {
|
||||||
|
title: string;
|
||||||
|
description?: string;
|
||||||
|
projectId?: string;
|
||||||
|
missionId?: string;
|
||||||
|
priority?: string;
|
||||||
|
};
|
||||||
|
type Priority = 'low' | 'medium' | 'high' | 'critical';
|
||||||
|
const task = await brain.tasks.create({
|
||||||
|
...p,
|
||||||
|
priority: p.priority as Priority | undefined,
|
||||||
|
});
|
||||||
|
return {
|
||||||
|
content: [{ type: 'text' as const, text: JSON.stringify(task, null, 2) }],
|
||||||
|
details: undefined,
|
||||||
|
};
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
const updateTask: ToolDefinition = {
|
||||||
|
name: 'brain_update_task',
|
||||||
|
label: 'Update Task',
|
||||||
|
description: 'Update an existing task.',
|
||||||
|
parameters: Type.Object({
|
||||||
|
id: Type.String({ description: 'Task ID' }),
|
||||||
|
title: Type.Optional(Type.String()),
|
||||||
|
description: Type.Optional(Type.String()),
|
||||||
|
status: Type.Optional(
|
||||||
|
Type.String({ description: 'not-started, in-progress, blocked, done, cancelled' }),
|
||||||
|
),
|
||||||
|
priority: Type.Optional(Type.String()),
|
||||||
|
}),
|
||||||
|
async execute(_toolCallId, params) {
|
||||||
|
const { id, ...updates } = params as {
|
||||||
|
id: string;
|
||||||
|
title?: string;
|
||||||
|
description?: string;
|
||||||
|
status?: string;
|
||||||
|
priority?: string;
|
||||||
|
};
|
||||||
|
type TaskStatus = 'not-started' | 'in-progress' | 'blocked' | 'done' | 'cancelled';
|
||||||
|
type Priority = 'low' | 'medium' | 'high' | 'critical';
|
||||||
|
const task = await brain.tasks.update(id, {
|
||||||
|
...updates,
|
||||||
|
status: updates.status as TaskStatus | undefined,
|
||||||
|
priority: updates.priority as Priority | undefined,
|
||||||
|
});
|
||||||
|
return {
|
||||||
|
content: [
|
||||||
|
{
|
||||||
|
type: 'text' as const,
|
||||||
|
text: task ? JSON.stringify(task, null, 2) : `Task not found: ${id}`,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
details: undefined,
|
||||||
|
};
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
const listMissions: ToolDefinition = {
|
||||||
|
name: 'brain_list_missions',
|
||||||
|
label: 'List Missions',
|
||||||
|
description: 'List all missions, optionally filtered by project.',
|
||||||
|
parameters: Type.Object({
|
||||||
|
projectId: Type.Optional(Type.String({ description: 'Filter by project ID' })),
|
||||||
|
}),
|
||||||
|
async execute(_toolCallId, params) {
|
||||||
|
const p = params as { projectId?: string };
|
||||||
|
const missions = p.projectId
|
||||||
|
? await brain.missions.findByProject(p.projectId)
|
||||||
|
: await brain.missions.findAll();
|
||||||
|
return {
|
||||||
|
content: [{ type: 'text' as const, text: JSON.stringify(missions, null, 2) }],
|
||||||
|
details: undefined,
|
||||||
|
};
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
const listConversations: ToolDefinition = {
|
||||||
|
name: 'brain_list_conversations',
|
||||||
|
label: 'List Conversations',
|
||||||
|
description: 'List conversations for a user.',
|
||||||
|
parameters: Type.Object({
|
||||||
|
userId: Type.String({ description: 'User ID' }),
|
||||||
|
}),
|
||||||
|
async execute(_toolCallId, params) {
|
||||||
|
const { userId } = params as { userId: string };
|
||||||
|
const conversations = await brain.conversations.findAll(userId);
|
||||||
|
return {
|
||||||
|
content: [{ type: 'text' as const, text: JSON.stringify(conversations, null, 2) }],
|
||||||
|
details: undefined,
|
||||||
|
};
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
return [
|
||||||
|
listProjects,
|
||||||
|
getProject,
|
||||||
|
listTasks,
|
||||||
|
createTask,
|
||||||
|
updateTask,
|
||||||
|
listMissions,
|
||||||
|
listConversations,
|
||||||
|
];
|
||||||
|
}
|
||||||
1
apps/gateway/src/agent/tools/index.ts
Normal file
1
apps/gateway/src/agent/tools/index.ts
Normal file
@@ -0,0 +1 @@
|
|||||||
|
export { createBrainTools } from './brain-tools.js';
|
||||||
3
pnpm-lock.yaml
generated
3
pnpm-lock.yaml
generated
@@ -95,6 +95,9 @@ importers:
|
|||||||
'@opentelemetry/semantic-conventions':
|
'@opentelemetry/semantic-conventions':
|
||||||
specifier: ^1.40.0
|
specifier: ^1.40.0
|
||||||
version: 1.40.0
|
version: 1.40.0
|
||||||
|
'@sinclair/typebox':
|
||||||
|
specifier: ^0.34.48
|
||||||
|
version: 0.34.48
|
||||||
better-auth:
|
better-auth:
|
||||||
specifier: ^1.5.5
|
specifier: ^1.5.5
|
||||||
version: 1.5.5(drizzle-kit@0.31.9)(drizzle-orm@0.45.1(@opentelemetry/api@1.9.0)(@types/pg@8.15.6)(kysely@0.28.11)(postgres@3.4.8))(mongodb@7.1.0(socks@2.8.7))(next@16.1.6(@opentelemetry/api@1.9.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4))(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(vitest@2.1.9(@types/node@22.19.15))
|
version: 1.5.5(drizzle-kit@0.31.9)(drizzle-orm@0.45.1(@opentelemetry/api@1.9.0)(@types/pg@8.15.6)(kysely@0.28.11)(postgres@3.4.8))(mongodb@7.1.0(socks@2.8.7))(next@16.1.6(@opentelemetry/api@1.9.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4))(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(vitest@2.1.9(@types/node@22.19.15))
|
||||||
|
|||||||
Reference in New Issue
Block a user