path: root/Software/Visual_Studio_22/Tango.Portal.Chat.Web/Controllers/ChatController.cs

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;

        public ChatController(SchemaRegistry schema, KqlGuard guard, KustoQueryService adx, LlmClient llm, ChatMessageLogger logger)
        {
            _schema = schema;
            _guard = guard;
            _adx = adx;
            _llm = llm;
            _logger = logger;
        }

        [HttpPost("ask")]
        public async Task<ActionResult<ChatResponse>> 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;

                // 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 = _schema.GetPlannerPrompt();
                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(),
                            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",
                            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<ChatResponse> 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<ChatResponse> 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<ChatResponse> 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
            };
        }
    }
}