Update agent.py

agent.py

@@ -1,42 +1,36 @@
-import requests
-from typing import Dict
-from retriever import retrieve_context  # ✅ Corrected import
-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="tiiuae/falcon-7b-instruct")
-from transformers import pipeline
-
-# Load once globally for efficiency
-qa_pipeline = pipeline("text2text-generation", model="Qwen/Qwen1.5-1.8B-Chat")
-
-
-    
-def simple_llm_call(prompt: str) -> str:
-    """
-    Send the prompt to the model and return the generated response.
-    """
-    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()
-
+from tools import get_tools
+from retriever import retrieve_context
+from config import LLM_MODEL
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+class Agent:
+    def __init__(self):
+        self.model = AutoModelForCausalLM.from_pretrained(
+            LLM_MODEL,
+            device_map="auto",
+            trust_remote_code=True
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
+        self.generator = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer)
+        self.tools = get_tools()
+
+    def generate_answer(self, question: str, context: str = "") -> str:
+        prompt = f"""
+You are an expert AI agent answering academic and logical questions concisely.
+Use the context below to help answer the user's question.
+
+Context:
+{context}
+
+Question:
+{question}
+
+Answer:
+"""
+        outputs = self.generator(prompt, max_new_tokens=100, do_sample=False)
+        return outputs[0]['generated_text'].split("Answer:")[-1].strip()
+
+    def run(self, task: dict) -> str:
+        question = task.get("question", "")
+        context = retrieve_context(task)
+        return self.generate_answer(question, context)