⚡ [Enhance] HuggingchatStreamer: logger auth messages

networks/huggingchat_streamer.py

@@ -1,43 +1,27 @@
 import copy
 import json
 import re
+
 import requests
-import uuid
+from curl_cffi import requests as cffi_requests
 
-# from curl_cffi import requests
 from tclogger import logger
-from transformers import AutoTokenizer
-
-from constants.models import (
-    MODEL_MAP,
-    STOP_SEQUENCES_MAP,
-    TOKEN_LIMIT_MAP,
-    TOKEN_RESERVED,
-)
+
+from constants.models import MODEL_MAP
 from constants.envs import PROXIES
-from constants.headers import (
-    REQUESTS_HEADERS,
-    HUGGINGCHAT_POST_HEADERS,
-    HUGGINGCHAT_SETTINGS_POST_DATA,
-)
+from constants.headers import HUGGINGCHAT_POST_HEADERS, HUGGINGCHAT_SETTINGS_POST_DATA
 from messagers.message_outputer import OpenaiStreamOutputer
+from messagers.message_composer import MessageComposer
+from messagers.token_checker import TokenChecker
 
 
-class HuggingchatStreamer:
+class HuggingchatRequester:
     def __init__(self, model: str):
         if model in MODEL_MAP.keys():
             self.model = model
         else:
-            self.model = "mixtral-8x7b"
+            self.model = "nous-mixtral-8x7b"
         self.model_fullname = MODEL_MAP[self.model]
-        self.message_outputer = OpenaiStreamOutputer(model=self.model)
-        # self.tokenizer = AutoTokenizer.from_pretrained(self.model_fullname)
-
-    # def count_tokens(self, text):
-    #     tokens = self.tokenizer.encode(text)
-    #     token_count = len(tokens)
-    #     logger.note(f"Prompt Token Count: {token_count}")
-    #     return token_count
 
     def get_hf_chat_id(self):
         request_url = "https://huggingface.co/chat/settings"
@@ -48,12 +32,13 @@ class HuggingchatStreamer:
         request_body.update(extra_body)
         logger.note(f"> hf-chat ID:", end=" ")
 
