Update agent.py

agent.py

@@ -11,11 +11,7 @@ from transformers import pipeline
 # Load once globally for efficiency
 qa_pipeline = pipeline("text2text-generation", model="Qwen/Qwen1.5-1.8B-Chat")
 
-def run_agent_on_question(question_obj):
-    question = question_obj["question"]
-    prompt = f"Answer this question:\n{question}"
-    result = qa_pipeline(prompt, max_new_tokens=50, do_sample=False)[0]["generated_text"]
-    return result.strip().replace("Answer: ", "").split("\n")[0]
+
     
 def simple_llm_call(prompt: str) -> str:
     """
@@ -24,4 +20,23 @@ def simple_llm_call(prompt: str) -> str:
     out = pipe(prompt, max_new_tokens=512)[0]['generated_text']
     return out
 
+# agent.py
+def run_agent_on_question(task: dict) -> str:
+    """
+    task = {
+        'task_id': 'abc123',
+        'question': 'What is the capital of France?'
+    }
+    """
+    question = task['question']
+    
+    # ✅ Your LLM prompt logic goes here (minimal working example)
+    import transformers
+    from transformers import pipeline
+
+    llm = pipeline("text-generation", model="tiiuae/falcon-7b-instruct", max_new_tokens=100)
+    result = llm(question)[0]['generated_text']
+
+    # Return a trimmed response (just the answer, no explanation, no prefix)
+    return result.strip()