ChatService.cs

using Microsoft.Graph;
using Microsoft.SemanticKernel;
using Microsoft.SemanticKernel.ChatCompletion;
using Microsoft.SemanticKernel.Connectors.OpenAI;
using Microsoft.SemanticKernel.Plugins.MsGraph;
using Microsoft.SemanticKernel.Plugins.MsGraph.Connectors;
using System.Net.Http.Headers;
using System.Reflection;
using System.Runtime.CompilerServices;
using System.Text;
using System.Text.Json;

namespace Office365SelfHostedLLMConnector;

// One chat-capable model advertised by the upstream gateway. ContextWindow is the
// best-known token capacity (0 = unknown). Surfaced to the frontend model picker
// via GET /models so the user can choose which model the backend calls per turn.
public sealed record LlmModelInfo(string Id, int ContextWindow);

// Per-request token accounting. Two sources of truth, marked by IsEstimated:
//
//   IsEstimated=false  → exact counts pulled from the upstream SSE chunk that
//                        OpenAI's spec attaches when `stream_options.include_usage`
//                        is requested. SK's OpenAI connector requests this for us.
//
//   IsEstimated=true   → upstream proxy didn't emit a usage chunk (many hosted
//                        and several other gateways strip it from streaming responses).
//                        We fall back to a chars/4 heuristic computed in
//                        RunChatStreamingAsync from the prompt history + streamed bytes.
//                        Accuracy is roughly ±25% — fine for cost visibility, NOT for
//                        billing. UI renders these with a "≈" prefix so the user knows.
public sealed record UsageInfo(int PromptTokens, int CompletionTokens, int TotalTokens, bool IsEstimated = false);

// Per-model context window (max input + output tokens the model accepts).
// Upstream gateway's /v1/models response strips this metadata, so we maintain
// our own lookup. Values come from each vendor's published spec; bump them
// when you add a new model to the gateway. Default is conservative (32K) so
// an unknown model still renders SOMETHING meaningful rather than 0 or null.
//
// LLM_CONTEXT_WINDOW env var, when set, OVERRIDES every per-model entry below.
// That's the right call for a single-LLM deployment (which is the production
// shape today) — operators set it to whatever the upstream gateway is
// configured for, and the UI shows the truth. The per-model dictionary stays
// as the multi-model fallback for environments that don't set the env var.
public static class ContextWindow
{
    private static readonly Dictionary<string, int> _windows = new(StringComparer.OrdinalIgnoreCase)
    {
        // Claude family — Anthropic publishes 200K for the Sonnet 4 series.
        ["claude-Sonnet-4-6"]  = 200_000,
        ["claude-Sonnet-4-5"]  = 200_000,
        ["claude-sonnet-4-6"]  = 200_000,
        ["claude-sonnet-4-5"]  = 200_000,
        // Grok 4.30 family — 1M token context window per xAI spec.
        ["grok-4.30-auto"]                       = 1_000_000,
        ["grok-4.30-fast"]                       = 1_000_000,
        ["grok-4.30-heavy"]                      = 1_000_000,
        ["grok-4.30-expert"]                     = 1_000_000,
        ["grok-4.30-0309"]                       = 1_000_000,
        ["grok-4.30-0309-reasoning"]             = 1_000_000,
        ["grok-4.30-0309-reasoning-heavy"]       = 1_000_000,
        ["grok-4.30-0309-reasoning-super"]       = 1_000_000,
        ["grok-4.30-0309-non-reasoning"]         = 1_000_000,
        ["grok-4.30-0309-non-reasoning-heavy"]   = 1_000_000,
        ["grok-4.30-0309-non-reasoning-super"]   = 1_000_000,
        ["grok-4.30-0309-super"]                 = 1_000_000,
        ["grok-4.30-0309-heavy"]                 = 1_000_000,
        ["grok-4.30-multi-agent-0309"]           = 1_000_000,
        // Image-only models have no chat context; map to 0 so the UI can skip the badge.
        ["grok-imagine-image"]      = 0,
        ["grok-imagine-image-lite"] = 0,
        ["grok-imagine-image-edit"] = 0,
        ["grok-imagine-image-pro"]  = 0,
        ["grok-imagine-video"]      = 0,
    };
    public const int Default = 32_000;

    // Parsed ONCE at type init from LLM_CONTEXT_WINDOW. null = no override; use the
    // per-model dictionary. Invalid / non-positive values fall through to the dict
    // (we don't want a typo'd env var to silently zero out the badge).
    private static readonly int? _override = ParseOverride(Environment.GetEnvironmentVariable("LLM_CONTEXT_WINDOW"));

    private static int? ParseOverride(string? raw)
    {
        if (string.IsNullOrWhiteSpace(raw)) return null;
        if (!int.TryParse(raw, out var n)) return null;
        return n > 0 ? n : (int?)null;
    }

    public static int For(string? modelId)
    {
        // Image-only models are dictionary-pinned at 0 regardless of env override —
        // they truly have no chat context, and showing 1M next to a generated image
        // would be misleading. So we honor the dict for explicit-zero entries first.
        if (!string.IsNullOrEmpty(modelId)
            && _windows.TryGetValue(modelId, out var dictVal) && dictVal == 0)
            return 0;

        if (_override is int n) return n;

        if (string.IsNullOrEmpty(modelId)) return Default;
        return _windows.TryGetValue(modelId, out var w) ? w : Default;
    }
}

// Extracted from Program.cs so both the legacy /chat (anonymous, stateless) and the
// new /conversations/{id}/chat (authorized, persisted) share one path to the LLM.
// Single source of truth for: kernel construction, Graph plugin wiring, streaming
// loop, and progress logging. Anything LLM-related changes here once, not twice.
public sealed class ChatService
{
    private readonly string _modelId;
    private readonly string _baseUrl;
    private readonly string _apiKey;
    private readonly HttpClient _httpClient;

    public ChatService(string modelId, string baseUrl, string apiKey, HttpClient httpClient)
    {
        _modelId = modelId;
        _baseUrl = baseUrl;
        _apiKey = apiKey;
        _httpClient = httpClient;
    }

    public string ModelId => _modelId;

    // Per-request kernel: Graph plugins are bound to the caller's access token, so we
    // can't share a kernel across users. The HttpClient (UA-rewriting) IS shared.
    // modelOverride lets the frontend model picker select a per-turn model; null/blank
    // falls back to the operator's configured default (_modelId). Same endpoint+key —
    // only the model name in the request body changes (standard OpenAI semantics).
    public Kernel CreateKernel(string accessToken, string? modelOverride = null)
    {
        var model = string.IsNullOrWhiteSpace(modelOverride) ? _modelId : modelOverride.Trim();
        var kernel = Kernel.CreateBuilder()
            .AddOpenAIChatCompletion(
                modelId: model,
                endpoint: new Uri(_baseUrl),
                apiKey: _apiKey,
                httpClient: _httpClient)
            .Build();

        var credential = new BearerTokenCredential(accessToken);
        var graphClient = new GraphServiceClient(credential);

        kernel.ImportPluginFromObject(new CalendarPlugin(new OutlookCalendarConnector(graphClient)), "Calendar");
        kernel.ImportPluginFromObject(new EmailPlugin(new OutlookMailConnector(graphClient)), "Email");
        kernel.ImportPluginFromObject(new CloudDrivePlugin(new OneDriveConnector(graphClient)), "OneDrive");

        return kernel;
    }

    // Fetches the chat-capable models the upstream gateway advertises (OpenAI
    // GET /v1/models). Image/video generation models are filtered out — they
    // aren't valid for /chat. Each id is annotated with its known context window
    // (0 when unknown) for the frontend picker. Caller is expected to cache the
    // result (see GET /models) so we don't hit the gateway on every page load.
    public async Task<IReadOnlyList<LlmModelInfo>> ListChatModelsAsync(CancellationToken ct = default)
    {
        using var req = new HttpRequestMessage(HttpMethod.Get, $"{_baseUrl}/models");
        req.Headers.Authorization = new AuthenticationHeaderValue("Bearer", _apiKey);
        using var resp = await _httpClient.SendAsync(req, ct);
        resp.EnsureSuccessStatusCode();

        var json = await resp.Content.ReadAsStringAsync(ct);
        using var doc = JsonDocument.Parse(json);

        var models = new List<LlmModelInfo>();
        // OpenAI shape: { "data": [ { "id": "...", ... }, ... ] }. Some gateways
        // return a bare array; handle both.
        var arr = doc.RootElement.ValueKind == JsonValueKind.Array
            ? doc.RootElement
            : (doc.RootElement.TryGetProperty("data", out var data) ? data : default);
        if (arr.ValueKind == JsonValueKind.Array)
        {
            foreach (var item in arr.EnumerateArray())
            {
                if (!item.TryGetProperty("id", out var idEl)) continue;
                var id = idEl.GetString();
                if (string.IsNullOrWhiteSpace(id)) continue;
                // Drop non-chat models (image/video generation) so they never
                // land in the chat picker where they'd 400 on use.
                if (id.Contains("image", StringComparison.OrdinalIgnoreCase) ||
                    id.Contains("video", StringComparison.OrdinalIgnoreCase) ||
                    id.Contains("embedding", StringComparison.OrdinalIgnoreCase)) continue;
                models.Add(new LlmModelInfo(id, ContextWindow.For(id)));
            }
        }
        return models;
    }

