Merge pull request #1039 from ionite34/prompt-amplify-n-stuff

Prompt amplify n stuff

Author: mohnjiles

.husky/task-runner.json

@@ -19,6 +19,12 @@
          "group": "generate-openapi",
          "command": "dotnet",
          "args": ["refitter", "--settings-file", "./StabilityMatrix.Core/Api/LykosAuthApi/.refitter"]
+      },
+      {
+         "name": "Run refitter for PromptGenApi",
+         "group": "generate-promptgen-openapi",
+         "command": "dotnet",
+         "args": ["refitter", "--settings-file", "./StabilityMatrix.Core/Api/PromptGen/.refitter"]
       }
    ]
 }

StabilityMatrix.Avalonia/App.axaml.cs

@@ -59,6 +59,7 @@
 using StabilityMatrix.Avalonia.Views;
 using StabilityMatrix.Core.Api;
 using StabilityMatrix.Core.Api.LykosAuthApi;
+using StabilityMatrix.Core.Api.PromptGenApi;
 using StabilityMatrix.Core.Attributes;
 using StabilityMatrix.Core.Converters.Json;
 using StabilityMatrix.Core.Database;
@@ -139,6 +140,14 @@ public sealed class App : Application
 #else
     public const string LykosAccountApiBaseUrl = "https://account.lykos.ai/";
 #endif
+#if DEBUG
+    // ReSharper disable twice LocalizableElement
+    // ReSharper disable once ConditionalAccessQualifierIsNonNullableAccordingToAPIContract
+    public static string PromptGenApiBaseUrl =>
+        Config?["PromptGenApiBaseUrl"] ?? "https://promptgen.lykos.ai/api";
+#else
+    public const string PromptGenApiBaseUrl = "https://promptgen.lykos.ai/api";
+#endif
 
     // ReSharper disable once MemberCanBePrivate.Global
     public IClassicDesktopStyleApplicationLifetime? DesktopLifetime =>
@@ -375,7 +384,7 @@ internal static void ConfigurePageViewModels(IServiceCollection services)
                 new MainWindowViewModel(
                     provider.GetRequiredService<ISettingsManager>(),
                     provider.GetRequiredService<IDiscordRichPresenceService>(),
-                    provider.GetRequiredService<ServiceManager<ViewModelBase>>(),
+                    provider.GetRequiredService<IServiceManager<ViewModelBase>>(),
                     provider.GetRequiredService<ITrackedDownloadService>(),
                     provider.GetRequiredService<IModelIndexService>(),
                     provider.GetRequiredService<Lazy<IModelDownloadLinkHandler>>(),
@@ -641,6 +650,21 @@ internal static IServiceCollection ConfigureServices()
                     new TokenAuthHeaderHandler(serviceProvider.GetRequiredService<LykosAuthTokenProvider>())
             );
 
