diff --git a/llm_vlm_api.py b/llm_vlm_api.py index 7d369f8..537cf1c 100644 --- a/llm_vlm_api.py +++ b/llm_vlm_api.py @@ -12,10 +12,10 @@ class LLM_VLM_API: "base_url": ("STRING", {"default": ""}), "api_key": ("STRING", {"default": ""}), "model": ("STRING", {"default": ""}), + "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.01}), + "max_tokens": ("INT", {"default": 4096, "min": 1, "max": 8192, "step": 1}), }, "optional": { - "temperature": ("FLOAT", {"forceInput": True}), - "max_tokens": ("INT", {"forceInput": True}), "image": ("IMAGE",), "system prompt": ("STRING", {"forceInput": True}), "text": ("STRING", {"forceInput": True}), @@ -27,12 +27,10 @@ class LLM_VLM_API: FUNCTION = "execute" CATEGORY = "LLM(VLM) API" - def execute(self, base_url, api_key, model, **kwargs): + def execute(self, base_url, api_key, model, temperature, max_tokens, **kwargs): system_prompt = kwargs.get("system prompt", None) image = kwargs.get("image", None) text = kwargs.get("text", None) - temperature = kwargs.get("temperature", None) - max_tokens = kwargs.get("max_tokens", None) if not text and image is None: return ("Error: No text or image provided.",) @@ -54,23 +52,19 @@ class LLM_VLM_API: client = OpenAI(api_key=api_key, base_url=base_url) - api_params = { - "model": model, - "messages": [] - } - - if temperature is not None: - api_params["temperature"] = temperature - if max_tokens is not None: - api_params["max_tokens"] = max_tokens - try: + messages = [] if system_prompt: - api_params["messages"].append({"role": "system", "content": system_prompt}) + messages.append({"role": "system", "content": system_prompt}) - api_params["messages"].append({"role": "user", "content": content}) + messages.append({"role": "user", "content": content}) - resp = client.chat.completions.create(**api_params) + resp = client.chat.completions.create( + model=model, + temperature=temperature, + max_tokens=max_tokens, + messages=messages + ) result = resp.choices[0].message.content except Exception as e: result = str(e)