From 6c0af3ac35b493b74b0bc1f866d45049d47a5197 Mon Sep 17 00:00:00 2001
From: ACBBZ <tpcacbbz@gmail.com>
Date: Fri, 24 May 2024 03:14:52 +0000
Subject: [PATCH] update chat

---
 src/blackbox/chat.py | 102 ++++++++++++++++++++++---------------------
 1 file changed, 53 insertions(+), 49 deletions(-)

diff --git a/src/blackbox/chat.py b/src/blackbox/chat.py
index 0baef46..7ddc52b 100644
--- a/src/blackbox/chat.py
+++ b/src/blackbox/chat.py
@@ -28,54 +28,10 @@ class Chat(Blackbox):
 
     # model_name有 Qwen1.5-14B-Chat , internlm2-chat-20b
     @logging_time(logger=logger)
-    def processing(self, model_name, prompt, template, context: list, temperature, top_p, n, max_tokens,stop,frequency_penalty,presence_penalty)  -> str:
-        if context == None:
-            context = []
-
-        # gpt-4， gpt-3.5-turbo
-        if re.search(r"gpt", model_name):
-            url = 'https://api.openai.com/v1/completions' 
-            key = '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'
-        else:
-            url = 'http://120.196.116.194:48892/v1/chat/completions'
-            key = 'YOUR_API_KEY'
-              
-        prompt_template = [
-           {"role": "system", "content": template},
-        ]
-
-        chat_inputs={
-            "model": model_name,
-            "messages": prompt_template + context + [
-                {
-                    "role": "user",
-                    "content": prompt
-                }
-            ],
-            "temperature": temperature,
-            "top_p": top_p,
-            "n": n,
-            "max_tokens": max_tokens,
-            "frequency_penalty": frequency_penalty,
-            "presence_penalty": presence_penalty,
-            "stop": stop
-        }
-
-        header = {
-            'Content-Type': 'application/json',
-            'Authorization': "Bearer " + key
-        }
-
-        fastchat_response = requests.post(url, json=chat_inputs, headers=header)
-
-        return fastchat_response.json()["choices"][0]["message"]["content"]
-    
-    async def fast_api_handler(self, request: Request) -> Response:
-        try:
-            data = await request.json()
-        except:
-            return JSONResponse(content={"error": "json parse error"}, status_code=status.HTTP_400_BAD_REQUEST)
+    def processing(self, *args, **kwargs)  -> str:
 
+        data = args[0]
+ 
         user_model_name = data.get("model_name")
         user_context = data.get("context")
         user_question = data.get("question")
@@ -88,6 +44,9 @@ class Chat(Blackbox):
         user_frequency_penalty = data.get("frequency_penalty")
         user_presence_penalty = data.get("presence_penalty")
 
+        if user_context == None:
+            user_context = []
+
         if user_question is None:
             return JSONResponse(content={"error": "question is required"}, status_code=status.HTTP_400_BAD_REQUEST)
 
@@ -119,5 +78,50 @@ class Chat(Blackbox):
             user_presence_penalty = 0.8
 
 
-        return JSONResponse(content={"response": self.processing(user_model_name, user_question, user_template, user_context, 
-        user_temperature, user_top_p, user_n, user_max_tokens,user_stop,user_frequency_penalty,user_presence_penalty)}, status_code=status.HTTP_200_OK)
\ No newline at end of file
+        # gpt-4， gpt-3.5-turbo
+        if re.search(r"gpt", user_model_name):
+            url = 'https://api.openai.com/v1/completions' 
+            key = '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'
+        else:
+            url = 'http://120.196.116.194:48892/v1/chat/completions'
+            key = 'YOUR_API_KEY'
+              
+        prompt_template = [
+           {"role": "system", "content": user_template},
+        ]
+
+        chat_inputs={
+            "model": user_model_name,
+            "messages": prompt_template + user_context + [
+                {
+                    "role": "user",
+                    "content": user_question
+                }
+            ],
+            "temperature": user_temperature,
+            "top_p": user_top_p,
+            "n": user_n,
+            "max_tokens": user_max_tokens,
+            "frequency_penalty": user_frequency_penalty,
+            "presence_penalty": user_presence_penalty,
+            "stop": user_stop
+        }
+
+        header = {
+            'Content-Type': 'application/json',
+            'Authorization': "Bearer " + key
+        }
+
+        fastchat_response = requests.post(url, json=chat_inputs, headers=header)
+
+        return fastchat_response.json()["choices"][0]["message"]["content"]
+    
+    async def fast_api_handler(self, request: Request) -> Response:
+        try:
+            data = await request.json()
+        except:
+            return JSONResponse(content={"error": "json parse error"}, status_code=status.HTTP_400_BAD_REQUEST)
+
+        setting: dict = data.get("settings")
+        
+        return JSONResponse(content={"response": self.processing(setting)}, status_code=status.HTTP_200_OK)
\ No newline at end of file