    // Runs the chat completion to completion (buffered) and returns the full text.
    // Buffered (not streamed) on purpose: the persisted endpoint needs to write the
    // assistant message in the same transaction as the user message. Streaming would
    // mean either (a) flush half a row on client disconnect, or (b) buffer server-side
    // anyway. We do log per-chunk progress so the operator sees liveness.
    public async Task<string> RunChatAsync(
        Kernel kernel,
        ChatHistory history,
        string logTag,
        CancellationToken ct = default)
    {
        var chatService = kernel.GetRequiredService<IChatCompletionService>();
        var settings = new OpenAIPromptExecutionSettings
        {
            ToolCallBehavior = ToolCallBehavior.AutoInvokeKernelFunctions,
        };

        var sw = System.Diagnostics.Stopwatch.StartNew();
        var buffer = new StringBuilder();
        int chunks = 0;
        long lastLogMs = 0;

        await foreach (var chunk in chatService.GetStreamingChatMessageContentsAsync(history, settings, kernel, ct))
        {
            var text = chunk.Content;
            if (string.IsNullOrEmpty(text)) continue;
            chunks++;
            buffer.Append(text);
            if (chunks == 1 || chunks % 25 == 0 || sw.ElapsedMilliseconds - lastLogMs > 2000)
            {
                Console.WriteLine($"[{logTag}] chunk #{chunks}, buffered={buffer.Length} chars, elapsed={sw.ElapsedMilliseconds}ms");
                lastLogMs = sw.ElapsedMilliseconds;
            }
        }
        Console.WriteLine($"[{logTag}] DONE: {chunks} chunks, {buffer.Length} chars, {sw.ElapsedMilliseconds}ms");
        return buffer.ToString();
    }

    // === Attachment handling ===
    // Called RIGHT AFTER history.AddUserMessage(prompt). If the request has no
    // attachments this is a no-op. Otherwise it rewrites the just-added user
    // message into a multimodal ChatMessageContentItemCollection that carries:
    //   - text item: the original prompt + any inlined text-file contents
    //   - image item(s): each binary attachment as a data-URI ImageContent
    //
    // We rewrite the LAST message (instead of building it inline at the caller)
    // because both /chat and /conversations/{id}/chat construct the history
    // through `history.AddUserMessage(prompt)` and we want a single seam.
    //
    // The OpenAI connector in SK 1.76 serializes ImageContent as `image_url` with
    // a `data:{mime};base64,{data}` URL. Most upstream proxies accept the
    // same format for application/pdf as for image/*, so we don't branch by mime.
    public static void AttachToLastUserMessage(ChatHistory history, ChatAttachment[]? attachments)
    {
        if (attachments is null || attachments.Length == 0) return;
        if (history.Count == 0) return;
        var last = history[^1];
        if (last.Role != AuthorRole.User) return;

        var originalText = last.Content ?? "";
        var textBuilder = new StringBuilder(originalText);

        var items = new ChatMessageContentItemCollection();
        var mediaCount = 0;
        foreach (var att in attachments)
        {
            if (att is null || string.IsNullOrEmpty(att.Mime) || string.IsNullOrEmpty(att.Data))
                continue;

            if (att.Mime.StartsWith("text/", StringComparison.OrdinalIgnoreCase)
                || att.Mime.Equals("application/json", StringComparison.OrdinalIgnoreCase))
            {
                string decoded;
                try { decoded = Encoding.UTF8.GetString(Convert.FromBase64String(att.Data)); }
                catch { continue; }
                textBuilder.Append("\n\n[Uploaded file: ").Append(att.Name).Append("]\n").Append(decoded);
            }
            else
            {
                byte[] bytes;
                try { bytes = Convert.FromBase64String(att.Data); }
                catch { continue; }
                items.Add(new ImageContent(bytes, att.Mime));
                mediaCount++;
                textBuilder.Append("\n\n[Attached: ").Append(att.Name).Append(" (").Append(att.Mime).Append(")]");
            }
        }

        if (mediaCount == 0)
        {
            history[^1] = new ChatMessageContent(AuthorRole.User, textBuilder.ToString());
            return;
        }

        items.Insert(0, new TextContent(textBuilder.ToString()));
        history[^1] = new ChatMessageContent(AuthorRole.User, items);
    }

