diff --git a/flowy/nodes_omost.py b/flowy/nodes_omost.py index e29f716..66df388 100644 --- a/flowy/nodes_omost.py +++ b/flowy/nodes_omost.py @@ -18,6 +18,8 @@ from .lib_omost.greedy_encode import ( CATEGORY = "Comflowy/Omost" from .utils import llm_request, logger from .types import LLM_MODELS, STRING +from .api_key_manager import load_api_key # Add this import + CANVAS_SIZE = 90 system_prompt = r"""You are a helpful AI assistant to compose images using the a json-based canvas system: @@ -102,7 +104,6 @@ class OmostLLMNode: "required": { "prompt": ("STRING", {"multiline": True}), "llm_model": (LLM_MODELS,), - "api_key": STRING, "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), }, } @@ -117,12 +118,16 @@ class OmostLLMNode: FUNCTION = "run_llm" CATEGORY = CATEGORY - def run_llm(self, prompt: str, llm_model: str, seed: int, api_key: str) -> Tuple[list[OmostCanvasCondition]]: + def run_llm(self, prompt: str, llm_model: str, seed: int) -> Tuple[list[OmostCanvasCondition]]: """Run LLM to generate area conditioning.""" if seed > 0xFFFFFFFF: seed = seed & 0xFFFFFFFF logger.warning("Seed is too large. Truncating to 32-bit: %d", seed) + api_key = load_api_key() + if not api_key: + raise ValueError("API Key is not set. Please use the 'Comflowy Set API Key' node to set a global API Key before using this node.") + try: generated_text = llm_request(prompt=prompt, llm_model=llm_model, system_prompt=system_prompt, api_key=api_key, max_tokens=4000, timeout=10) # 如果生成的字符中包含了多余的字符,比如 "```json" 或者 "```",则需要去掉改行 diff --git a/flowy/utils.py b/flowy/utils.py index ac912ae..9db80dc 100644 --- a/flowy/utils.py +++ b/flowy/utils.py @@ -47,3 +47,26 @@ def llm_request(prompt, system_prompt, llm_model, api_key, max_tokens=3000, time raise Exception(f"Error: {ret.get('error')}") except Exception as e: raise Exception(f"Failed to get response from LLM model with {API_HOST}/api/open/v0/prompt, error: {str(e)}") + +def get_nested_value(obj, path, default=None): + """ + Get a nested value from a dictionary using a dot-separated path. + + Args: + obj (dict): The dictionary to search in. + path (str): The dot-separated path to the desired value. + default: The value to return if the path is not found. + + Returns: + The value at the specified path, or the default value if not found. + """ + keys = path.split('.') + for key in keys: + if isinstance(obj, dict): + obj = obj.get(key, default) + else: + return default + return obj + +# Make sure to export the functions +__all__ = ['llm_request', 'get_nested_value']