// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Text;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
namespace Microsoft.AspNetCore.Components.AI.SourceGenerators;
internal static class ToolBlockEmitter
{
public static void EmitHandler(SourceProductionContext spc, ToolBlockCandidate candidate)
{
var blockType = candidate.BlockTypeGlobal;
var sb = new StringBuilder();
sb.AppendLine("// <auto-generated/>");
sb.AppendLine("#nullable enable");
sb.AppendLine();
if (!string.IsNullOrEmpty(candidate.Namespace))
{
sb.AppendLine($"namespace {candidate.Namespace};");
sb.AppendLine();
}
var className = candidate.ClassName + "Handler";
sb.AppendLine($"[global::System.ComponentModel.EditorBrowsable(global::System.ComponentModel.EditorBrowsableState.Never)]");
sb.AppendLine($"internal sealed class {className}");
sb.AppendLine($" : global::Microsoft.AspNetCore.Components.AI.ContentBlockHandler<{blockType}>");
sb.AppendLine("{");
sb.AppendLine($" public override global::Microsoft.AspNetCore.Components.AI.BlockMappingResult<{blockType}> Handle(");
sb.AppendLine($" global::Microsoft.AspNetCore.Components.AI.BlockMappingContext context,");
sb.AppendLine($" {blockType} state)");
sb.AppendLine(" {");
sb.AppendLine(" if (state.Result is not null)");
sb.AppendLine(" {");
sb.AppendLine($" return global::Microsoft.AspNetCore.Components.AI");
sb.AppendLine($" .BlockMappingResult<{blockType}>.Complete();");
sb.AppendLine(" }");
sb.AppendLine();
sb.AppendLine(" var shouldEmit = false;");
sb.AppendLine();
// Phase 1: Match FunctionCallContent by tool name
sb.AppendLine(" if (state.Call is null)");
sb.AppendLine(" {");
sb.AppendLine(" global::Microsoft.Extensions.AI.FunctionCallContent? callContent = null;");
sb.AppendLine(" foreach (var content in context.UnhandledContents)");
sb.AppendLine(" {");
sb.AppendLine(" if (content is global::Microsoft.Extensions.AI.FunctionCallContent fc");
sb.AppendLine($" && fc.Name == \"{EscapeString(candidate.ToolName)}\")");
sb.AppendLine(" {");
sb.AppendLine(" callContent = fc;");
sb.AppendLine(" break;");
sb.AppendLine(" }");
sb.AppendLine(" }");
sb.AppendLine();
sb.AppendLine(" if (callContent is not null)");
sb.AppendLine(" {");
sb.AppendLine(" context.MarkHandled(callContent);");
sb.AppendLine(" state.Call = callContent;");
sb.AppendLine(" shouldEmit = true;");
// Deserialize arguments
if (candidate.Parameters.Count > 0)
{
sb.AppendLine();
sb.AppendLine(" if (callContent.Arguments is { } args)");
sb.AppendLine(" {");
foreach (var param in candidate.Parameters)
{
var varName = LocalName(param.PropertyName);
var member = "state." + EscapeIdentifier(param.PropertyName);
sb.AppendLine($" if (args.TryGetValue(\"{EscapeString(param.ArgumentKey)}\", out var {varName}) && {varName} is not null)");
sb.AppendLine(" {");
sb.AppendLine($" {member} = {varName} switch");
sb.AppendLine(" {");
EmitDeserialization(sb, param, varName);
sb.AppendLine(" };");
sb.AppendLine(" }");
}
sb.AppendLine(" }");
}
sb.AppendLine(" }");
sb.AppendLine(" }");
sb.AppendLine();
// Phase 2: Match FunctionResultContent by CallId
sb.AppendLine(" global::Microsoft.Extensions.AI.FunctionResultContent? resultContent = null;");
sb.AppendLine(" foreach (var content in context.UnhandledContents)");
sb.AppendLine(" {");
sb.AppendLine(" if (content is global::Microsoft.Extensions.AI.FunctionResultContent frc");
sb.AppendLine(" && state.Call is not null");
sb.AppendLine(" && frc.CallId == state.Call.CallId)");
sb.AppendLine(" {");
sb.AppendLine(" resultContent = frc;");
sb.AppendLine(" break;");
sb.AppendLine(" }");
sb.AppendLine(" }");
sb.AppendLine();
sb.AppendLine(" if (resultContent is not null)");
sb.AppendLine(" {");
sb.AppendLine(" context.MarkHandled(resultContent);");
sb.AppendLine(" state.Result = resultContent;");
// Deserialize [ToolResult] properties
if (candidate.ResultProperties.Count > 0)
{
sb.AppendLine();
sb.AppendLine(" if (resultContent.Result is not null)");
sb.AppendLine(" {");
if (candidate.ResultProperties.Count == 1)
{
// Single result property: map the entire Result value directly
var rp = candidate.ResultProperties[0];
var varName = "__result";
var member = "state." + EscapeIdentifier(rp.PropertyName);
sb.AppendLine($" var {varName} = resultContent.Result;");
sb.AppendLine($" {member} = {varName} switch");
sb.AppendLine(" {");
EmitResultDeserialization(sb, rp, varName);
sb.AppendLine(" };");
}
else
{
// Multiple result properties: treat Result as a JSON object
sb.AppendLine(" var __resultObj = resultContent.Result switch");
sb.AppendLine(" {");
sb.AppendLine(" global::System.Text.Json.JsonElement __element => __element,");
sb.AppendLine(" string __json => global::System.Text.Json.JsonSerializer.Deserialize<global::System.Text.Json.JsonElement>(__json),");
sb.AppendLine(" _ => global::System.Text.Json.JsonSerializer.SerializeToElement(resultContent.Result),");
sb.AppendLine(" };");
sb.AppendLine();
sb.AppendLine(" if (__resultObj.ValueKind == global::System.Text.Json.JsonValueKind.Object)");
sb.AppendLine(" {");
foreach (var rp in candidate.ResultProperties)
{
var varName = "__r_" + LocalSuffix(rp.PropertyName);
var member = "state." + EscapeIdentifier(rp.PropertyName);
sb.AppendLine($" if (__resultObj.TryGetProperty(\"{EscapeString(rp.ResultKey)}\", out var {varName}))");
sb.AppendLine(" {");
sb.AppendLine($" {member} = {varName} switch");
sb.AppendLine(" {");
EmitJsonElementDeserialization(sb, rp, varName);
sb.AppendLine(" };");
sb.AppendLine(" }");
}
sb.AppendLine(" }");
}
sb.AppendLine(" }");
}
sb.AppendLine(" return shouldEmit");
sb.AppendLine($" ? global::Microsoft.AspNetCore.Components.AI.BlockMappingResult<{blockType}>.Emit(state, state)");
sb.AppendLine($" : global::Microsoft.AspNetCore.Components.AI.BlockMappingResult<{blockType}>.Complete();");
sb.AppendLine(" }");
sb.AppendLine();
sb.AppendLine(" return shouldEmit");
sb.AppendLine($" ? global::Microsoft.AspNetCore.Components.AI.BlockMappingResult<{blockType}>.Emit(state, state)");
sb.AppendLine($" : global::Microsoft.AspNetCore.Components.AI.BlockMappingResult<{blockType}>.Pass();");
sb.AppendLine(" }");
sb.AppendLine("}");
spc.AddSource(HintName(candidate), sb.ToString());
}
public static void EmitRegistration(SourceProductionContext spc, ImmutableArray<ToolBlockCandidate> candidates)
{
if (candidates.IsEmpty)
{
return;
}
// Check for duplicate tool names
var seen = new Dictionary<string, ToolBlockCandidate>();
foreach (var c in candidates)
{
if (seen.TryGetValue(c.ToolName, out var existing))
{
spc.ReportDiagnostic(Diagnostic.Create(
DiagnosticDescriptors.DuplicateToolName,
Location.None,
c.ToolName, existing.ClassName, c.ClassName));
}
else
{
seen[c.ToolName] = c;
}
}
var sb = new StringBuilder();
sb.AppendLine("// <auto-generated/>");
sb.AppendLine("#nullable enable");
sb.AppendLine();
sb.AppendLine("namespace Microsoft.AspNetCore.Components.AI;");
sb.AppendLine();
sb.AppendLine("[global::System.ComponentModel.EditorBrowsable(global::System.ComponentModel.EditorBrowsableState.Never)]");
sb.AppendLine("internal static class GeneratedToolBlockRegistrations");
sb.AppendLine("{");
sb.AppendLine(" internal static void AddGeneratedToolBlocks(this UIAgentOptions options)");
sb.AppendLine(" {");
foreach (var candidate in candidates)
{
var handlerName = candidate.ClassName + "Handler";
var fullClass = string.IsNullOrEmpty(candidate.Namespace)
? "global::" + handlerName
: "global::" + candidate.Namespace + "." + handlerName;
sb.AppendLine($" options.AddBlockHandler(new {fullClass}());");
}
sb.AppendLine(" }");
sb.AppendLine("}");
spc.AddSource("GeneratedToolBlockRegistrations.g.cs", sb.ToString());
}
// Hint names must be unique per generated file; two ToolBlocks with the same simple name in
// different namespaces would otherwise collide. Qualify with the (sanitized) namespace while
// keeping "<ClassName>Handler.g.cs" as the suffix.
private static string HintName(ToolBlockCandidate candidate)
{
var suffix = candidate.ClassName + "Handler.g.cs";
return string.IsNullOrEmpty(candidate.Namespace)
? suffix
: Sanitize(candidate.Namespace) + "." + suffix;
}
private static string Sanitize(string value)
{
var sb = new StringBuilder(value.Length);
foreach (var ch in value)
{
sb.Append(char.IsLetterOrDigit(ch) || ch == '.' || ch == '_' ? ch : '_');
}
return sb.ToString();
}
private static string EscapeIdentifier(string name)
=> SyntaxFacts.GetKeywordKind(name) != SyntaxKind.None
|| SyntaxFacts.GetContextualKeywordKind(name) != SyntaxKind.None
? "@" + name
: name;
private static string LocalName(string propertyName)
=> "__" + LocalSuffix(propertyName);
private static string LocalSuffix(string propertyName)
=> propertyName.Substring(0, 1).ToLowerInvariant() + propertyName.Substring(1);
private static void EmitDeserialization(StringBuilder sb, ToolParameterInfo param, string varName)
{
switch (param.TypeKind)
{
case ParameterTypeKind.String:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => __je.GetString()!,");
sb.AppendLine($" string __s => __s,");
sb.AppendLine($" _ => (string){varName}!");
break;
case ParameterTypeKind.Int32:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => __je.GetInt32(),");
sb.AppendLine($" _ => global::System.Convert.ToInt32({varName})");
break;
case ParameterTypeKind.Int64:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => __je.GetInt64(),");
sb.AppendLine($" _ => global::System.Convert.ToInt64({varName})");
break;
case ParameterTypeKind.Double:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => __je.GetDouble(),");
sb.AppendLine($" _ => global::System.Convert.ToDouble({varName})");
break;
case ParameterTypeKind.Single:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => __je.GetSingle(),");
sb.AppendLine($" _ => global::System.Convert.ToSingle({varName})");
break;
case ParameterTypeKind.Decimal:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => __je.GetDecimal(),");
sb.AppendLine($" _ => global::System.Convert.ToDecimal({varName})");
break;
case ParameterTypeKind.Boolean:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => __je.GetBoolean(),");
sb.AppendLine($" _ => global::System.Convert.ToBoolean({varName})");
break;
case ParameterTypeKind.Complex:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => global::System.Text.Json.JsonSerializer.Deserialize<{param.TypeName}>(__je)!,");
sb.AppendLine($" _ => ({param.TypeName}){varName}!");
break;
}
}
private static string EscapeString(string value)
{
return value.Replace("\\", "\\\\").Replace("\"", "\\\"");
}
private static void EmitResultDeserialization(StringBuilder sb, ToolResultPropertyInfo prop, string varName)
{
// Single-property result: map the entire Result value directly
switch (prop.TypeKind)
{
case ParameterTypeKind.String:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => __je.GetString()!,");
sb.AppendLine($" string __s => __s,");
sb.AppendLine($" _ => {varName}!.ToString()!");
break;
case ParameterTypeKind.Int32:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => __je.GetInt32(),");
sb.AppendLine($" _ => global::System.Convert.ToInt32({varName})");
break;
case ParameterTypeKind.Int64:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => __je.GetInt64(),");
sb.AppendLine($" _ => global::System.Convert.ToInt64({varName})");
break;
case ParameterTypeKind.Double:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => __je.GetDouble(),");
sb.AppendLine($" _ => global::System.Convert.ToDouble({varName})");
break;
case ParameterTypeKind.Single:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => __je.GetSingle(),");
sb.AppendLine($" _ => global::System.Convert.ToSingle({varName})");
break;
case ParameterTypeKind.Decimal:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => __je.GetDecimal(),");
sb.AppendLine($" _ => global::System.Convert.ToDecimal({varName})");
break;
case ParameterTypeKind.Boolean:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => __je.GetBoolean(),");
sb.AppendLine($" _ => global::System.Convert.ToBoolean({varName})");
break;
case ParameterTypeKind.Complex:
sb.AppendLine($" global::System.Text.Json.JsonElement __je => global::System.Text.Json.JsonSerializer.Deserialize<{prop.TypeName}>(__je)!,");
sb.AppendLine($" string __json => global::System.Text.Json.JsonSerializer.Deserialize<{prop.TypeName}>(__json)!,");
sb.AppendLine($" _ => ({prop.TypeName}){varName}!");
break;
}
}
private static void EmitJsonElementDeserialization(StringBuilder sb, ToolResultPropertyInfo prop, string varName)
{
// Multi-property result: deserialize from a JSON object property
switch (prop.TypeKind)
{
case ParameterTypeKind.String:
sb.AppendLine($" global::System.Text.Json.JsonElement {{ ValueKind: global::System.Text.Json.JsonValueKind.String }} => {varName}.GetString()!,");
sb.AppendLine($" _ => {varName}.GetRawText()");
break;
case ParameterTypeKind.Int32:
sb.AppendLine($" _ => {varName}.GetInt32()");
break;
case ParameterTypeKind.Int64:
sb.AppendLine($" _ => {varName}.GetInt64()");
break;
case ParameterTypeKind.Double:
sb.AppendLine($" _ => {varName}.GetDouble()");
break;
case ParameterTypeKind.Single:
sb.AppendLine($" _ => {varName}.GetSingle()");
break;
case ParameterTypeKind.Decimal:
sb.AppendLine($" _ => {varName}.GetDecimal()");
break;
case ParameterTypeKind.Boolean:
sb.AppendLine($" _ => {varName}.GetBoolean()");
break;
case ParameterTypeKind.Complex:
sb.AppendLine($" _ => global::System.Text.Json.JsonSerializer.Deserialize<{prop.TypeName}>({varName})!");
break;
}
}
}