terra-llm/Assets/LLM/Editor/Client/GeminiApiClient.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using LLM.Editor.Analysis;
using LLM.Editor.Api;
using LLM.Editor.Core;
using LLM.Editor.Data;
using LLM.Editor.Helper;
using UnityEditor;
using UnityEngine.Networking;
using Debug = UnityEngine.Debug;
using Object = UnityEngine.Object;

namespace LLM.Editor.Client
{
    /// <summary>
    /// The client responsible for communicating with the live Google Gemini API.
    /// </summary>
    public class GeminiApiClient : ILlmApiClient
    {
        [Serializable]
        private class CommandResponse { public List<CommandData> commands; }

        private Settings.MCPSettings _settings;
        private string _systemPrompt;
        private bool _isInitialized;

        public async Task SendPrompt(string prompt, List<Object> stagedContext)
        {
            if (!await Initialize()) return;

            var authToken = GetAuthToken();
            if (string.IsNullOrEmpty(authToken))
            {
                Debug.LogError("[GeminiApiClient] Failed to get authentication token.");
                return;
            }
            
            var fullPrompt = BuildInitialPrompt(prompt, stagedContext);
            Debug.Log("[GeminiApiClient] Sending prompt: \n" + fullPrompt);

            var chatHistory = SessionManager.LoadChatHistory();
            chatHistory.Add(new ChatEntry { role = "user", content = fullPrompt });
            
            await SendApiRequest(chatHistory, authToken);
        }
        
        public async Task SendFollowUp(string detailedContext)
        {
            if (!await Initialize()) return;

            var authToken = GetAuthToken();
            if (string.IsNullOrEmpty(authToken))
            {
                Debug.LogError("[GeminiApiClient] Failed to get authentication token.");
                return;
            }
            
            var chatHistory = SessionManager.LoadChatHistory();
            chatHistory.Add(new ChatEntry { role = "user", content = detailedContext });
            
            await SendApiRequest(chatHistory, authToken);
        }

        private Task<bool> Initialize()
        {
            if (_isInitialized) return Task.FromResult(true);

            LoadSettings();
            LoadSystemPrompt();

            if (!_settings || string.IsNullOrEmpty(_systemPrompt))
            {
                Debug.LogError("[GeminiApiClient] Initialization failed. Check if MCPSettings and MCP_SystemPrompt.txt exist.");
                return Task.FromResult(false);
            }

            _isInitialized = true;
            return Task.FromResult(true);
        }

        private static string BuildInitialPrompt(string prompt, List<Object> stagedContext)
        {
            var promptBuilder = new StringBuilder();
            promptBuilder.AppendLine("User Request:");
            promptBuilder.AppendLine(prompt);

            if (stagedContext == null || stagedContext.All(o => o == null)) return promptBuilder.ToString();
            
            var tier1Summary = ContextBuilder.BuildTier1Summary(stagedContext);
            promptBuilder.AppendLine("\n--- Staged Context ---");
            promptBuilder.AppendLine(tier1Summary);
            promptBuilder.AppendLine("--- End Context ---");

            return promptBuilder.ToString();
        }

