Update agents.py

agents.py

@@ -1,32 +1,41 @@
-from transformers import pipeline
-from prompts import *
-
-generator = pipeline(
-    "text-generation",
-    model="HuggingFaceH4/zephyr-7b-beta",
-    max_new_tokens=300,
-    temperature=0.7
+from huggingface_hub import InferenceClient
+import os
+
+# HF automatically injects token
+client = InferenceClient(
+    token=os.getenv("HF_TOKEN")
 )
 
+MODEL = "HuggingFaceH4/zephyr-7b-beta"
+
+
 def call_llm(prompt):
-    return generator(prompt)[0]["generated_text"]
 
-def market_agent(problem, memory):
-    prompt = MARKET_PROMPT.format(problem=problem, memory=memory)
-    return call_llm(prompt)
+    response = client.text_generation(
+        model=MODEL,
+        prompt=prompt,
+        max_new_tokens=400,
+        temperature=0.7
+    )
+
+    return response
+
+
+def market_agent(problem, memory, prompt):
+    return call_llm(prompt.format(problem=problem, memory=memory))
+
+
+def finance_agent(problem, memory, prompt):
+    return call_llm(prompt.format(problem=problem, memory=memory))
+
+
+def risk_agent(problem, memory, prompt):
+    return call_llm(prompt.format(problem=problem, memory=memory))
 
-def finance_agent(problem, memory):
-    prompt = FINANCE_PROMPT.format(problem=problem, memory=memory)
-    return call_llm(prompt)
 
-def risk_agent(problem, memory):
-    prompt = RISK_PROMPT.format(problem=problem, memory=memory)
-    return call_llm(prompt)
+def ethics_agent(problem, memory, prompt):
+    return call_llm(prompt.format(problem=problem, memory=memory))
 
-def ethics_agent(problem, memory):
-    prompt = ETHICS_PROMPT.format(problem=problem, memory=memory)
-    return call_llm(prompt)
 
-def synthesiser_agent(problem, memory):
-    prompt = SYNTHESIS_PROMPT.format(problem=problem, memory=memory)
-    return call_llm(prompt)
+def synthesis_agent(problem, memory, prompt):
+    return call_llm(prompt.format(problem=problem, memory=memory))