Bare-Metal AOT Ollama Shell

fast-ollama.cs

using System;
using System.Buffers;
using System.Collections.Generic;
using System.IO;
using System.Net.Http;
using System.Runtime.CompilerServices;
using System.Text;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;

// ==========================================
// TOP-LEVEL ENTRY POINT (INTERACTIVE LOOP)
// ==========================================

using var client = new BareMetalOllamaClient();
string modelName = "gpt-oss:120b-cloud";

Console.WriteLine($"--- Bare-Metal AOT Ollama Shell ({modelName}) ---");
Console.WriteLine("Type your prompt and press Enter. Type 'exit' to quit.\n");

while (true)
{
    Console.ForegroundColor = ConsoleColor.Green;
    Console.Write("User > ");
    Console.ResetColor();

    string? userQuery = Console.ReadLine();
    
    // Breakout conditions
    if (string.IsNullOrWhiteSpace(userQuery)) continue;
    if (userQuery.Equals("exit", StringComparison.OrdinalIgnoreCase)) break;

    Console.ForegroundColor = ConsoleColor.Cyan;
    Console.Write("Ollama > ");
    Console.ResetColor();

    using var cts = new CancellationTokenSource();
    
    try
    {
        // Streams tokens continuously using pooled memory arrays
        await foreach (var token in client.StreamChatAsync(
            prompt: userQuery, 
            model: modelName, 
            cancellationToken: cts.Token))
        {
            Console.Write(token);
        }
    }
    catch (HttpRequestException ex) when (ex.StatusCode == System.Net.HttpStatusCode.BadRequest)
    {
        Console.WriteLine("\n[Error 400]: Payload structure rejected by Ollama.");
    }
    catch (HttpRequestException ex) when (ex.StatusCode == System.Net.HttpStatusCode.NotFound)
    {
        Console.WriteLine($"\n[Error 404]: Could not find model '{modelName}'.");
    }
    catch (Exception ex)
    {
        Console.WriteLine($"\n[Error]: {ex.Message}");
    }

    Console.WriteLine("\n"); // Clear turn separation
}


// ==========================================
// BARE-METAL OLLAMA CLIENT (TRUE ZERO ALLOCATION)
// ==========================================
public sealed class BareMetalOllamaClient : IDisposable
{
    private readonly HttpClient _httpClient;
    
    private static ReadOnlySpan<byte> MessagePropName => "message"u8;
    private static ReadOnlySpan<byte> ContentPropName => "content"u8;
    private static ReadOnlySpan<byte> DonePropName => "done"u8;
    private static byte LineFeed => (byte)'\n';

    public BareMetalOllamaClient(string baseUrl = "http://localhost:11434")
    {
        var handler = new SocketsHttpHandler
        {
            PooledConnectionLifetime = TimeSpan.FromMinutes(15),
            PooledConnectionIdleTimeout = TimeSpan.FromMinutes(2),
            MaxConnectionsPerServer = 100,
            EnableMultipleHttp2Connections = true
        };

        _httpClient = new HttpClient(handler) { BaseAddress = new Uri(baseUrl) };
    }

    public async IAsyncEnumerable<string> StreamChatAsync(
        string prompt, 
        string model, 
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        using var memoryStream = new MemoryStream();
        using (var writer = new Utf8JsonWriter(memoryStream))
        {
            writer.WriteStartObject();
            writer.WriteString("model"u8, model);
            
            writer.WriteStartArray("messages"u8);
            
            writer.WriteStartObject();
            writer.WriteString("role"u8, "system");
            writer.WriteString("content"u8, "You are a helpful assistant. You must answer directly and succinctly.");
            writer.WriteEndObject();

            writer.WriteStartObject();
            writer.WriteString("role"u8, "user");
            writer.WriteString("content"u8, prompt);
            writer.WriteEndObject();
            
            writer.WriteEndArray();
            
            writer.WriteStartObject("options"u8);
            writer.WriteNumber("temperature"u8, 0.0f); // Lock creativity for speed and deterministic logic
            writer.WriteEndObject();

            writer.WriteBoolean("stream"u8, true);
            writer.WriteEndObject();
        }

        memoryStream.Position = 0;

        var request = new HttpRequestMessage(HttpMethod.Post, "/api/chat")
        {
            Content = new StreamContent(memoryStream),
            Version = System.Net.HttpVersion.Version20 
        };
        request.Content.Headers.ContentType = new System.Net.Http.Headers.MediaTypeHeaderValue("application/json");

        var response = await _httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, cancellationToken).ConfigureAwait(false);
        response.EnsureSuccessStatusCode();

        using var responseStream = await response.Content.ReadAsStreamAsync(cancellationToken).ConfigureAwait(false);
        
        // Rent a reusable byte block from the system array pool. Eliminates heap allocation churn.
        byte[] buffer = ArrayPool<byte>.Shared.Rent(4096);
        int bufferOffset = 0;

        try
        {
            while (true)
            {
                int bytesRead = await responseStream.ReadAsync(buffer.AsMemory(bufferOffset, buffer.Length - bufferOffset), cancellationToken).ConfigureAwait(false);
                if (bytesRead == 0) break;

                int totalBytes = bufferOffset + bytesRead;
                int scanIndex = 0;

                // Slice lines straight from raw byte segments without allocating intermediate line strings
                while (scanIndex < totalBytes)
                {
                    int lfIndex = Array.IndexOf(buffer, LineFeed, scanIndex, totalBytes - scanIndex);
                    if (lfIndex == -1) break; // Token stream is cut off across buffers; loop back to stream more socket data

                    int lineLength = lfIndex - scanIndex;
                    if (lineLength > 0)
                    {
                        ReadOnlySpan<byte> rawJsonLine = new ReadOnlySpan<byte>(buffer, scanIndex, lineLength);
                        string? token = ParseRawBytesFast(rawJsonLine);
                        if (token != null) yield return token;
                    }

                    scanIndex = lfIndex + 1;
                }

                // Carry over incomplete trailing line segments to the start of the next read loop pass
                if (scanIndex < totalBytes)
                {
                    bufferOffset = totalBytes - scanIndex;
                    Array.Copy(buffer, scanIndex, buffer, 0, bufferOffset);
                }
                else
                {
                    bufferOffset = 0;
                }
            }
        }
        finally
        {
            // Always return the array back to the kernel pool
            ArrayPool<byte>.Shared.Return(buffer);
        }
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private static string? ParseRawBytesFast(ReadOnlySpan<byte> jsonBytes)
    {
        // Parses completely on the thread stack via ReadOnlySpans. Completely zero GC overhead.
        var jsonReader = new Utf8JsonReader(jsonBytes);
        string? extractedToken = null;

        while (jsonReader.Read())
        {
            if (jsonReader.TokenType == JsonTokenType.PropertyName)
            {
                if (jsonReader.ValueTextEquals(MessagePropName))
                {
                    while (jsonReader.Read() && jsonReader.TokenType != JsonTokenType.EndObject)
                    {
                        if (jsonReader.TokenType == JsonTokenType.PropertyName && jsonReader.ValueTextEquals(ContentPropName))
                        {
                            jsonReader.Read();
                            extractedToken = jsonReader.GetString();
                            break;
                        }
                    }
                }
                else if (jsonReader.ValueTextEquals(DonePropName))
                {
                    jsonReader.Read();
                    if (jsonReader.GetBoolean()) break;
                }
            }
        }
        return extractedToken;
    }

    public void Dispose() => _httpClient.Dispose();
}