Update agents.py

agents.py

@@ -6,7 +6,6 @@ from langgraph.prebuilt import tools_condition
 from langgraph.prebuilt import ToolNode
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_groq import ChatGroq
-from langchain_community.llms import RWKV
 from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader
@@ -14,7 +13,10 @@ 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
+from rwkv.model import RWKV
+from rwkv.utils import PIPELINE, PIPELINE_ARGS
 
 load_dotenv()
 
@@ -166,15 +168,12 @@ def build_graph(provider: str = "rwkv"):
         pth = hf_hub_download(repo_id="BlinkDL/rwkv7-g1", filename=f"{title}.pth")
         model_path = pth.replace(".pth", "")
 
-        raw_llm = RWKV(
-            model=model_path,
-            strategy="cpu fp32",
-            tokens_path="./20B_tokenizer.json"
-        )
+        raw_llm = RWKV(model=model_path, strategy='cuda fp32')
+        pipeline = PIPELINE(raw_llm, "rwkv_vocab_v20230424")
 
         class RWKVWithTools:
-            def __init__(self, llm, system_prompt: str):
-                self.llm = llm
+            def __init__(self, pipeline, system_prompt: str):
+                self.pipeline = pipeline
                 self.system_prompt = system_prompt
                 self.tools = []
 
@@ -204,8 +203,9 @@ def build_graph(provider: str = "rwkv"):
                 )
 
                 prompt = header + convo
-                # delegate to LangChain’s invoke()
-                return self.llm.invoke(prompt)
+                
+                # delegate to RWKV invoke()
+                return self.pipeline.generate(prompt, token_count=200)
 
         llm = RWKVWithTools(raw_llm, system_prompt=system_prompt)
         # --- END RWKV SETUP ---