test: closed llm output

agiresearch · dongyuanjushi · Apr 8, 2024 · Apr 2, 2024 · Apr 2, 2024 · Apr 2, 2024
commit 4343d21d3b2946240c2ef0a72e616ee4274eccd9
diff --git a/src/__init__.py b/src/__init__.py
diff --git a/src/llms/llms.py b/src/llms/llms.py
@@ -24,7 +24,7 @@
 import time
 
 class LLMKernel:
-    def __init__(self, llm_name: str, max_gpu_memory: dict, eval_device: str, max_new_tokens: int = 256):
+    def __init__(self, llm_name: str, max_gpu_memory: dict = None, eval_device: str = None, max_new_tokens: int = 256):
         print("Initialize AIOS powered by LLM: {}".format(llm_name))
         self.config = self.load_config(llm_name)
         self.max_gpu_memory = max_gpu_memory

diff --git a/tests/test_llms.py b/tests/test_llms.py
@@ -0,0 +1,13 @@
+import sys
+
+import os
+
+from src.llms.llms import LLMKernel
+
+def test_closed_llm():
+    llm_type = "gemini-pro"
+    llm = LLMKernel(llm_type)
+    prompt = "Craft a tale about a valiant warrior on a quest to uncover priceless treasures hidden within a mystical island."
+    response = llm.address_request(prompt)
+    assert isinstance(response, str)
+    print(response)