From 6c0af1d32fe7e1fb6253411e6c6f6485e7535db5 Mon Sep 17 00:00:00 2001
From: Rodrigo Ribeiro Gomes <rodrigor.gomes@hotmail.com>
Date: Tue, 1 Oct 2024 12:23:10 -0300
Subject: [PATCH] Enhacements of JSON support. Related to issue #20

---
 powershai/powershai.psm1       |  6 +++---
 powershai/providers/google.ps1 |  8 ++++++++
 util/TestPowershai.ps1         | 11 +++++++++++
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/powershai/powershai.psm1 b/powershai/powershai.psm1
index a20c3a9..08baef7 100644
--- a/powershai/powershai.psm1
+++ b/powershai/powershai.psm1
@@ -859,7 +859,7 @@ function Get-AiChat {
 		,#Formato da resposta 
 		 #Os formatos aceitáveis, e comportamento, devem seguir o mesmo da OpenAI: https://platform.openai.com/docs/api-reference/chat/create#chat-create-response_format
 		 #Atalhos:
-		 #	"json", equivale a {"type": "json_object"}
+		 #	"json"|"json_object", equivale a {"type": "json_object"}
 			$ResponseFormat = $null
 		
 		,#Lista de tools que devem ser invocadas!
@@ -888,8 +888,8 @@ function Get-AiChat {
 	$Provider = Get-AiCurrentProvider
 	$FuncParams = $PsBoundParameters;
 	
-	if($ResponseFormat -eq "json"){
-		$ResponseFormat = @{type = "json_object"}
+	if($ResponseFormat -in "json","json_object"){
+		$FuncParams.ResponseFormat = @{type = "json_object"}
 	}
 
 	if(!$model){
diff --git a/powershai/providers/google.ps1 b/powershai/providers/google.ps1
index f0008fd..2ec828a 100644
--- a/powershai/providers/google.ps1
+++ b/powershai/providers/google.ps1
@@ -162,6 +162,12 @@ function Invoke-GoogleGenerateContent {
 	#Note que reaproveitamos o convert to openai message do provider!
 	$GoogleContent = @(ConvertTo-GoogleContentMessage $messages)
 	
+	$Config = @{}
+	
+	if($ResponseFormat.type -eq "json_object"){
+		$Config.responseMimeType = "application/json"
+	}
+	
 	$Data = @{
 		contents = $GoogleContent.content
 		#tools = @()
@@ -170,6 +176,8 @@ function Invoke-GoogleGenerateContent {
 						@{text = $GoogleContent.SystemMessage }
 					)
 			}
+			
+		generationConfig = $Config
 	}
 	
 	$ReqParams = @{
diff --git a/util/TestPowershai.ps1 b/util/TestPowershai.ps1
index 53fb34e..234f80b 100644
--- a/util/TestPowershai.ps1
+++ b/util/TestPowershai.ps1
@@ -85,6 +85,17 @@ try {
 		}
 		
 		write-host "	TEMPORARY TEST PASSED";
+		
+		
+		# Json test!
+		$ExpectedA = Get-Random -Min 1 -Max 100;
+		$ExpectedB = [Guid]::NewGuid().Guid;
+		$prompt = @(
+			"Return an array with 5 objects"
+			"Each object must contains two props: a and b. a must contain value $ExpectedA, b must contains value $ExpectedB"
+		)
+		$JsonResult = @(iat -Json -Lines $prompt);
+		write-host "JsonOutput:`n$JsonResult";
 	}
 	
 } catch {