FunnyGirl/Source/FunnyGirl/Private/LLM/BrainComponentBase.cpp

// Copyright Epic Games, Inc. All Rights Reserved.

#include "LLM/BrainComponentBase.h"

#include "Dom/JsonObject.h"
#include "Dom/JsonValue.h"
#include "Interfaces/IHttpRequest.h"
#include "Kismet/GameplayStatics.h"
#include "LLM/BrainToolBase.h"
#include "LLM/LLMProviderDataAsset.h"
#include "Misc/FileHelper.h"
#include "Misc/Paths.h"
#include "Serialization/JsonSerializer.h"
#include "Serialization/JsonWriter.h"
#include "HAL/FileManager.h"
#include "World/WorldContextActor.h"

UBrainComponentBase::UBrainComponentBase()
	: bAutoFindWorldContextActor(true)
	, bIsThinking(false)
{
	PrimaryComponentTick.bCanEverTick = false;
}

void UBrainComponentBase::BeginPlay()
{
	Super::BeginPlay();

	AutoFindWorldContextActor();
}

void UBrainComponentBase::CancelRequest()
{
	if (ActiveRequest.IsValid())
	{
		ActiveRequest->CancelRequest();
		ActiveRequest.Reset();
	}

	SetIsThinking(false);
}

bool UBrainComponentBase::IsThinking() const
{
	return bIsThinking;
}

FString UBrainComponentBase::GetChatIdentityId() const
{
	BroadcastBrainError(FString::Printf(TEXT("%s does not provide a chat identity id."), *GetBrainDebugName()));
	return FString();
}

void UBrainComponentBase::BroadcastBrainError(const FString& ErrorMessage) const
{
	UE_LOG(LogTemp, Warning, TEXT("%s error: %s"), *GetBrainDebugName(), *ErrorMessage);
}

FString UBrainComponentBase::GetBrainDebugName() const
{
	return GetNameSafe(GetClass());
}

void UBrainComponentBase::SetIsThinking(bool bNewIsThinking)
{
	bIsThinking = bNewIsThinking;
}

void UBrainComponentBase::AutoFindWorldContextActor()
{
	if (WorldContextActor || !bAutoFindWorldContextActor)
	{
		return;
	}

	WorldContextActor = Cast<AWorldContextActor>(UGameplayStatics::GetActorOfClass(this, AWorldContextActor::StaticClass()));
}

bool UBrainComponentBase::ValidateProvider(const ULLMProviderDataAsset* EffectiveProvider) const
{
	if (!EffectiveProvider)
	{
		BroadcastBrainError(FString::Printf(TEXT("%s LLM provider is not set."), *GetBrainDebugName()));
		return false;
	}

	if (EffectiveProvider->ApiBaseUrl.IsEmpty() || EffectiveProvider->ChatCompletionsPath.IsEmpty())
	{
		BroadcastBrainError(FString::Printf(TEXT("%s LLM provider endpoint is not configured."), *GetBrainDebugName()));
		return false;
	}

	if (EffectiveProvider->Model.IsEmpty())
	{
		BroadcastBrainError(FString::Printf(TEXT("%s LLM provider model is not configured."), *GetBrainDebugName()));
		return false;
	}

	return true;
}

bool UBrainComponentBase::ConfigureRequestWithBody(
	const TSharedRef<IHttpRequest, ESPMode::ThreadSafe>& Request,
	const ULLMProviderDataAsset* EffectiveProvider,
	const FString& RequestBody,
	const FString& AcceptHeader) const
{
	if (!ValidateProvider(EffectiveProvider))
	{
		return false;
	}

	Request->SetURL(EffectiveProvider->GetChatCompletionsUrl());
	Request->SetVerb(TEXT("POST"));
	Request->SetHeader(TEXT("Content-Type"), TEXT("application/json"));
	Request->SetHeader(TEXT("Accept"), AcceptHeader);

	if (!EffectiveProvider->ApiKey.IsEmpty())
	{
		Request->SetHeader(TEXT("Authorization"), FString::Printf(TEXT("Bearer %s"), *EffectiveProvider->ApiKey));
	}

	for (const TPair<FString, FString>& Header : EffectiveProvider->ExtraHeaders)
	{
		Request->SetHeader(Header.Key, Header.Value);
	}

	Request->SetTimeout(EffectiveProvider->TimeoutSeconds);
	Request->SetContentAsString(RequestBody);
	return true;
}

