Implement conversation context and streaming with OllamaSharp

shell/agents/AIShell.Ollama.Agent/AIShell.Ollama.Agent.csproj

@@ -3,7 +3,8 @@
   <PropertyGroup>
     <TargetFramework>net8.0</TargetFramework>
     <ImplicitUsings>enable</ImplicitUsings>
-    <SuppressNETCoreSdkPreviewMessage>true</SuppressNETCoreSdkPreviewMessage>
+	<SuppressNETCoreSdkPreviewMessage>true</SuppressNETCoreSdkPreviewMessage>
+	<CopyLocalLockFileAssemblies>true</CopyLocalLockFileAssemblies>
 
     <!-- Disable deps.json generation -->
     <GenerateDependencyFile>false</GenerateDependencyFile>
@@ -16,6 +17,10 @@
   </PropertyGroup>
 
   <ItemGroup>
+    <PackageReference Include="OllamaSharp" Version="4.0.7" />
+  </ItemGroup>
+
+	<ItemGroup>
     <ProjectReference Include="..\..\AIShell.Abstraction\AIShell.Abstraction.csproj">
       <!-- Disable copying AIShell.Abstraction.dll to output folder -->
       <Private>false</Private>

shell/agents/AIShell.Ollama.Agent/OllamaAgent.cs

@@ -1,10 +1,32 @@
 using System.Diagnostics;
+using System.Text;
+using System.Text.Json;
 using AIShell.Abstraction;
