Add files via upload

2025-08-05 13:52:03 +08:00
parent 43c42cd182
commit d16ec14445
4 changed files with 795 additions and 0 deletions
@@ -0,0 +1,119 @@
+import requests
+import json
+import re
+
+class OpenAINode:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "prompt": ("STRING", {
+                    "multiline": True,
+                    "default": "user prompts"
+                }),
+                "system_prompt": ("STRING", {
+                    "multiline": True,
+                    "default": "You are a prompt generation AI. your task is to take a user input for a stable difusion prompt and output and expand the supplied prompt in a stable difusion format to provide better output. Do not deviate from the format. Do not output anything other than a stable diffusion prompt."
+                }),
+                "model": ("STRING", {
+                    "multiline": True,
+                    "default": "DeepSeek-R1-Distill-Qwen-32B-910A"
+                }),
+                "api_url": ("STRING", {
+                    "multiline": False,
+                    "default": "http://127.0.0.1:1935/v1"
+                }),
+                "api_key": ("STRING", {
+                    "multiline": False,
+                    "default": "123"
+                }),
+                "temperature": ("FLOAT", {
+                    "default": 0.6,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01,
+                    "round": 0.01,
+                    "display": "number"
+                }),
+                "max_tokens": ("INT", {
+                        "default": 1024,
+                        "min": 1, 
+                        "max": 2048,
+                        "display": "number"
+                })
+            }
+        }
+    
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "get_completion"
+    CATEGORY = "OpenAIapi"
+    
+    def remove_think_tags(self, text):
+        """
+        移除文本中的 <think> 和 </think> 标签及其之间的内容
+        """
+        # 使用正则表达式匹配 <think> 到 </think> 之间的所有内容（包括标签本身）
+        # re.DOTALL 使 . 匹配包括换行符在内的任何字符
+        pattern = r'<think>.*?</think>'
+        cleaned_text = re.sub(pattern, '', text, flags=re.DOTALL | re.IGNORECASE)
+        
+        # 清理多余的空行和空白字符
+        cleaned_text = re.sub(r'\n\s*\n', '\n', cleaned_text.strip())
+        
+        return cleaned_text
+    
+    def get_completion(self, prompt, api_url, temperature, max_tokens, model, system_prompt, api_key):
+        try:
+            # 确保API URL格式正确
+            if not api_url.endswith('/chat/completions'):
+                if api_url.endswith('/v1'):
+                    api_url += '/chat/completions'
+                else:
+                    api_url += '/v1/chat/completions'
+            
+            headers = {
+                "Content-Type": "application/json"
+            }
+            
+            # 如果提供了有效的API密钥，添加认证头
+            if api_key and api_key != "123":
+                headers["Authorization"] = f"Bearer {api_key}"
+            
+            payload = {
+                "messages": [
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": prompt}
+                ],
+                "max_tokens": max_tokens,
+                "temperature": temperature,
+                "model": model,
+                "top_p": 0.95,
+                "top_k": 20,
+                "repetition_penalty": 1.1
+            }
+            
+            response = requests.post(api_url, headers=headers, data=json.dumps(payload), timeout=100000)
+            
+            if response.status_code == 200:
+                data = response.json()
+                # 兼容返回结构
+                if "choices" in data and len(data["choices"]) > 0:
+                    content = data["choices"][0]["message"]["content"]
+                    # 移除 think 标签内容
+                    cleaned_content = self.remove_think_tags(content)
+                    return (cleaned_content,)
+                else:
+                    return ("No choices found in response.",)
+            else:
+                return (f"Error: HTTP {response.status_code} - {response.text}",)
+                
+        except Exception as e:
+            return (f"Exception: {str(e)}",)
+
+NODE_CLASS_MAPPINGS = {
+    "OpenAINode": OpenAINode
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "OpenAINode": "OpenAI Node"
+}