        private async Task SendApiRequest(List<ChatEntry> chatHistory, string authToken)
        {
            var url = $"https://{_settings.gcpRegion}-aiplatform.googleapis.com/v1/projects/{_settings.gcpProjectId}/locations/{_settings.gcpRegion}/publishers/google/models/{_settings.modelName}:generateContent";
            
            var systemPromptWithContext = GetWorkingContext();
            
            var apiRequest = new ApiRequest
            {
                system_instruction = new SystemInstruction
                {
                    parts = new List<Part> { new() { text = systemPromptWithContext } }
                },
                contents = chatHistory.Select(entry => new Content
                {
                    role = entry.role == "assistant" ? "model" : entry.role,
                    parts = new List<Part> { new() { text = entry.content } }
                }).ToList()
            };
            
            var jsonPayload = apiRequest.ToJson();

            using var request = new UnityWebRequest(url, "POST");
            var bodyRaw = Encoding.UTF8.GetBytes(jsonPayload);
            request.uploadHandler = new UploadHandlerRaw(bodyRaw);
            request.downloadHandler = new DownloadHandlerBuffer();
            request.SetRequestHeader("Content-Type", "application/json");
            request.SetRequestHeader("Authorization", $"Bearer {authToken}");

            var operation = request.SendWebRequest();
            while (!operation.isDone)
            {
                await Task.Yield();
            }

            if (request.result == UnityWebRequest.Result.Success)
            {
                var responseJson = request.downloadHandler.text;
                Debug.Log($"[GeminiApiClient] API Response: \n{responseJson}");
                
                var apiResponse = responseJson.FromJson<ApiResponse>();
                if (apiResponse.candidates != null && apiResponse.candidates.Any())
                {
                    var rawText = apiResponse.candidates[0].content.parts[0].text;
                    var commandJson = ExtractJsonFromString(rawText);
                    if (string.IsNullOrEmpty(commandJson))
                    {
                        Debug.LogError($"[GeminiApiClient] Could not extract valid JSON from the LLM response text: \n{rawText}");
                        return;
                    }
                    chatHistory.Add(new ChatEntry { role = "model", content = commandJson });
                    SessionManager.SaveChatHistory(chatHistory);
                    
                    CommandResponse commandResponse;

                    try
                    {
                        commandResponse = commandJson.FromJson<CommandResponse>();
                    }
                    catch (Exception exception)
                    {
                        Debug.LogException(exception);
                        return;
                    }
                    if (commandResponse is { commands: not null })
                    {
                        CommandExecutor.SetQueue(commandResponse.commands);
                    }
                    else
                    {
                        Debug.LogError($"[GeminiApiClient] Failed to parse command structure from LLM response text: {commandJson}");
                    }
                }
            }
            else
            {
                Debug.LogError($"[GeminiApiClient] API Error: {request.error}\n{request.downloadHandler.text}");
            }
        }

        private void LoadSettings()
        {
            if (_settings) return;
            var guids = AssetDatabase.FindAssets("t:MCPSettings");
            if (guids.Length == 0) return;
            var path = AssetDatabase.GUIDToAssetPath(guids[0]);
            _settings = AssetDatabase.LoadAssetAtPath<Settings.MCPSettings>(path);
        }

        private void LoadSystemPrompt()
        {
            if (!string.IsNullOrEmpty(_systemPrompt)) return;
            var guids = AssetDatabase.FindAssets("MCP_SystemPrompt");
            if (guids.Length == 0) return;
            var path = AssetDatabase.GUIDToAssetPath(guids[0]);
            _systemPrompt = File.ReadAllText(path);
        }
        
        private static string ExtractJsonFromString(string text)
        {
            if (string.IsNullOrWhiteSpace(text)) return null;

            var firstBrace = text.IndexOf('{');
            var lastBrace = text.LastIndexOf('}');

            if (firstBrace == -1 || lastBrace == -1 || lastBrace < firstBrace)
            {
                return null;
            }

            return text.Substring(firstBrace, lastBrace - firstBrace + 1);
        }

        private string GetAuthToken()
        {
            if (!_settings) return null;
            try
            {
                var startInfo = new ProcessStartInfo
                {
                    FileName = _settings.gcloudPath,
                    Arguments = "auth print-access-token",
                    RedirectStandardOutput = true,
                    RedirectStandardError = true,
                    UseShellExecute = false,
                    CreateNoWindow = true
                };
                using var process = Process.Start(startInfo);
                if (process == null) return null;
                var accessToken = process.StandardOutput.ReadToEnd().Trim();
                var error = process.StandardError.ReadToEnd();
                process.WaitForExit();
                if (process.ExitCode == 0) return accessToken;
                Debug.LogError($"[GeminiApiClient] gcloud auth error: {error}");
                return null;
            }
            catch (Exception e)
            {
                Debug.LogError($"[GeminiApiClient] Exception while getting auth token: {e.Message}");
                return null;
            }
        }

        private string GetWorkingContext()
        {
            var workingContext = SessionManager.LoadWorkingContext();
            var systemPromptWithContext = new StringBuilder(_systemPrompt);
            systemPromptWithContext.AppendLine("\n\n## CURRENT WORKING CONTEXT");
            systemPromptWithContext.AppendLine("This is your short-term memory. Use the data here before asking for it again.");
            systemPromptWithContext.AppendLine("```json");
            systemPromptWithContext.AppendLine(workingContext.ToString(Newtonsoft.Json.Formatting.Indented));
            systemPromptWithContext.AppendLine("```");
            return systemPromptWithContext.ToString();
        }
    }
}