    // Build the string that goes into Messages.Content for an attachment-bearing
    // user turn. v1 trade-off: we keep the persisted row small by recording only
    // the file metadata, NOT the full inlined text/binary. That means a replay
    // (loading the conversation later) gives the LLM the prompt text + a note
    // "[attached: foo.pdf]" — enough context for the assistant to reference what
    // was attached, but not enough to re-read it. The user can re-attach if needed.
    public static string BuildPersistedContent(string prompt, ChatAttachment[]? attachments)
    {
        if (attachments is null || attachments.Length == 0) return prompt;
        var sb = new StringBuilder(prompt ?? "");
        sb.Append("\n\n[attached: ");
        for (int i = 0; i < attachments.Length; i++)
        {
            if (i > 0) sb.Append(", ");
            sb.Append(attachments[i].Name);
        }
        sb.Append(']');
        return sb.ToString();
    }

    // Streaming variant: yields token chunks as they arrive from the upstream LLM.
    // Used by the SSE endpoints (/chat, /conversations/{id}/chat) so bytes flow to
    // the browser continuously — Cloudflare's idle-timeout (~100s on Free/Pro) won't
    // fire as long as tokens keep coming. Each yielded chunk is a flush trigger
    // upstream of the SSE writer.
    //
    // The persisted-chat endpoint still gets atomic DB writes: it buffers the
    // yielded chunks server-side into its own StringBuilder, then runs the
    // multi-row insert after the stream completes successfully. Cancellation
    // (client disconnect) bubbles up via `ct` and we abandon the half-stream
    // without persisting — same semantics as the buffered path.
    public async IAsyncEnumerable<string> RunChatStreamingAsync(
        Kernel kernel,
        ChatHistory history,
        string logTag,
        Action<UsageInfo>? onUsage = null,
        [EnumeratorCancellation] CancellationToken ct = default)
    {
        var chatService = kernel.GetRequiredService<IChatCompletionService>();
        var settings = new OpenAIPromptExecutionSettings
        {
            ToolCallBehavior = ToolCallBehavior.AutoInvokeKernelFunctions,
        };

        // Snapshot of prompt-side text chars BEFORE we hit the wire. Used only as
        // input to the fallback estimator at the end of the stream. Counts text
        // content (system + user + prior assistant turns); skips multimodal image
        // bytes because tokens-per-byte for images is fixed-tile (~85/tile), not
        // linear with bytes — adding image bytes here would skew the estimate.
        var promptChars = CountTextChars(history);

        var sw = System.Diagnostics.Stopwatch.StartNew();
        int chunks = 0;
        int totalChars = 0;
        long lastLogMs = 0;
        bool gotUsageFromStream = false;

        await foreach (var chunk in chatService.GetStreamingChatMessageContentsAsync(history, settings, kernel, ct))
        {
            // The SK OpenAI connector auto-requests stream_options.include_usage,
            // so a compliant upstream sends a final chunk with empty Content but
            // populated Metadata["Usage"]. Some self-hosted proxies (LiteLLM in
            // certain configs, hosted gateways with custom middleware) STRIP that chunk — see fallback below.
            if (onUsage is not null && chunk.Metadata is not null && !gotUsageFromStream)
            {
                var usage = TryExtractUsage(chunk.Metadata);
                if (usage is not null)
                {
                    gotUsageFromStream = true;
                    onUsage(usage);
                }
            }

            var text = chunk.Content;
            if (string.IsNullOrEmpty(text)) continue;
            chunks++;
            totalChars += text.Length;
            if (chunks == 1 || chunks % 25 == 0 || sw.ElapsedMilliseconds - lastLogMs > 2000)
            {
                Console.WriteLine($"[{logTag}] stream chunk #{chunks}, total={totalChars} chars, elapsed={sw.ElapsedMilliseconds}ms");
                lastLogMs = sw.ElapsedMilliseconds;
            }
            yield return text;
        }
        Console.WriteLine($"[{logTag}] stream DONE: {chunks} chunks, {totalChars} chars, {sw.ElapsedMilliseconds}ms");

        // === Fallback estimator ===
        // If the upstream never emitted a usage chunk we still want the user to see
        // SOME token+context-window badge and admin analytics to record SOMETHING.
        // OpenAI's published rule-of-thumb is ~4 chars/token across English+code
        // (tiktoken cl100k_base). Real spread is roughly 3-5 for natural language
        // and 2-3 for dense JSON, so worst-case error is about ±35% — labelled
        // IsEstimated=true so callers can render "≈" and the column lets admin
        // queries split estimated vs exact totals.
        if (onUsage is not null && !gotUsageFromStream)
        {
            var p = EstimateTokens(promptChars);
            var c = EstimateTokens(totalChars);
            Console.WriteLine($"[{logTag}] upstream returned no usage chunk → ESTIMATED prompt={p} completion={c} (chars/4 heuristic)");
            onUsage(new UsageInfo(p, c, p + c, IsEstimated: true));
        }
    }