void UBrainComponentBase::AddSystemMessage(TArray<TSharedPtr<FJsonValue>>& Messages, const FString& Content) const
{
	const FString TrimmedContent = Content.TrimStartAndEnd();
	if (TrimmedContent.IsEmpty())
	{
		return;
	}

	Messages.Add(MakeShared<FJsonValueObject>(MakeMessageObject(TEXT("system"), TrimmedContent)));
}

TSharedRef<FJsonObject> UBrainComponentBase::MakeMessageObject(const FString& Role, const FString& Content) const
{
	TSharedRef<FJsonObject> MessageObject = MakeShared<FJsonObject>();
	MessageObject->SetStringField(TEXT("role"), Role);
	MessageObject->SetStringField(TEXT("content"), Content);
	return MessageObject;
}

FString UBrainComponentBase::SerializeJsonObject(const TSharedRef<FJsonObject>& RootObject) const
{
	FString OutputJson;
	TSharedRef<TJsonWriter<>> Writer = TJsonWriterFactory<>::Create(&OutputJson);
	FJsonSerializer::Serialize(RootObject, Writer);
	return OutputJson;
}

void UBrainComponentBase::AppendAllowedTools(TSharedRef<FJsonObject> RootObject, const TArray<TSubclassOf<UBrainToolBase>>& AllowedTools) const
{
	TArray<TSharedPtr<FJsonValue>> Tools;
	const TArray<TSharedPtr<FJsonValue>>* ExistingTools = nullptr;
	if (RootObject->TryGetArrayField(TEXT("tools"), ExistingTools) && ExistingTools)
	{
		Tools = *ExistingTools;
	}

	for (const TSubclassOf<UBrainToolBase>& ToolClass : AllowedTools)
	{
		if (!ToolClass)
		{
			continue;
		}

		const UBrainToolBase* ToolDefaultObject = ToolClass->GetDefaultObject<UBrainToolBase>();
		TSharedPtr<FJsonObject> ToolDefinition = BuildToolDefinition(ToolDefaultObject);
		if (ToolDefinition.IsValid())
		{
			Tools.Add(MakeShared<FJsonValueObject>(ToolDefinition.ToSharedRef()));
		}
	}

	if (!Tools.IsEmpty())
	{
		RootObject->SetArrayField(TEXT("tools"), Tools);
		RootObject->SetStringField(TEXT("tool_choice"), TEXT("auto"));
	}
}

TSharedPtr<FJsonObject> UBrainComponentBase::BuildToolDefinition(const UBrainToolBase* ToolDefaultObject) const
{
	if (!ToolDefaultObject)
	{
		return nullptr;
	}

	const FString ToolName = ToolDefaultObject->GetResolvedToolName();
	if (ToolName.TrimStartAndEnd().IsEmpty())
	{
		BroadcastBrainError(FString::Printf(TEXT("Tool has empty ToolName: %s"), *GetNameSafe(ToolDefaultObject->GetClass())));
		return nullptr;
	}

	TSharedRef<FJsonObject> ParametersObject = MakeShared<FJsonObject>();
	if (ToolDefaultObject->ParametersJsonSchema.IsEmpty())
	{
		BroadcastBrainError(FString::Printf(TEXT("Tool %s has empty ParametersJsonSchema."), *ToolName));
		return nullptr;
	}

	{
		TSharedPtr<FJsonObject> ParsedParametersObject;
		TSharedRef<TJsonReader<>> ParametersReader = TJsonReaderFactory<>::Create(ToolDefaultObject->ParametersJsonSchema.ToString());
		if (FJsonSerializer::Deserialize(ParametersReader, ParsedParametersObject) && ParsedParametersObject.IsValid())
		{
			ParametersObject = ParsedParametersObject.ToSharedRef();
		}
		else
		{
			BroadcastBrainError(FString::Printf(TEXT("Tool %s has invalid ParametersJsonSchema."), *ToolName));
			return nullptr;
		}
	}

	if (ParametersObject->Values.IsEmpty())
	{
		BroadcastBrainError(FString::Printf(TEXT("Tool %s has empty ParametersJsonSchema object."), *ToolName));
		return nullptr;
	}

	TSharedRef<FJsonObject> FunctionObject = MakeShared<FJsonObject>();
	FunctionObject->SetStringField(TEXT("name"), ToolName);
	FunctionObject->SetStringField(TEXT("description"), ToolDefaultObject->Description.ToString());
	FunctionObject->SetObjectField(TEXT("parameters"), ParametersObject);

	TSharedRef<FJsonObject> ToolObject = MakeShared<FJsonObject>();
	ToolObject->SetStringField(TEXT("type"), TEXT("function"));
	ToolObject->SetObjectField(TEXT("function"), FunctionObject);
	return ToolObject;
}

