From 73948006b7da97fdd9726bfc3d7d5b3a75a4d2fd Mon Sep 17 00:00:00 2001
From: sugarkwork <sugarkwork@sugar-knight.com>
Date: Tue, 30 Apr 2024 09:08:30 +0900
Subject: [PATCH] init

---
 README.md        | 12 ++++++
 __init__.py      |  3 ++
 nodes.py         | 99 ++++++++++++++++++++++++++++++++++++++++++++++++
 requirements.txt |  1 +
 4 files changed, 115 insertions(+)
 create mode 100644 README.md
 create mode 100644 __init__.py
 create mode 100644 nodes.py
 create mode 100644 requirements.txt

diff --git a/README.md b/README.md
new file mode 100644
index 0000000..582ad64
--- /dev/null
+++ b/README.md
@@ -0,0 +1,12 @@
+# comfyui_cohere
+
+ComfyUI から cohere （Command R+）を使用するためのノードです。
+
+
+ComfyUI の起動 bat ファイルを編集して、以下のように Cohere から取得した API キーを記述する必要があります。
+
+    @echo off
+    set COHERE_API_KEY=COHERE_API_KEY
+    .\python_embeded\python.exe -s ComfyUI\main.py --windows-standalone-build
+    pause
+
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..4a30922
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,3 @@
+from .nodes import *
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
diff --git a/nodes.py b/nodes.py
new file mode 100644
index 0000000..6d6f80b
--- /dev/null
+++ b/nodes.py
@@ -0,0 +1,99 @@
+import os
+import pickle
+import time
+import threading
+
+import cohere
+
+
+lock = threading.Lock()
+
+
+cache_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),"cache.pkl")
+
+
+cohere_client = cohere.Client(os.getenv("COHERE_API_KEY"))
+
+
+def save_memory(key, val):
+    while True:
+        try:
+            with lock:
+                if os.path.exists(cache_path):
+                    with open(cache_path, 'rb') as f:
+                        memory = pickle.load(f)
+                else:
+                    memory = {}
+                memory[str(key).strip()] = val
+                with open(cache_path, 'wb') as f:
+                    pickle.dump(memory, f)
+            break
+        except Exception as e:
+            print(f"Error saving memory: {e}")
+            time.sleep(1)
+            continue
+
+def load_memory(key, defval=None):
+    while True:
+        try:
+            with lock:
+                if os.path.exists(cache_path):
+                    with open(cache_path, 'rb') as f:
+                        memory = pickle.load(f)
+                else:
+                    memory = {}
+            return memory.get(str(key).strip(), defval)
+        except Exception as e:
+            print(f"Error loading memory: {e}")
+            time.sleep(1)
+            continue
+
+
+class SimpleCohereNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "system": ("STRING", {"default": "You are a friendly AI assistant."}),
+                "text": ("STRING", {"default": "Hello, how are you?"}),
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("result",)
+
+    FUNCTION = "cohere"
+
+    CATEGORY = "text"
+
+    OUTPUT_NODE = True
+
+    def cohere(self, system:str, text:str):
+        key = f"cohere: system={system} , text={text}"
+        response = load_memory(key)
+        if response is not None:
+            return (response, )
+        
+        if not response:
+            response = self.client.chat(
+                chat_history=[
+                    {"role": "SYSTEM", "message": str(system)},
+                ],
+                message=str(text),
+                model="command-r-plus"
+            ).text
+        
+        save_memory(key, response)
+        return (response, )
+
+
+NODE_CLASS_MAPPINGS = {
+    "SimpleCohereNode": SimpleCohereNode,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SimpleCohereNode": "SimpleCohereNode"
+}
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..2e6b115
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1 @@
+cohere