using System.Data; using System.Text.Json; using Tango.Portal.Chat.Web.Models; using Tango.Portal.Chat.Web.Services; using Kusto.Data.Data; using Microsoft.AspNetCore.Mvc; using Newtonsoft.Json.Linq; using System.Text.Json.Nodes; using System.Collections; using Tango.Portal.Chat.Web.Utils; namespace Tango.Portal.Chat.Web.Controllers { [ApiController] [Route("api/[controller]")] public sealed class ChatController : ControllerBase { private readonly SchemaRegistry _schema; private readonly KqlGuard _guard; private readonly KustoQueryService _adx; private readonly LlmClient _llm; private readonly ChatMessageLogger _logger; private readonly AIInstructionService _instructionService; public ChatController(SchemaRegistry schema, KqlGuard guard, KustoQueryService adx, LlmClient llm, ChatMessageLogger logger, AIInstructionService instructionService) { _schema = schema; _guard = guard; _adx = adx; _llm = llm; _logger = logger; _instructionService = instructionService; } [HttpPost("ask")] public async Task> Ask([FromBody] ChatRequest req, CancellationToken ct, LlmProvider? provider = null) { try { if (!SessionUtils.IsUserAuthenticated(HttpContext)) { return new ChatResponse { Answer = "User is not authenticated or session expired", ThreadId = req.ThreadId }; } var sessionUser = SessionUtils.GetSessionUser(HttpContext); var sessionId = HttpContext.Session.Id; // Handle SYSTEM commands for roy@twine-s.com if (req.Question.StartsWith("SYSTEM:", StringComparison.OrdinalIgnoreCase) && sessionUser?.Email?.Equals("roy@twine-s.com", StringComparison.OrdinalIgnoreCase) == true) { return await HandleSystemCommandAsync(req.Question, sessionUser.Email, ct); } // Log the question _ = Task.Run(async () => { try { await _logger.LogQuestionAsync( sessionId, sessionUser?.Email ?? "unknown", sessionUser?.FullName ?? "unknown", req.Question, ct); } catch { // Ignore logging failures } }, ct); var schemaJson = _schema.GetSchemaJson(); var plannerPrompt = await _schema.GetPlannerPromptAsync(); var plotySample = _schema.GetPlotySample(); // 1) Ask the model for KQL ProposeKqlResult? plan = null; if (provider == LlmProvider.Claude) { plan = await _llm.ProposeKqlWithClaudeAsync(plannerPrompt, plotySample, req.Question, schemaJson, req.History, ct); } else { plan = await _llm.ProposeKqlAsync(plannerPrompt, plotySample, req.Question, schemaJson, req.History, ct); } ChatResponse response; if (plan.Assistant == "data" || plan.Assistant == "ploty") { response = await AnswerWithDataAssistant(req, plan, ct); } else if (plan.Assistant == "docs") { response = await AnswerWithDocsAssistant(req, plan, ct); } else { response = AnswerWithPlannerConversation(req, plan); } // Log the answer _ = Task.Run(async () => { try { await _logger.LogAnswerAsync( sessionId, sessionUser?.Email ?? "unknown", sessionUser?.FullName ?? "unknown", response, plan.Assistant, plan.Provider.ToString(), plan.Assumptions, ct); } catch { // Ignore logging failures } }, ct); return response; } catch (Exception ex) { var errorResponse = new ChatResponse { Answer = $"Ooops something went wrong...\n{ex.Message}", ThreadId = req.ThreadId }; // Log the error response var sessionUser = SessionUtils.GetSessionUser(HttpContext); var sessionId = req.ThreadId ?? Guid.NewGuid().ToString(); _ = Task.Run(async () => { try { await _logger.LogAnswerAsync( sessionId, sessionUser?.Email ?? "unknown", sessionUser?.FullName ?? "unknown", errorResponse, "error", "Unknown", null, ct); } catch { // Ignore logging failures } }, ct); return errorResponse; } } private ChatResponse AnswerWithPlannerConversation(ChatRequest req, ProposeKqlResult plan) { return new ChatResponse { Answer = plan.ConversationAnswer, ThreadId = req.ThreadId }; } private async Task AnswerWithDocsAssistant(ChatRequest req, ProposeKqlResult plan, CancellationToken ct) { // AFTER var run = await _llm.AnswerWithAssistantAsync( LlmClient.AssistantType.Docs, req.Question, string.Empty, plan.Kql, req.ThreadId, // <-- reuse if provided ct); return new ChatResponse { Answer = run.Answer, ThreadId = run.ThreadId }; } private async Task AnswerWithDataAssistant(ChatRequest req, ProposeKqlResult plan, CancellationToken ct) { // 2) Guardrail validation var val = _guard.Validate(plan.Kql); if (!val.IsOk) { // Return error to the client so they can iterate return new ChatResponse { Answer = $"The generated kusto query contains invalid tokens..\n{val.Error}", ThreadId = req.ThreadId, UsedKql = plan.Kql }; } // 4) Execute in ADX DataTable table; try { table = await _adx.QueryAsync(plan.Kql, plan.Parameters, ct); } catch (Exception ex) { if (plan.Provider == LlmProvider.OpenAI) { var fallbackResult = await Ask(req, ct, LlmProvider.Claude); return fallbackResult.Value ?? new ChatResponse { Answer = "Fallback to Claude failed", ThreadId = req.ThreadId }; } else { // Return error to the client so they can iterate return new ChatResponse { Answer = $"Seems like my kusto query ran into some issue..\n{ex.Message}", ThreadId = req.ThreadId, UsedKql = plan.Kql }; } } if (table.Columns.Contains("ploty") && table.Rows.Count > 0) { return AnswerWithPloty(req, plan, table); } // 5) Build compact facts (limit rows/cols) if (table.Rows.Count <= 200) { return await AnswerWithDataAssistantInternal(req, plan, table, ct); } else { return AnswerWithMarkdownTable(req, plan, table); } } private async Task AnswerWithDataAssistantInternal(ChatRequest req, ProposeKqlResult plan, DataTable table, CancellationToken ct) { var preview = DataHelper.ToPreview(table, 200); var facts = JsonSerializer.Serialize(preview); // 6) Ask model for final answer //var answer = await _llm.AnswerFromFactsAsync(req.Question, facts, plan.Kql, ct); var run = await _llm.AnswerWithAssistantAsync( LlmClient.AssistantType.Data, req.Question, facts, plan.Kql, req.ThreadId, // <-- reuse if provided ct); return new ChatResponse { Answer = run.Answer, UsedKql = plan.Kql, Preview = preview, ThreadId = run.ThreadId // <-- echo back the thread id used/created }; } private ChatResponse AnswerWithMarkdownTable(ChatRequest req, ProposeKqlResult plan, DataTable table) { var markdown = DataHelper.ToMarkdownTable(table); return new ChatResponse { Answer = $"The result set was too big for me to analyze so I'm just going to drop the entire result set.\n{markdown}", UsedKql = plan.Kql, Preview = markdown, ThreadId = req.ThreadId // <-- echo back the thread id used/created }; } private ChatResponse AnswerWithPloty(ChatRequest req, ProposeKqlResult plan, DataTable table) { String? ploty = table.Rows[0]["ploty"]?.ToString(); return new ChatResponse { Answer = plan.ConversationAnswer, ThreadId = req.ThreadId, UsedKql = plan.Kql, Ploty = ploty ?? String.Empty }; } private async Task> HandleSystemCommandAsync(string question, string userEmail, CancellationToken ct) { var instruction = question.Substring(7).Trim(); // Remove "SYSTEM:" prefix if (instruction.Equals("delete", StringComparison.OrdinalIgnoreCase)) { var success = await _instructionService.DeleteLastInstructionAsync(); return new ChatResponse { Answer = success ? "The last instruction has been successfully deleted." : "No instructions found to delete or deletion failed.", ThreadId = null }; } else if (!string.IsNullOrWhiteSpace(instruction)) { var success = await _instructionService.AddInstructionAsync(instruction, userEmail); return new ChatResponse { Answer = success ? "The new instruction has been successfully added." : "Failed to add the instruction.", ThreadId = null }; } else { return new ChatResponse { Answer = "Invalid SYSTEM command. Use 'SYSTEM: ' to add or 'SYSTEM: delete' to remove the last instruction.", ThreadId = null }; } } } }