From a0b4473e87beab8381d32c0b7d7f7231cf0d89e1 Mon Sep 17 00:00:00 2001
From: Marpioux <Marius.pingaud@gmail.com>
Date: Mon, 4 Nov 2024 11:18:54 +0100
Subject: [PATCH] add test for num_ctx / add prequery for request construction

---
 src/Pharo-Ollama-Tests/OllamaAPITest.class.st | 12 +++++++
 src/Pharo-Ollama/OllamaAPI.class.st           | 35 ++++++++++++-------
 2 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/src/Pharo-Ollama-Tests/OllamaAPITest.class.st b/src/Pharo-Ollama-Tests/OllamaAPITest.class.st
index 52706fa..44e8468 100644
--- a/src/Pharo-Ollama-Tests/OllamaAPITest.class.st
+++ b/src/Pharo-Ollama-Tests/OllamaAPITest.class.st
@@ -17,6 +17,18 @@ OllamaAPITest >> setUp [
 	ollama := OllamaAPI new.
 ]
 
+{ #category : 'tests' }
+OllamaAPITest >> testNumTokenVariables [
+	
+	| ollamatest |
+	ollamatest := OllamaAPI new.
+	self assert: (ollamatest num_ctx isNil).
+	self assert: (((ollamatest prequery: 'test') at: #options) isEmptyOrNil).
+	ollamatest num_ctx: 50.
+	self assert: (ollamatest num_ctx == 50).
+	self assert: (((ollamatest prequery: 'test') at: #options) isNotEmpty).
+]
+
 { #category : 'running' }
 OllamaAPITest >> testOllamaAPIHasDefaultModelInitialized [
 
diff --git a/src/Pharo-Ollama/OllamaAPI.class.st b/src/Pharo-Ollama/OllamaAPI.class.st
index efc501f..b239930 100644
--- a/src/Pharo-Ollama/OllamaAPI.class.st
+++ b/src/Pharo-Ollama/OllamaAPI.class.st
@@ -191,6 +191,27 @@ OllamaAPI >> port: anObject [
 	port := anObject
 ]
 
+{ #category : 'api' }
+OllamaAPI >> prequery: aTextQuery [
+
+    | paramDic optionsDic |
+    paramDic := Dictionary new.
+    paramDic at: #model put: self model fullName.
+    paramDic at: #prompt put: (self buildQuery: aTextQuery).
+    paramDic at: #stream put: self stream.
+    self format ifNotNil: [ :f | paramDic at: #format put: f ].
+    self addParameterTo: paramDic.
+
+    self num_ctx ifNotNil: [
+        optionsDic := Dictionary new.
+        optionsDic at: #num_ctx put: self num_ctx.
+        paramDic at: #options put: optionsDic.
+    ].
+
+    ^ paramDic
+
+]
+
 { #category : 'api' }
 OllamaAPI >> pull [
 	"Pull the model linked to this API"
@@ -211,7 +232,7 @@ OllamaAPI >> pull [
 { #category : 'api' }
 OllamaAPI >> query: aTextQuery [
 
-	| znClient paramDic writer optionsDic |
+	| znClient  writer  |
 	znClient := ZnClient new.
 	writer := ZnUtils defaultJSONWriter.
 	znClient
@@ -221,18 +242,8 @@ OllamaAPI >> query: aTextQuery [
 	znClient port: self port.
 	znClient path: 'api/generate'.
 
-	paramDic := Dictionary new.
-	paramDic at: #model put: self model fullName.
-	paramDic at: #prompt put: (self buildQuery: aTextQuery).
-	paramDic at: #stream put: self stream.
-	self format ifNotNil: [ :f | paramDic at: #format put: f ].
-	self addParameterTo: paramDic.
-
-	optionsDic := Dictionary new.
-	optionsDic at: #num_ctx put: self num_ctx.
-	paramDic at: #options put: optionsDic.
 
-	znClient contents: paramDic.
+	znClient contents: (self prequery: aTextQuery).
 
 	self stream 
     	ifTrue: [