sexy update

nodes/any.py

@@ -1,30 +1,28 @@
 """
-AnyNode v0.1
+@author: newsbubbles
+@title: AnyNode v0.1
+@nickname: AnyNode
+@description: This single node uses an LLM to generate a functionality based on your request. You can make the node do anything.
 author: newsbubbles
-inspired by: rgthree AnySwitch
-shouts and thanks: MachineLearners Discord
 """
+# inspired by: rgthree AnySwitch, shouts and thanks to MachineLearners Discord
 
 # TODO: finish inputs section which shows python types of inputs for the final prompt
 # - Store the json in some sort of hidden file which stores in the saved workflow
 # - How do I use just the "internal variables" since it seems scope on AnyNode class params are global (not good for storing script)
 
-import os, json, random, string, sys, math, datetime, collections, itertools, functools, urllib, shutil, re, numpy, torch
+import os, json, random, string, sys, math, datetime, collections, itertools, functools, urllib, shutil, re, torch
+import numpy as np
+import collections.abc
 import traceback
 import os
 import openai
 from openai import OpenAI
-from .context_utils import is_context_empty, _create_context_data
-from .constants import get_category, get_name
-from .utils import any_type
-
-def is_none(value):
-    """Checks if a value is none. Pulled out in case we want to expand what 'None' means."""
-    if value is not None:
-        if isinstance(value, dict) and 'model' in value and 'clip' in value:
-            return is_context_empty(value)
-    return value is None
+# from .context_utils import is_context_empty, _create_context_data
+# from .constants import get_category, get_name
+from .utils import any_type, is_none, variable_info
 
+# The template for the system message sent to ChatCompletions
 SYSTEM_TEMPLATE = """
 # Coding a Python Function
 You are an expert python coder who specializes in writing custom nodes for ComfyUI.
@@ -33,6 +31,10 @@ You are an expert python coder who specializes in writing custom nodes for Comfy
 It is not required to use any of these libraries, but if you do use any import in your code it must be on this list:
 [[IMPORTS]]
 
+## Input Data
+Here is some important information about the input data:
+- input_data: [[INPUT]]
+
 ## Coding Instructions
 - Your job is to code the user's requested node given the input and desired output type.
 - Code only the contents of the function itself.
@@ -47,13 +49,14 @@ class AnyNode:
 
   NAME = "AnyNode"
   CATEGORY = "utils"
-  
-  script = None
-  last_prompt = None
-  imports = []
-  
+
   ALLOWED_IMPORTS = {"os", "re", "json", "random", "string", "sys", "math", "datetime", "collections", "itertools", "functools", "urllib", "shutil", "numpy", "openai", "traceback", "torch"}
 
+  def __init__(self):
+      self.script = None
+      self.last_prompt = None
+      self.imports = []  
+  
   @classmethod
   def INPUT_TYPES(cls):  # pylint: disable = invalid-name, missing-function-docstring
     return {
@@ -75,16 +78,17 @@ class AnyNode:
   RETURN_NAMES = ('any',)
   FUNCTION = "go"
 
-  def get_openai_response(self, prompt):
+  def get_openai_response(self, prompt:str, any=None) -> str:
       try:
           client = OpenAI(
               # This is the default and can be omitted
               api_key=os.environ.get("OPENAI_API_KEY"),
           )
+          final_template = SYSTEM_TEMPLATE.replace('[[IMPORTS]]', ",".join(list(self.ALLOWED_IMPORTS))).replace('[[INPUT]]', variable_info(any))
           response = client.chat.completions.create(
               model="gpt-4o",  # Use the model of your choice, e.g., gpt-4 or gpt-3.5-turbo
               messages=[
-                  {"role": "system", "content": SYSTEM_TEMPLATE},
+                  {"role": "system", "content": final_template},
                   {"role": "user", "content": prompt}
               ]
           )
@@ -125,8 +129,7 @@ class AnyNode:
       if not is_none(any):
           if self.script is None or self.last_prompt != prompt:
               # Generate the function code using OpenAI
-              r = self.get_openai_response(prompt.replace('[[IMPORTS]]', ",".join(list(self.ALLOWED_IMPORTS))))
-              print(f"Generated code:\n{r}")
+              r = self.get_openai_response(prompt, any=any)
               
               # Store the script for future use
               self.script = self.extract_imports(r)

nodes/utils.py

@@ -1,7 +1,9 @@
 import json
 import os
 import re
-
+import numpy as np
+import torch
+from .context_utils import is_context_empty
 
 class AnyType(str):
   """A special class that is always equal in not equal comparisons. Credit to pythongosssss"""
@@ -82,4 +84,61 @@ def path_exists(path):
   """Checks if a path exists, accepting None type."""
   if path is not None:
     return os.path.exists(path)
-  return False
+  return False
+
+def is_none(value):
+    """Checks if a value is none. Pulled out in case we want to expand what 'None' means."""
+    if value is not None:
+        if isinstance(value, dict) and 'model' in value and 'clip' in value:
+            return is_context_empty(value)
+    return value is None
+
+def get_variable_info(variable):
+    info = {}
+    
+    # Get the exact type of the variable
+    info['type'] = type(variable).__name__
+
+    # Check for common types and add relevant information
+    if isinstance(variable, (np.ndarray, torch.Tensor)):
+        info['shape'] = tuple(variable.shape)
+        info['dtype'] = str(variable.dtype)
+    elif isinstance(variable, (list, tuple, set)):
+        info['length'] = len(variable)
+        if len(variable) > 128:
+            info['structure'] = f"{info['type']} of length {info['length']}"
+        else:
+            info['structure'] = [type(v).__name__ for v in variable]
+    elif isinstance(variable, dict):
+        info['length'] = len(variable)
+        if len(variable) > 128:
+            info['structure'] = f"{info['type']} of length {info['length']}"
+        else:
+            info['keys'] = list(variable.keys())
+            info['values'] = [type(variable[k]).__name__ for k in variable]
+    elif isinstance(variable, str):
+        info['length'] = len(variable)
+    elif hasattr(variable, '__dict__'):
+        # For user-defined classes, show their attributes
+        info['attributes'] = vars(variable)
+
+    return info
+
+def variable_info(variable):
+    info = get_variable_info(variable)
+    info_str = f"Type: {info['type']}"
+    
+    if 'shape' in info:
+        info_str += f", Shape: {info['shape']}, Dtype: {info['dtype']}"
+    if 'length' in info:
+        info_str += f", Length: {info['length']}"
+    if 'structure' in info:
+        info_str += f", Structure: {info['structure']}"
+    if 'keys' in info:
+        info_str += f", Keys: {info['keys']}"
+    if 'values' in info:
+        info_str += f", Values: {info['values']}"
+    if 'attributes' in info:
+        info_str += f", Attributes: {info['attributes']}"
+
+    return info_str