agent.py

import os
from utils import sanitize_answer, download_file
from typing import Optional

class LLMInterface:
    def __init__(self):
        provider = os.environ.get("LLM_PROVIDER", "openai")
        self.provider = provider
        if provider == "openai":
            import openai
            self.client = openai
        else:
            # optionally support HF Inference here
            raise NotImplementedError("HF Inference not implemented")

    def chat(self, system: str, user: str, max_tokens: int = 256) -> str:
        if self.provider == "openai":
            resp = self.client.ChatCompletion.create(
                model=os.environ.get("OPENAI_MODEL", "gpt-4o-mini"),
                messages=[
                    {"role": "system", "content": system},
                    {"role": "user", "content": user},
                ],
                temperature=0.0,
                max_tokens=max_tokens,
            )
            return resp.choices[0].message.content
        else:
            raise RuntimeError("LLM provider not supported")

class Agent:
    def __init__(self, llm: Optional[LLMInterface] = None):
        self.llm = llm or LLMInterface()

    def build_prompt(self, task: dict) -> (str, str):
        system = (
            "You are a strict QA assistant. Return exactly the correct answer, "
            "no explanation, no prefixes, no suffixes, only one line."
        )
        user = f"Task ID: {task.get('task_id')}\nQuestion: {task.get('question')}\n"
        if task.get("has_file"):
            user += "This task has a file; download it if needed and process it.\n"
        user += "Output only the final answer on one line."
        return system, user

    def answer_task(self, task: dict) -> str:
        system, user = self.build_prompt(task)
        raw = self.llm.chat(system, user, max_tokens=256)
        return sanitize_answer(raw)