Update agent.py

agent.py

@@ -1,13 +1,14 @@
 import requests
 from typing import Dict
-from retriever import retrieve_relevant_context
+from retriever import retrieve_context
 from tools import clean_answer_with_prompt, build_prompt
+from transformers import pipeline
+
+# Load the model once to save time
+pipe = pipeline("text-generation", model="Qwen/Qwen1.5-1.8B-Chat")
 
-# You can later replace this with a more sophisticated tool-calling agent
 def simple_llm_call(prompt: str) -> str:
     """Send the prompt to the model and return the generated response"""
-    from transformers import pipeline
-    pipe = pipeline("text-generation", model="Qwen/Qwen1.5-1.8B-Chat")
     out = pipe(prompt, max_new_tokens=512)[0]['generated_text']
     return out
 
@@ -15,9 +16,11 @@ def run_agent_on_question(task: Dict) -> str:
     """Main agent loop: processes a single task and returns an answer"""
     task_id = task['task_id']
     question = task['question']
+    file_name = task.get('file_name', '')
 
-    # Step 1: Retrieve relevant context
-    context = retrieve_relevant_context(question)
+    # Step 1: Retrieve relevant context (as list of strings)
+    context_list = retrieve_context(task_id, question)
+    context = "\n".join(context_list)
 
     # Step 2: Build a prompt using system instruction + context + question
     prompt = build_prompt(question, context)