+using OllamaSharp;
+using OllamaSharp.Models;
 
 namespace AIShell.Ollama.Agent;
 
 public sealed class OllamaAgent : ILLMAgent
 {
+    private bool _reloadSettings;
+    private bool _isDisposed;
+    private string _configRoot;
+    private Settings _settings;
+    private OllamaApiClient _client;
+    private GenerateRequest _request;
+    private FileSystemWatcher _watcher;
+
+    /// <summary>
+    /// The name of setting file
+    /// </summary>
+    private const string SettingFileName = "ollama.config.json";
+
+    /// <summary>
+    /// Gets the settings.
+    /// </summary>
+    internal Settings Settings => _settings;
+
     /// <summary>
     /// The name of the agent
     /// </summary>
@@ -13,7 +35,7 @@ public sealed class OllamaAgent : ILLMAgent
     /// <summary>
     /// The description of the agent to be shown at start up
     /// </summary>
-    public string Description => "This is an AI assistant to interact with a language model running locally by utilizing the Ollama CLI tool. Be sure to follow all prerequisites in aka.ms/aish/ollama";
+    public string Description => "This is an AI assistant to interact with a language model running locally by utilizing the Ollama CLI tool. Be sure to follow all prerequisites in https://github.com/PowerShell/AIShell/tree/main/shell/agents/AIShell.Ollama.Agent";
 
     /// <summary>
     /// This is the company added to /like and /dislike verbiage for who the telemetry helps.
@@ -30,19 +52,25 @@ public sealed class OllamaAgent : ILLMAgent
     /// <summary>
     /// These are any optional legal/additional information links you want to provide at start up
     /// </summary>
-    public Dictionary<string, string> LegalLinks { private set; get; }
-
-    /// <summary>
-    /// This is the chat service to call the API from
-    /// </summary>
-    private OllamaChatService _chatService;
+    public Dictionary<string, string> LegalLinks { private set; get; } = new(StringComparer.OrdinalIgnoreCase)
+    {
+        ["Ollama Docs"] = "https://github.com/ollama/ollama",
+        ["Prerequisites"] = "https://github.com/PowerShell/AIShell/tree/main/shell/agents/AIShell.Ollama.Agent"
+    };
 
     /// <summary>
     /// Dispose method to clean up the unmanaged resource of the chatService
     /// </summary>
     public void Dispose()
     {
-        _chatService?.Dispose();
+        if (_isDisposed)
+        {
+            return;
+        }
+
+        GC.SuppressFinalize(this);
+        _watcher.Dispose();
+        _isDisposed = true;
     }
 
     /// <summary>
@@ -51,13 +79,30 @@ public void Dispose()
     /// <param name="config">Agent configuration for any configuration file and other settings</param>
     public void Initialize(AgentConfig config)
     {
-        _chatService = new OllamaChatService();
+        _configRoot = config.ConfigurationRoot;
+
+        SettingFile = Path.Combine(_configRoot, SettingFileName);
+        _settings = ReadSettings();
+
+        if (_settings is null)
+        {
+            // Create the setting file with examples to serve as a template for user to update.
+            NewExampleSettingFile();
+            _settings = ReadSettings();
+        }
+
+        _request = new GenerateRequest()
+        {
+            Model = _settings.Model,
+            Stream = _settings.Stream
+        };
 
-        LegalLinks = new(StringComparer.OrdinalIgnoreCase)
+        _watcher = new FileSystemWatcher(_configRoot, SettingFileName)
         {
-            ["Ollama Docs"] = "https://github.com/ollama/ollama",
-            ["Prerequisites"] = "https://aka.ms/ollama/readme"
+            NotifyFilter = NotifyFilters.LastWrite,
+            EnableRaisingEvents = true,
         };
+        _watcher.Changed += OnSettingFileChange;
     }
 
     /// <summary>
@@ -68,7 +113,7 @@ public void Initialize(AgentConfig config)
     /// <summary>
     /// Gets the path to the setting file of the agent.
     /// </summary>
-    public string SettingFile { private set; get; } = null;
+    public string SettingFile { private set; get; }
 
     /// <summary>
     /// Gets a value indicating whether the agent accepts a specific user action feedback.
@@ -87,7 +132,19 @@ public void OnUserAction(UserActionPayload actionPayload) {}
     /// Refresh the current chat by starting a new chat session.
     /// This method allows an agent to reset chat states, interact with user for authentication, print welcome message, and more.
     /// </summary>
-    public Task RefreshChatAsync(IShell shell, bool force) => Task.CompletedTask;
+    public Task RefreshChatAsync(IShell shell, bool force)
+    {
+        if (force)
+        {
+            // Reload the setting file if needed.
+            ReloadSettings();
+
+            // Reset context
+            _request.Context = null;
+        }
+
+        return Task.CompletedTask;
+    }
 
     /// <summary>
     /// Main chat function that takes the users input and passes it to the LLM and renders it.
@@ -103,23 +160,169 @@ public async Task<bool> ChatAsync(string input, IShell shell)
         // get the cancellation token
         CancellationToken token = shell.CancellationToken;
 
+        // Reload the setting file if needed.
+        ReloadSettings();
+
         if (Process.GetProcessesByName("ollama").Length is 0)
         {
-            host.RenderFullResponse("Please be sure the Ollama is installed and server is running. Check all the prerequisites in the README of this agent are met.");
+            host.MarkupWarningLine($"[[{Name}]]: Please be sure the Ollama is installed and server is running. Check all the prerequisites in the README of this agent are met.");
+            return false;
+        }
+
+        // Self check settings Model and Endpoint
+        if (!SelfCheck(host))
+        {
             return false;
         }
 
-        ResponseData ollamaResponse = await host.RunWithSpinnerAsync(
-            status: "Thinking ...",
-            func: async context => await _chatService.GetChatResponseAsync(context, input, token)
-        ).ConfigureAwait(false);
+        // Update request
+        _request.Prompt = input;
+        _request.Model = _settings.Model;
+        _request.Stream = _settings.Stream;
+
+        // Ollama client is created per chat with reloaded settings
+        _client = new OllamaApiClient(_settings.Endpoint);
+
+        try
+        {
+            if (_request.Stream)
+            {
+                using IStreamRender streamingRender = host.NewStreamRender(token);
+
+                // Last stream response has context value
+                GenerateDoneResponseStream ollamaLastStream = null;
+
+                // Directly process the stream when no spinner is needed
+                await foreach (var ollamaStream in _client.GenerateAsync(_request, token))
+                {
+                    // Update the render
+                    streamingRender.Refresh(ollamaStream.Response);
+                    if (ollamaStream.Done)
+                    {
+                       ollamaLastStream = (GenerateDoneResponseStream)ollamaStream;
+                    }
+                }
+
+                // Update request context
+                _request.Context = ollamaLastStream.Context;
+            }
+            else
+            {
+                // Build single response with spinner
+                var ollamaResponse = await host.RunWithSpinnerAsync(
+                    status: "Thinking ...",
+                    func: async () => { return await _client.GenerateAsync(_request, token).StreamToEndAsync(); }
+                ).ConfigureAwait(false);
 
-        if (ollamaResponse is not null)
+                // Update request context
+                _request.Context = ollamaResponse.Context;
+
+                // Render the full response
+                host.RenderFullResponse(ollamaResponse.Response);
+            }
+        }
+        catch (OperationCanceledException)
         {
-            // render the content
-            host.RenderFullResponse(ollamaResponse.response);
+            // Ignore the cancellation exception.
         }
-
+        catch (HttpRequestException e)
+        {
+            host.WriteErrorLine($"[{Name}]: {e.Message}");
+            host.WriteErrorLine($"[{Name}]: Selected Model    : \"{_settings.Model}\"");
+            host.WriteErrorLine($"[{Name}]: Selected Endpoint : \"{_settings.Endpoint}\"");
+            host.WriteErrorLine($"[{Name}]: Configuration File: \"{SettingFile}\"");
+            return false;
+        }
+
         return true;
     }
