OpenQuery/OpenQuery.cs

using System.Text;
using System.Text.Json;
using System.Text.RegularExpressions;
using OpenQuery.Models;
using OpenQuery.Services;
using OpenQuery.Tools;

namespace OpenQuery;

public class OpenQueryApp
{
    private readonly OpenRouterClient _client;
    private readonly SearchTool _searchTool;
    private readonly string _model;

    public OpenQueryApp(
        OpenRouterClient client,
        SearchTool searchTool,
        string model)
    {
        _client = client;
        _searchTool = searchTool;
        _model = model;
    }

    public async Task RunAsync(OpenQueryOptions options)
    {
        using var reporter = new StatusReporter(options.Verbose);
        reporter.StartSpinner();

        var queries = new List<string> { options.Question };

        if (options.Queries > 1)
        {
            if (options.Verbose)
            {
                reporter.WriteLine($"[Generating {options.Queries} search queries based on your question...]");
            }
            else
            {
                reporter.UpdateStatus("Generating search queries...");
            }

            var queryGenMessages = new List<Message>
            {
                new Message("system", """
                                      You are an expert researcher. The user will ask a question. Your task is to generate optimal search queries to gather comprehensive information to answer this question.

                                      Instructions:
                                      1. Break down complex questions into diverse search queries.
                                      2. Use synonyms and alternative phrasing to capture different sources.
                                      3. Target different aspects of the question (e.g., specific entities, mechanisms, pros/cons, historical context).

                                      Examples:
                                      User: "What are the environmental impacts of electric cars compared to gas cars?"
                                      Output: ["environmental impact of electric cars", "gas vs electric car carbon footprint", "EV battery production environmental cost", "lifecycle emissions electric vs gas vehicles"]

                                      User: "How does the mRNA vaccine technology work?"
                                      Output: ["how mRNA vaccines work", "mechanism of mRNA vaccination", "mRNA vaccine technology explained", "history of mRNA vaccines"]

                                      CRITICAL: Your output MUST strictly be a valid JSON array of strings. Do not include any markdown formatting (like ```json), explanations, preambles, or other text. Just the raw JSON array.
                                      """),
                new Message("user", $"Generate {options.Queries} distinct search queries for this question:\n{options.Question}")
            };

            try
            {
                var request = new ChatCompletionRequest(_model, queryGenMessages);
                var response = await _client.CompleteAsync(request);
                var content = response.Choices.FirstOrDefault()?.Message.Content;

                if (!string.IsNullOrEmpty(content))
                {
                    content = Regex.Replace(content, @"```json\s*|\s*```", "").Trim();

                    var generatedQueries = JsonSerializer.Deserialize(content, AppJsonContext.Default.ListString);
                    if (generatedQueries != null && generatedQueries.Count > 0)
                    {
                        queries = generatedQueries;
                        if (options.Verbose)
                        {
                            reporter.WriteLine($"[Generated queries: {string.Join(", ", queries)}]");
                        }
                    }
                }
            }
            catch (Exception ex)
            {
                if (options.Verbose)
                {
                    reporter.WriteLine($"[Failed to generate queries, falling back to original question. Error: {ex.Message}]");
                }
            }
        }

        reporter.UpdateStatus("Searching web...");
        var searchResult = await _searchTool.ExecuteAsync(
            options.Question,
            queries,
            options.Results,
            options.Chunks,
            (progress) => {
                if (options.Verbose)
                {
                    reporter.WriteLine(progress);
                }
                else
                {
                    // Parse progress messages for compact mode
                    if (progress.StartsWith("[Fetching article") && progress.Contains("/"))
                    {
                        // Extract "X/Y" from "[Fetching article X/Y: domain]"
                        var match = Regex.Match(progress, @"\[(\d+)/(\d+)");
                        if (match.Success)
                        {
                            reporter.UpdateStatus($"Fetching articles {match.Groups[1].Value}/{match.Groups[2].Value}...");
                        }
                    }
                    else if (progress.Contains("embeddings"))
                    {
                        reporter.UpdateStatus("Processing embeddings...");
                    }
                }
            },
            options.Verbose);

        if (!options.Verbose)
        {
            reporter.UpdateStatus("Asking AI...");
        }
        else
        {
            reporter.ClearStatus();
            Console.WriteLine();
        }

        var systemPrompt = "You are a helpful AI assistant. Answer the user's question in depth, based on the provided context. Be precise and accurate. You can mention sources or citations.";
        if (options.Short)
            systemPrompt += " Give a very short concise answer.";
        if (options.Long)
            systemPrompt += " Give a long elaborate detailed answer.";

        var messages = new List<Message>
        {
            new Message("system", systemPrompt),
            new Message("user", $"Context:\n{searchResult}\n\nQuestion: {options.Question}")
        };

        var requestStream = new ChatCompletionRequest(_model, messages);

        var assistantResponse = new StringBuilder();
        var isFirstChunk = true;

        try
        {
            using var streamCts = new CancellationTokenSource();
            await foreach (var chunk in _client.StreamAsync(requestStream, streamCts.Token))
            {
                if (chunk.TextDelta == null) continue;
                if (isFirstChunk)
                {
                    reporter.StopSpinner();
                    if (!options.Verbose)
                    {
                        reporter.ClearStatus();
                    }
                    else
                    {
                        Console.Write("Assistant: ");
                    }
                    isFirstChunk = false;
                }
                Console.Write(chunk.TextDelta);
                assistantResponse.Append(chunk.TextDelta);
            }
        }
        finally
        {
            reporter.StopSpinner();
        }

        Console.WriteLine();
    }
}