Use langchain instead of custom code

agents.py

@@ -6,6 +6,7 @@ from langgraph.prebuilt import tools_condition
 from langgraph.prebuilt import ToolNode
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_groq import ChatGroq
+from langchain.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
@@ -174,26 +175,11 @@ def build_graph(provider: str = "rwkv"):
         # --- BEGIN RWKV SETUP ---
         title = "rwkv7-g1-0.1b-20250307-ctx4096"
         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)
+        model_path = pth.replace(".pth", "")
+        llm = RWKV(
+            model=model_path,
+            strategy="cpu fp32", 
+        )
         # --- END RWKV SETUP ---
     else:
         raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'.")