All files / src/llm-orchestration tool-schema.service.ts

97.87% Statements 46/47
100% Branches 6/6
100% Functions 5/5
97.77% Lines 44/45

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 15518x     18x 18x               18x 18x     18x 18x                         4x   4x 5x 5x 4x   1x           4x 4x                                   8x   8x   8x   10x 2x     2x     8x 8x 8x               8x         8x 8x 7x     4x     4x 7x     4x 3x 3x           1x         8x                   13x   12x 12x   12x 4x       4x 4x       12x                            
import { Injectable, Logger } from '@nestjs/common';
import { ActionHandler } from './action-handlers/action-handler.interface';
import { LlmFunctionTool } from '../llm-provider/llm-provider.interface';
import { SystemPromptsService } from '../system-prompts/system-prompts.service';
import { McpService } from '../mcp/mcp.service';
 
/**
 * Service to generate LLM tool function definitions from ActionHandler metadata.
 * Converts structured ToolMetadata into JSON Schema format compatible with OpenAI/Z.AI APIs.
 * Also includes MCP tool definitions from connected servers.
 */
@Injectable()
export class ToolSchemaService {
  private readonly logger = new Logger(ToolSchemaService.name);
 
  constructor(
    private readonly systemPromptsService: SystemPromptsService,
    private readonly mcpService: McpService,
  ) {}
 
  /**
   * Converts the ActionHandler registry to an array of LlmFunctionTool definitions.
   * Uses the structured metadata from handlers instead of parsing markdown strings.
   *
   * @param handlerRegistry Map of tool_name to ActionHandler
   * @returns Array of LlmFunctionTool definitions
   */
  generateToolDefinitions(
    handlerRegistry: Map<string, ActionHandler>,
  ): LlmFunctionTool[] {
    const tools: LlmFunctionTool[] = [];
 
    for (const [toolName, handler] of handlerRegistry.entries()) {
      try {
        const toolDefinition = this.generateToolDefinition(handler);
        tools.push(toolDefinition);
      } catch (error) {
        this.logger.warn(
          `Failed to generate tool definition for ${toolName}: ${error.message}`,
        );
      }
    }
 
    this.logger.log(`Generated ${tools.length} tool definitions.`);
    return tools;
  }
 
  /**
   * Generates tool definitions for a specific system prompt.
   * Only includes built-in tools that are enabled for that prompt.
   * Also includes MCP tools filtered by the prompt's enabledMcpTools setting.
   *
   * @param handlerRegistry Map of all available handlers
   * @param systemPromptId The ID of the active system prompt
   * @returns Array of LlmFunctionTool definitions (built-in + MCP)
   */
  async generateToolDefinitionsForPrompt(
    handlerRegistry: Map<string, ActionHandler>,
    systemPromptId: string,
  ): Promise<LlmFunctionTool[]> {
    // Get enabled tools for this prompt
    const enabledToolNames =
      await this.systemPromptsService.getToolsForPrompt(systemPromptId);
 
    const tools: LlmFunctionTool[] = [];
 
    for (const [toolName, handler] of handlerRegistry.entries()) {
      // Skip tools not enabled for this prompt
      if (!enabledToolNames.includes(toolName)) {
        this.logger.debug(
          `Tool ${toolName} is not enabled for system prompt ${systemPromptId}, skipping.`,
        );
        continue;
      }
 
      try {
        const toolDefinition = this.generateToolDefinition(handler);
        tools.push(toolDefinition);
      } catch (error) {
        this.logger.warn(
          `Failed to generate tool definition for ${toolName}: ${error.message}`,
        );
      }
    }
 
    this.logger.log(
      `Generated ${tools.length} built-in tool definitions for system prompt ${systemPromptId}.`,
    );
 
    // Add MCP tool definitions filtered by prompt-level override
    try {
      const allMcpTools = await this.mcpService.getActiveMcpToolDefinitions();
      if (allMcpTools.length > 0) {
        // Get prompt-level MCP tool filter
        const enabledMcpToolNames =
          await this.systemPromptsService.getMcpToolsForPrompt(systemPromptId);
 
        // Filter MCP tools: only include those in the enabled list
        const filteredMcpTools = allMcpTools.filter((tool) =>
          enabledMcpToolNames.includes(tool.function.name),
        );
 
        if (filteredMcpTools.length > 0) {
          tools.push(...filteredMcpTools);
          this.logger.log(
            `Added ${filteredMcpTools.length} MCP tool definitions (filtered from ${allMcpTools.length} total).`,
          );
        }
      }
    } catch (error) {
      this.logger.warn(
        `Failed to load MCP tool definitions: ${error.message}. Continuing without MCP tools.`,
      );
    }
 
    return tools;
  }
 
  /**
   * Generates a single LlmFunctionTool from an ActionHandler's metadata.
   *
   * @param handler The ActionHandler to convert
   * @returns LlmFunctionTool definition
   */
  private generateToolDefinition(handler: ActionHandler): LlmFunctionTool {
    const metadata = handler.getMetadata();
 
    const properties: Record<string, any> = {};
    const required: string[] = [];
 
    for (const arg of metadata.arguments) {
      properties[arg.name] = {
        type: arg.type,
        description: arg.description,
      };
      if (arg.required) {
        required.push(arg.name);
      }
    }
 
    return {
      type: 'function',
      function: {
        name: metadata.name,
        description: metadata.description,
        parameters: {
          type: 'object',
          properties,
          required: required.length > 0 ? required : undefined,
        },
      },
    };
  }
}