agent.py

import os
from tools import TOOLS
from metadata import load_metadata
from mistral_hf_wrapper import MistralInference

# Load Mistral endpoint and token from environment
API_URL = os.getenv("HF_MISTRAL_ENDPOINT")
API_TOKEN = os.getenv("HF_TOKEN")

# Load all tasks from metadata.jsonl
def load_tasks():
    return load_metadata("metadata.jsonl")

# Initialize MistralInference wrapper
mistral = MistralInference(api_url=API_URL, api_token=API_TOKEN)

# Solve a single task
def solve_task(task, tools=TOOLS):
    system_prompt = "You are a helpful agent. Use reasoning, tools if needed, and return the answer only."
    user_prompt = task["question"]

    prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n{user_prompt} [/INST]"
    response = mistral.run(prompt)

    return {
        "task_id": task["question_id"],
        "submitted_answer": response.strip()
    }