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;
    private static readonly char[] Function = ['|', '/', '-', '\\'];

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

    public async Task RunAsync(OpenQueryOptions options)
    {
        var queries = new List<string> { options.Question };

        if (options.Queries > 1)
        {
            Console.WriteLine($"[Generating {options.Queries} search queries based on your question...]");
            
            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;
                        Console.WriteLine($"[Generated queries: {string.Join(", ", queries)}]");
                    }
                }
            }
            catch (Exception ex)
            {
                Console.WriteLine($"[Failed to generate queries, falling back to original question. Error: {ex.Message}]");
            }
        }

        var searchResult = await _searchTool.ExecuteAsync(options.Question, queries, options.Results, options.Chunks, msg => Console.WriteLine(msg));
        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;

        Console.Write("[Sending request to AI model...] ");
        
        using var cts = new CancellationTokenSource();
        var spinnerTask = Task.Run(async () =>
        {
            var spinner = Function;
            var index = 0;
            while (cts is { Token.IsCancellationRequested: false })
            {
                if (Console.CursorLeft > 0)
                {
                    Console.Write(spinner[index++ % spinner.Length]);
                    Console.SetCursorPosition(Console.CursorLeft - 1, Console.CursorTop);
                }
                try
                {
                    await Task.Delay(100, cts.Token);
                }
                catch (TaskCanceledException)
                {
                    break;
                }
            }
        }, cts.Token);
        
        try
        {
            await foreach (var chunk in _client.StreamAsync(requestStream, cts.Token))
            {
                if (chunk.TextDelta == null) continue;
                if (isFirstChunk)
                {
                    await cts.CancelAsync();
                    await spinnerTask;
                    Console.WriteLine();
                    Console.Write("Assistant: ");
                    isFirstChunk = false;
                }
                Console.Write(chunk.TextDelta);
                assistantResponse.Append(chunk.TextDelta);
            }
        }
        finally
        {
            if (!cts.IsCancellationRequested)
            {
                await cts.CancelAsync();
            }
        }

        Console.WriteLine();
    }
}