const UBrainToolBase* UBrainComponentBase::FindAllowedToolDefaultObject(const FString& ToolName, const TArray<TSubclassOf<UBrainToolBase>>& AllowedTools) const
{
	if (ToolName.TrimStartAndEnd().IsEmpty())
	{
		return nullptr;
	}

	for (const TSubclassOf<UBrainToolBase>& ToolClass : AllowedTools)
	{
		if (!ToolClass)
		{
			continue;
		}

		const UBrainToolBase* ToolDefaultObject = ToolClass->GetDefaultObject<UBrainToolBase>();
		if (ToolDefaultObject && ToolDefaultObject->GetResolvedToolName().Equals(ToolName, ESearchCase::IgnoreCase))
		{
			return ToolDefaultObject;
		}
	}

	return nullptr;
}

bool UBrainComponentBase::ExecuteBrainToolCalls(const TArray<TSharedPtr<FJsonValue>>& ToolCalls, const TArray<TSubclassOf<UBrainToolBase>>& AllowedTools)
{
	if (ToolCalls.IsEmpty())
	{
		return true;
	}

	bool bAllSucceeded = true;
	for (const TSharedPtr<FJsonValue>& ToolCallValue : ToolCalls)
	{
		const TSharedPtr<FJsonObject>* ToolCallObject = nullptr;
		if (!ToolCallValue.IsValid() || !ToolCallValue->TryGetObject(ToolCallObject) || !ToolCallObject || !ToolCallObject->IsValid())
		{
			bAllSucceeded = false;
			continue;
		}

		const TSharedPtr<FJsonObject>* FunctionObject = nullptr;
		if (!(*ToolCallObject)->TryGetObjectField(TEXT("function"), FunctionObject) || !FunctionObject || !FunctionObject->IsValid())
		{
			BroadcastBrainError(TEXT("Native tool call is missing function object."));
			bAllSucceeded = false;
			continue;
		}

		FString ToolName;
		if (!(*FunctionObject)->TryGetStringField(TEXT("name"), ToolName))
		{
			BroadcastBrainError(TEXT("Native tool call is missing function name."));
			bAllSucceeded = false;
			continue;
		}

		FString ArgumentsJson;
		if (!(*FunctionObject)->TryGetStringField(TEXT("arguments"), ArgumentsJson))
		{
			ArgumentsJson = TEXT("{}");
		}

		const UBrainToolBase* ToolDefaultObject = FindAllowedToolDefaultObject(ToolName, AllowedTools);
		if (!ToolDefaultObject)
		{
			BroadcastBrainError(FString::Printf(TEXT("Tried to call disallowed native tool: %s"), *ToolName));
			bAllSucceeded = false;
			continue;
		}

		UBrainToolBase* ToolInstance = NewObject<UBrainToolBase>(this, ToolDefaultObject->GetClass());
		if (!ToolInstance)
		{
			BroadcastBrainError(FString::Printf(TEXT("Could not create native tool: %s"), *ToolName));
			bAllSucceeded = false;
			continue;
		}

		FString ErrorMessage;
		if (!ToolInstance->ExecuteTool(this, ArgumentsJson, ErrorMessage))
		{
			BroadcastBrainError(ErrorMessage.IsEmpty() ? FString::Printf(TEXT("Native tool failed: %s"), *ToolName) : ErrorMessage);
			bAllSucceeded = false;
		}

		ToolInstance->CleanupTool();
	}

	return bAllSucceeded;
}

FString UBrainComponentBase::GetChatDirectory() const
{
	const FString Directory = FPaths::ProjectSavedDir() / TEXT("Chat") / GetChatIdentityId();
	IFileManager::Get().MakeDirectory(*Directory, true);
	return Directory;
}

FString UBrainComponentBase::GetChatFilePath() const
{
	return GetChatSiblingFilePath(TEXT("Chat.json"));
}

FString UBrainComponentBase::GetChatSiblingFilePath(const FString& FileName) const
{
	return GetChatDirectory() / FileName;
}

void UBrainComponentBase::SaveRequestBodyToFile(const FString& RequestBody, const FString& FileName) const
{
	FFileHelper::SaveStringToFile(RequestBody, *GetChatSiblingFilePath(FileName), FFileHelper::EEncodingOptions::ForceUTF8WithoutBOM);
}