+        services
+            .AddRefitClient<IPromptGenApi>(defaultRefitSettings)
+            .ConfigureHttpClient(c =>
+            {
+                c.BaseAddress = new Uri(PromptGenApiBaseUrl);
+                c.Timeout = TimeSpan.FromHours(1);
+                c.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", "");
+            })
+            .ConfigurePrimaryHttpMessageHandler(() => new HttpClientHandler { AllowAutoRedirect = false })
+            .AddPolicyHandler(retryPolicy)
+            .AddHttpMessageHandler(
+                serviceProvider =>
+                    new TokenAuthHeaderHandler(serviceProvider.GetRequiredService<LykosAuthTokenProvider>())
+            );
+
         services
             .AddRefitClient<ILykosAnalyticsApi>(defaultRefitSettings)
             .ConfigureHttpClient(c =>

StabilityMatrix.Avalonia/Controls/BetterDownloadableComboBox.cs

@@ -54,7 +54,7 @@ private static async Task PromptDownloadAsync(IDownloadableResource downloadable
         if (downloadable.DownloadableResource is not { } resource)
             return;
 
-        var vmFactory = App.Services.GetRequiredService<ServiceManager<ViewModelBase>>();
+        var vmFactory = App.Services.GetRequiredService<IServiceManager<ViewModelBase>>();
         var confirmDialog = vmFactory.Get<DownloadResourceViewModel>();
         confirmDialog.Resource = resource;
         confirmDialog.FileName = resource.FileName;

StabilityMatrix.Avalonia/Controls/FADownloadableComboBox.cs

@@ -55,7 +55,7 @@ private static async Task PromptDownloadAsync(IDownloadableResource downloadable
         if (downloadable.DownloadableResource is not { } resource)
             return;
 
-        var vmFactory = App.Services.GetRequiredService<ServiceManager<ViewModelBase>>();
+        var vmFactory = App.Services.GetRequiredService<IServiceManager<ViewModelBase>>();
         var confirmDialog = vmFactory.Get<DownloadResourceViewModel>();
         confirmDialog.Resource = resource;
         confirmDialog.FileName = resource.FileName;

StabilityMatrix.Avalonia/Controls/Inference/PromptCard.axaml

@@ -164,7 +164,7 @@ public static void Initialize()
 
         Services = services.BuildServiceProvider();
 
-        var dialogFactory = Services.GetRequiredService<ServiceManager<ViewModelBase>>();
+        var dialogFactory = Services.GetRequiredService<IServiceManager<ViewModelBase>>();
         var settingsManager = Services.GetRequiredService<ISettingsManager>();
         var downloadService = Services.GetRequiredService<IDownloadService>();
         var modelFinder = Services.GetRequiredService<ModelFinder>();
@@ -438,8 +438,8 @@ public static void Initialize()
     [NotNull]
     public static UpdateViewModel? UpdateViewModel { get; private set; }
 
-    public static ServiceManager<ViewModelBase> DialogFactory =>
-        Services.GetRequiredService<ServiceManager<ViewModelBase>>();
+    public static IServiceManager<ViewModelBase> DialogFactory =>
+        Services.GetRequiredService<IServiceManager<ViewModelBase>>();
 
     public static MainWindowViewModel MainWindowViewModel =>
         Services.GetRequiredService<MainWindowViewModel>();

StabilityMatrix.Avalonia/DesignData/DesignData.cs

@@ -21,7 +21,7 @@ public MockLaunchPageViewModel(
         IPyRunner pyRunner,
         INotificationService notificationService,
         ISharedFolders sharedFolders,
-        ServiceManager<ViewModelBase> dialogFactory
+        IServiceManager<ViewModelBase> dialogFactory
     )
         : base(
             logger,
@@ -34,34 +34,35 @@ ServiceManager<ViewModelBase> dialogFactory
         ) { }
 
     public override BasePackage? SelectedBasePackage =>
-        SelectedPackage?.PackageName != "dank-diffusion" ? base.SelectedBasePackage :
-        new DankDiffusion(null!, null!, null!, null!);
-    
+        SelectedPackage?.PackageName != "dank-diffusion"
+            ? base.SelectedBasePackage
+            : new DankDiffusion(null!, null!, null!, null!);
+
     protected override Task LaunchImpl(string? command)
     {
         IsLaunchTeachingTipsOpen = false;
 
-        RunningPackage = new PackagePair(
-            null!, 
-            new DankDiffusion(null!, null!, null!, null!));
-        
-        Console.Document.Insert(0, 
+        RunningPackage = new PackagePair(null!, new DankDiffusion(null!, null!, null!, null!));
+
+        Console.Document.Insert(
+            0,
             """
             Python 3.10.11 (tags/v3.10.11:7d4cc5a, Apr  5 2023, 00:38:17) [MSC v.1929 64 bit (AMD64)]
             Version: 1.5.0
             Commit hash: <none>
 
             Fetching updates for midas...
             Checking out commit for midas with hash: 2e42b7f...
-            """);
-        
+            """
+        );
+
         return Task.CompletedTask;
     }
 
     public override Task Stop()
     {
         RunningPackage = null;
-        
+
         return Task.CompletedTask;
     }
 }

StabilityMatrix.Avalonia/DesignData/MockLaunchPageViewModel.cs

@@ -0,0 +1,64 @@
+using System.ComponentModel;
+using Avalonia.Controls;
+using StabilityMatrix.Avalonia.Controls;
+using StabilityMatrix.Avalonia.Services;
+using StabilityMatrix.Core.Attributes;
+
+namespace StabilityMatrix.Avalonia.Extensions;
+
+[Localizable(false)]
+public static class ServiceManagerExtensions
+{
+    /// <summary>
+    /// Get a view model instance, set as DataContext of its View, and return
+    /// a BetterContentDialog with that View as its Content
+    /// </summary>
+    public static BetterContentDialog GetDialog<TService>(this IServiceManager<TService> serviceManager)
+    {
+        var instance = serviceManager.Get<TService>()!;
+
+        if (
+            Attribute.GetCustomAttribute(instance.GetType(), typeof(ViewAttribute))
+            is not ViewAttribute viewAttr
+        )
+        {
+            throw new InvalidOperationException($"View not found for {instance.GetType().FullName}");
+        }
+
+        if (Activator.CreateInstance(viewAttr.ViewType) is not Control view)
+        {
+            throw new NullReferenceException($"Unable to create instance for {instance.GetType().FullName}");
+        }
+
+        return new BetterContentDialog { Content = view };
+    }
+
+    /// <summary>
+    /// Get a view model instance with initializer, set as DataContext of its View, and return
+    /// a BetterContentDialog with that View as its Content
+    /// </summary>
+    public static BetterContentDialog GetDialog<TService>(
+        this IServiceManager<TService> serviceManager,
+        Action<TService> initializer
+    )
+    {
+        var instance = serviceManager.Get(initializer)!;
+
+        if (
+            Attribute.GetCustomAttribute(instance.GetType(), typeof(ViewAttribute))
+            is not ViewAttribute viewAttr
+        )
+        {
+            throw new InvalidOperationException($"View not found for {instance.GetType().FullName}");
+        }
+
+        if (Activator.CreateInstance(viewAttr.ViewType) is not Control view)
+        {
+            throw new NullReferenceException($"Unable to create instance for {instance.GetType().FullName}");
+        }
+
+        view.DataContext = instance;
+
+        return new BetterContentDialog { Content = view };
+    }
+}

StabilityMatrix.Avalonia/Extensions/ServiceManagerExtensions.cs

@@ -49,10 +49,21 @@ public static void AddServicesByAttributesSourceGen(IServiceCollection services)
         services.AddStabilityMatrixAvalonia();
     }
 
+    private static bool IsScoped(IServiceProvider serviceProvider)
+    {
+        // Hacky check for if service provider is scoped
+        var typeName = serviceProvider.GetType().Name;
+        if (typeName == "ServiceProviderEngineScope" || typeName.Contains("Scope"))
+        {
+            return true;
+        }
+        return false;
+    }
+
     /// <summary>
-    /// Adds a <see cref="ServiceManager{T}"/> to the <see cref="IServiceCollection"/>.
+    /// Adds a <see cref="IServiceManager{T}"/> to the <see cref="IServiceCollection"/>.
     /// </summary>
-    /// <param name="services">The <see cref="IServiceCollection"/> to which the <see cref="ServiceManager{T}"/> will be added.</param>
+    /// <param name="services">The <see cref="IServiceCollection"/> to which the <see cref="IServiceManager{T}"/> will be added.</param>
     /// <param name="serviceFilter">An optional filter for the services.</param>
     /// <typeparam name="TService">The base type of the services.</typeparam>
     /// <exception cref="InvalidOperationException"></exception>
@@ -61,13 +72,14 @@ public static IServiceCollection AddServiceManagerWithCurrentCollectionServices<
         Func<ServiceDescriptor, bool>? serviceFilter = null
     )
     {
-        return services.AddSingleton<ServiceManager<TService>>(provider =>
+        // Register main service manager as singleton
+        services.AddSingleton<ServiceManager<TService>>(provider =>
         {
             using var _ = CodeTimer.StartDebug(
                 callerName: $"{nameof(AddServiceManagerWithCurrentCollectionServices)}<{typeof(TService)}>"
             );
 
-            var serviceManager = new ServiceManager<TService>();
+            var serviceManager = new ServiceManager<TService>(provider);
 
             // Get registered services that are assignable to TService
             var serviceDescriptors = services.Where(s => s.ServiceType.IsAssignableTo(typeof(TService)));
@@ -84,10 +96,36 @@ public static IServiceCollection AddServiceManagerWithCurrentCollectionServices<
                 Debug.Assert(type is not null, "type is not null");
                 Debug.Assert(type.IsAssignableTo(typeof(TService)), "type is assignable to TService");
 
-                serviceManager.Register(type, () => (TService)provider.GetRequiredService(type));
+                if (service.Lifetime == ServiceLifetime.Scoped)
+                {
+                    serviceManager.RegisterScoped(type, sp => (TService)sp.GetRequiredService(type));
+                }
+                else
+                {
+                    serviceManager.Register(type, () => (TService)provider.GetRequiredService(type));
+                }
             }
 
             return serviceManager;
         });
+
+        // Register scoped for interface
+        services.AddScoped<IServiceManager<TService>>(provider =>
+        {
+            var rootServiceManager = provider.GetRequiredService<ServiceManager<TService>>();
+
+            // For non scoped, just return the singleton
+            if (!IsScoped(provider))
+            {
+                return rootServiceManager;
+            }
+
+            // For scoped, create a new instance using root and provider
+            var scopedServiceManager = new ScopedServiceManager<TService>(rootServiceManager, provider);
+
+            return scopedServiceManager;
+        });
+
+        return services;
     }
 }