Update agents.py

agents.py

@@ -1,37 +1,19 @@
-import os
-from huggingface_hub import InferenceClient
+from transformers import pipeline
 
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-MODEL_NAME = "google/gemma-2-2b-it"
-
-client = InferenceClient(
-    model=MODEL_NAME,
-    token=HF_TOKEN
+# Load once (cached by HF Spaces)
+generator = pipeline(
+    "text2text-generation",
+    model="google/flan-t5-base",
+    max_length=1024
 )
 
 
 def run_agent(system_prompt: str, user_input: str) -> str:
-    """
-    Runs a conversational LLM call.
-    This is REQUIRED for Gemma instruction models.
-    """
-
-    messages = [
-        {
-            "role": "system",
-            "content": system_prompt
-        },
-        {
-            "role": "user",
-            "content": user_input
-        }
-    ]
-
-    response = client.chat_completion(
-        messages=messages,
-        max_tokens=1800,
-        temperature=0.2
-    )
-
-    return response.choices[0].message.content
+    prompt = f"""
+{system_prompt}
+
+USER INPUT:
+{user_input}
+"""
+    output = generator(prompt, do_sample=False)[0]["generated_text"]
+    return output