+
+
+    private void ReloadSettings()
+    {
+        if (_reloadSettings)
+        {
+            _reloadSettings = false;
+            var settings = ReadSettings();
+            if (settings is null)
+            {
+                return;
+            }
+
+            _settings = settings;
+        }
+    }
+
+    private Settings ReadSettings()
+    {
+        Settings settings = null;
+        FileInfo file = new(SettingFile);
+
+        if (file.Exists)
+        {
+            try
+            {
+                using var stream = file.OpenRead();
+                var data = JsonSerializer.Deserialize(stream, SourceGenerationContext.Default.ConfigData);
+                settings = new Settings(data);
+            }
+            catch (Exception e)
+            {
+                throw new InvalidDataException($"Parsing settings from '{SettingFile}' failed with the following error: {e.Message}", e);
+            }
+        }
+
+        return settings;
+    }
+
+    private void OnSettingFileChange(object sender, FileSystemEventArgs e)
+    {
+        if (e.ChangeType is WatcherChangeTypes.Changed)
+        {
+            _reloadSettings = true;
+        }
+    }
+
+    private bool SelfCheck(IHost host)
+    {
+        var settings = new (string settingValue, string settingName)[]
+        {
+            (_settings.Model, "Model"),
+            (_settings.Endpoint, "Endpoint")
+        };
+
+        foreach (var (settingValue, settingName) in settings)
+        {
+            if (string.IsNullOrWhiteSpace(settingValue))
+            {
+                host.WriteErrorLine($"[{Name}]: {settingName} is undefined in the settings file: \"{SettingFile}\"");
+                return false;
+            }
+        }
+
+        return true;
+    }
+
+    private void NewExampleSettingFile()
+    {
+        string SampleContent = $$"""
+        {
+            // To use Ollama API service:
+            // 1. Install Ollama:
+            //      winget install Ollama.Ollama
+            // 2. Start Ollama API server:
+            //      ollama serve
+            // 3. Install Ollama model:
+            //      ollama pull phi3
+
+            // Declare Ollama model
+            "Model": "phi3",
+            // Declare Ollama endpoint
+            "Endpoint": "http://localhost:11434",
+            // Enable Ollama streaming
+            "Stream": false
+        }
+        """;
+        File.WriteAllText(SettingFile, SampleContent, Encoding.UTF8);
+    }
 }