    // === Helpers for the fallback estimator ===

    // Sums text chars across the whole ChatHistory. Iterates Items when present
    // (multimodal messages) and falls back to .Content for plain-text messages,
    // which avoids double-counting since SK's ChatMessageContent.Content surfaces
    // the first TextContent from Items if any.
    private static int CountTextChars(ChatHistory history)
    {
        int total = 0;
        foreach (var msg in history)
        {
            if (msg.Items is { Count: > 0 })
            {
                foreach (var item in msg.Items)
                {
                    if (item is TextContent t && !string.IsNullOrEmpty(t.Text))
                        total += t.Text.Length;
                }
            }
            else if (!string.IsNullOrEmpty(msg.Content))
            {
                total += msg.Content.Length;
            }
        }
        return total;
    }

    // ~4 chars/token is OpenAI's documented rule-of-thumb. Min of 1 so an empty
    // completion still records a nonzero row (zeros look like "missing data" in
    // admin queries; 1 token signals "we measured, it was tiny").
    private static int EstimateTokens(int chars) => Math.Max(1, (chars + 3) / 4);

    // Reflective usage extractor. SK 1.76's OpenAI connector packs ChatTokenUsage
    // (from the OpenAI SDK) under Metadata["Usage"], with properties
    // InputTokenCount / OutputTokenCount / TotalTokenCount. We use reflection so
    // a future SK upgrade renaming the key or the property names doesn't silently
    // break the token badge — we just fall back to null and the UI hides it.
    private static UsageInfo? TryExtractUsage(IReadOnlyDictionary<string, object?> metadata)
    {
        if (!metadata.TryGetValue("Usage", out var raw) || raw is null)
        {
            // Some connector versions lowercase the key.
            if (!metadata.TryGetValue("usage", out raw) || raw is null) return null;
        }
        var type = raw.GetType();
        int? Read(params string[] names)
        {
            foreach (var n in names)
            {
                var p = type.GetProperty(n, BindingFlags.Public | BindingFlags.Instance);
                if (p is null) continue;
                var v = p.GetValue(raw);
                if (v is int i) return i;
                if (v is long l) return (int)l;
                if (v is null) continue;
                if (int.TryParse(v.ToString(), out var parsed)) return parsed;
            }
            return null;
        }
        var prompt     = Read("InputTokenCount", "PromptTokens", "PromptTokenCount", "prompt_tokens");
        var completion = Read("OutputTokenCount", "CompletionTokens", "CompletionTokenCount", "completion_tokens");
        var total      = Read("TotalTokenCount", "TotalTokens", "total_tokens");
        if (prompt is null && completion is null && total is null) return null;
        var p = prompt ?? 0;
        var c = completion ?? 0;
        var t = total ?? (p + c);
        return new UsageInfo(p, c, t);
    }
}