fix: pytest issues

application/llm/premai.py

@@ -1,32 +1,35 @@
 from application.llm.base import BaseLLM
 from application.core.settings import settings
 
+
 class PremAILLM(BaseLLM):
 
     def __init__(self, api_key):
         from premai import Prem
-        
-        self.client = Prem(
-            api_key=api_key
-        )
+
+        self.client = Prem(api_key=api_key)
         self.api_key = api_key
         self.project_id = settings.PREMAI_PROJECT_ID
 
-    def gen(self, model, messages, stream=False, **kwargs):
-        response = self.client.chat.completions.create(model=model,
+    def _raw_gen(self, model, messages, stream=False, **kwargs):
+        response = self.client.chat.completions.create(
+            model=model,
             project_id=self.project_id,
             messages=messages,
             stream=stream,
-            **kwargs)
+            **kwargs
+        )
 
         return response.choices[0].message["content"]
 
-    def gen_stream(self, model, messages, stream=True, **kwargs):
-        response = self.client.chat.completions.create(model=model,
+    def _raw_gen_stream(self, model, messages, stream=True, **kwargs):
+        response = self.client.chat.completions.create(
+            model=model,
             project_id=self.project_id,
             messages=messages,
             stream=stream,
-            **kwargs)
+            **kwargs
+        )
 
         for line in response:
             if line.choices[0].delta["content"] is not None: