Spaces:
Runtime error
Runtime error
File size: 1,889 Bytes
6b3dac4 85fd043 6b3dac4 85fd043 6b3dac4 85fd043 6b3dac4 85fd043 6b3dac4 85fd043 6b3dac4 85fd043 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | 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)
|