-        res = requests.post(
+        res = cffi_requests.post(
             request_url,
             headers=HUGGINGCHAT_POST_HEADERS,
             json=request_body,
             proxies=PROXIES,
             timeout=10,
+            impersonate="chrome",
         )
         self.hf_chat_id = res.cookies.get("hf-chat")
         if self.hf_chat_id:
@@ -61,9 +46,9 @@ class HuggingchatStreamer:
         else:
             logger.warn(f"[{res.status_code}]")
             logger.warn(res.text)
-            raise ValueError("Failed to get hf-chat ID!")
+            raise ValueError(f"Failed to get hf-chat ID: {res.text}")
 
-    def get_conversation_id(self, preprompt: str = ""):
+    def get_conversation_id(self, system_prompt: str = ""):
         request_url = "https://huggingface.co/chat/conversation"
         request_headers = HUGGINGCHAT_POST_HEADERS
         extra_headers = {
@@ -72,7 +57,7 @@ class HuggingchatStreamer:
         request_headers.update(extra_headers)
         request_body = {
             "model": self.model_fullname,
-            "preprompt": preprompt,
+            "preprompt": system_prompt,
         }
         logger.note(f"> Conversation ID:", end=" ")
 
@@ -120,7 +105,7 @@ class HuggingchatStreamer:
             logger.success(f"[{message_id}]")
         else:
             logger.warn(f"[{res.status_code}]")
-            raise ValueError("Failed to get conversation ID!")
+            raise ValueError("Failed to get message ID!")
 
         return message_id
 
@@ -139,9 +124,8 @@ class HuggingchatStreamer:
         else:
             logger_func = logger.warn
 
-        logger_func(status_code_str)
-
         logger.enter_quiet(not verbose)
+        logger_func(status_code_str)
 
         if status_code != 200:
             logger_func(res.text)
@@ -176,17 +160,17 @@ class HuggingchatStreamer:
 
         logger.exit_quiet(not verbose)
 
-    def chat_response(
-        self,
-        prompt: str = None,
-        temperature: float = 0.5,
-        top_p: float = 0.95,
-        max_new_tokens: int = None,
-        api_key: str = None,
-        use_cache: bool = False,
-    ):
+    def chat_completions(self, messages: list[dict], iter_lines=False, verbose=False):
+        composer = MessageComposer(model=self.model)
+        system_prompt, input_prompt = composer.decompose_to_system_and_input_prompt(
+            messages
+        )
+
+        checker = TokenChecker(input_str=system_prompt + input_prompt, model=self.model)
+        checker.check_token_limit()
+
         self.get_hf_chat_id()
-        self.get_conversation_id()
+        self.get_conversation_id(system_prompt=system_prompt)
         message_id = self.get_last_message_id()
 
         request_url = f"https://huggingface.co/chat/conversation/{self.conversation_id}"
@@ -200,7 +184,7 @@ class HuggingchatStreamer:
         request_body = {
             "files": [],
             "id": message_id,
-            "inputs": prompt,
+            "inputs": input_prompt,
             "is_continue": False,
             "is_retry": False,
             "web_search": False,
@@ -214,20 +198,106 @@ class HuggingchatStreamer:
             proxies=PROXIES,
             stream=True,
         )
-        self.log_response(res, stream=True, iter_lines=True, verbose=True)
+        self.log_response(res, stream=True, iter_lines=iter_lines, verbose=verbose)
         return res
 
-    def chat_return_dict(self, stream_response):
-        pass
 
-    def chat_return_generator(self, stream_response):
-        pass
+class HuggingchatStreamer:
+    def __init__(self, model: str):
+        if model in MODEL_MAP.keys():
+            self.model = model
+        else:
+            self.model = "nous-mixtral-8x7b"
+        self.model_fullname = MODEL_MAP[self.model]
+        self.message_outputer = OpenaiStreamOutputer(model=self.model)
+
+    def chat_response(self, messages: list[dict], verbose=False):
+        requester = HuggingchatRequester(model=self.model)
+        return requester.chat_completions(
+            messages=messages, iter_lines=False, verbose=verbose
+        )
+
+    def chat_return_generator(self, stream_response: requests.Response, verbose=False):
+        is_finished = False
+        for line in stream_response.iter_lines():
+            line = line.decode("utf-8")
+            line = re.sub(r"^data:\s*", "", line)
+            line = line.strip()
+            if not line:
+                continue
+
+            content = ""
+            content_type = "Completions"
+            try:
+                data = json.loads(line, strict=False)
+                msg_type = data.get("type")
+                if msg_type == "status":
+                    msg_status = data.get("status")
+                    continue
+                elif msg_type == "stream":
+                    content_type = "Completions"
+                    content = data.get("token", "")
+                    if verbose:
+                        logger.success(content, end="")
+                elif msg_type == "finalAnswer":
+                    content_type = "Finished"
+                    content = ""
+                    full_content = data.get("text")
+                    if verbose:
+                        logger.success("\n[Finished]")
+                    is_finished = True
+                    break
+                else:
+                    continue
+            except Exception as e:
+                logger.warn(e)
+
+            output = self.message_outputer.output(
+                content=content, content_type=content_type
+            )
+            yield output
+
+        if not is_finished:
+            yield self.message_outputer.output(content="", content_type="Finished")
+
+    def chat_return_dict(self, stream_response: requests.Response):
+        final_output = self.message_outputer.default_data.copy()
+        final_output["choices"] = [
+            {
+                "index": 0,
+                "finish_reason": "stop",
+                "message": {"role": "assistant", "content": ""},
+            }
+        ]
+        final_content = ""
+        for item in self.chat_return_generator(stream_response):
+            try:
+                data = json.loads(item)
+                delta = data["choices"][0]["delta"]
+                delta_content = delta.get("content", "")
+                if delta_content:
+                    final_content += delta_content
+            except Exception as e:
+                logger.warn(e)
+        final_output["choices"][0]["message"]["content"] = final_content.strip()
+        return final_output
 
 
 if __name__ == "__main__":
-    # model = "llama3-70b"
-    model = "command-r-plus"
+    # model = "command-r-plus"
+    model = "llama3-70b"
+    # model = "zephyr-141b"
+
     streamer = HuggingchatStreamer(model=model)
-    prompt = "what is your model?"
-    streamer.chat_response(prompt=prompt)
+    messages = [
+        {
+            "role": "system",
+            "content": "You are an LLM developed by CloseAI.\nYour name is Niansuh-Copilot.",
+        },
+        {"role": "user", "content": "Hello, what is your role?"},
+        {"role": "assistant", "content": "I am an LLM."},
+        {"role": "user", "content": "What is your name?"},
+    ]
+
+    streamer.chat_response(messages=messages)
     # HF_ENDPOINT=https://hf-mirror.com python -m networks.huggingchat_streamer