Add RWKV support for Agent

agents.py

@@ -13,6 +13,15 @@ from langchain_community.document_loaders import ArxivLoader
 from langchain_community.vectorstores import SupabaseVectorStore
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
+from huggingface_hub import hf_hub_download
+
+# RWKV setup flags — must come before importing RWKV
+os.environ["RWKV_JIT_ON"]   = "1"
+os.environ["RWKV_V7_ON"]    = "1"   # enable RWKV-7
+os.environ["RWKV_CUDA_ON"]  = "0"   # set to "1"
+
+from rwkv.model import RWKV
+from rwkv.utils import PIPELINE
 
 load_dotenv()
 
@@ -144,7 +153,7 @@ tools = [
 ]
 
 # Build graph function
-def build_graph(provider: str = "groq"):
+def build_graph(provider: str = "rwkv"):
     """Build the graph"""
     # Load environment variables from .env file
     if provider == "google":
@@ -161,6 +170,31 @@ def build_graph(provider: str = "groq"):
                 temperature=0,
             ),
         )
+    elif provider == "rwkv":
+        # --- BEGIN RWKV SETUP ---
+        title = "rwkv7-g1"
+        pth = hf_hub_download(repo_id="BlinkDL/rwkv7-g1", filename=f"{title}.pth")
+        # 2) Load RWKV (drop .pth extension for RWKV loader)
+        rwkv_model = RWKV(model=pth.replace(".pth", ""), strategy="cpu fp32")
+        # 3) Build the tokenization + generation pipeline
+        rwkv_pipe = PIPELINE(rwkv_model, "rwkv_vocab_v20230424")
+        # 4) Wrap into a Chat-style interface
+        class ChatRWKV:
+            def __init__(self, pipe):
+                self.pipe = pipe
+            def invoke(self, messages):
+                prompt = "\n".join(m.content for m in messages)
+                return self.pipe(
+                    prompt,
+                    temperature=0.0,
+                    top_p=0.95,
+                    max_tokens=256,
+                )
+            def bind_tools(self, tools):
+                return self
+
+        llm = ChatRWKV(rwkv_pipe)
+        # --- END RWKV SETUP ---
     else:
         raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'.")
     # Bind tools to LLM