Final_Assignment_Template / tools /attachments.py
jwlee-ai's picture
Upload folder using huggingface_hub
4a5f5e9 verified
Raw History Blame Contribute Delete
12.7 kB
"""GAIA ์ฒจ๋ถ€ ํŒŒ์ผ ์ฒ˜๋ฆฌ + ์งˆ๋ฌธโ†”task_id ์ธ๋ฑ์Šค.
CodeAgent์˜ ์‹œ๊ทธ๋‹ˆ์ฒ˜ ์ œ์•ฝ(__call__์ด question๋งŒ ๋ฐ›์Œ) ๋•Œ๋ฌธ์— task_id๋ฅผ ์ง์ ‘
์ฃผ์ž…ํ•  ์ˆ˜ ์—†์–ด, ๋ชจ๋“ˆ ์ „์—ญ mutable ์ปจํ…Œ์ด๋„ˆ + prefetch ์ธ๋ฑ์Šค๋กœ ์šฐํšŒํ•œ๋‹ค.
ํ๋ฆ„:
1) BasicAgent.__init__ ๋‹จ๊ณ„์— prefetch_question_index() โ†’ /questions 1ํšŒ ํ˜ธ์ถœ
ํ•ด์„œ {์งˆ๋ฌธ๋ณธ๋ฌธ: task_id} ์‚ฌ์ „์„ ๋งŒ๋“ค๊ณ  set_question_index() ๋กœ ๋“ฑ๋ก.
2) BasicAgent.__call__ ์ง„์ž… ์‹œ set_current_task(question) ์œผ๋กœ ํ˜„์žฌ ๋ฌธ์ œ์˜
task_id์™€ ์งˆ๋ฌธ ๋ณธ๋ฌธ์„ _CURRENT_TASK ์— ์„ธํŒ….
3) ์—์ด์ „ํŠธ๊ฐ€ get_attached_file() ์„ ์ธ์ž ์—†์ด ํ˜ธ์ถœํ•˜๋ฉด _CURRENT_TASK ์˜
task_id๋กœ ์ฑ„์  ์„œ๋ฒ„์—์„œ ํŒŒ์ผ์„ ๋ฐ›์•„์˜ค๊ณ , ํƒ€์ž…๋ณ„๋กœ ์ฒ˜๋ฆฌ:
- ํ…์ŠคํŠธ/CSV/JSON/code: UTF-8 ๋””์ฝ”๋”ฉ
- Excel(.xlsx): ์‹œํŠธ๋ณ„ CSV
- PDF: ํŽ˜์ด์ง€๋ณ„ ํ…์ŠคํŠธ ์ถ”์ถœ (pypdf)
- ์ด๋ฏธ์ง€: VLM(Qwen2.5-VL-7B)์œผ๋กœ ํ˜„์žฌ ์งˆ๋ฌธ ์ปจํ…์ŠคํŠธ์— ๋งž์ถฐ ๋ถ„์„
- ์˜ค๋””์˜ค: Whisper(large-v3) ์ „์‚ฌ
"""
import io
import re
import requests
from smolagents import tool
# ์ฑ„์  ์„œ๋ฒ„ URL์„ ์—ฌ๊ธฐ์„œ๋„ ํ•œ ๋ฒˆ ์ •์˜ (app.py์™€ ๋™์ผ ๊ฐ’).
# tools ๋ชจ๋“ˆ์„ ๋…๋ฆฝ์ ์œผ๋กœ ์‚ฌ์šฉํ•˜๋”๋ผ๋„ ์˜๋ฏธ๊ฐ€ ํ†ตํ•˜๋„๋ก ๋ถ„๋ฆฌํ•ด ๋‘”๋‹ค.
_DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
# BasicAgent.__call__ ์ง„์ž… ์‹œ ๊ฐฑ์‹ ๋˜๋Š” mutable ์ปจํ…Œ์ด๋„ˆ.
# question์€ ์ด๋ฏธ์ง€ VLM ํ˜ธ์ถœ ์‹œ ์ปจํ…์ŠคํŠธ(prompt)๋กœ ์‚ฌ์šฉ๋œ๋‹ค.
_CURRENT_TASK = {"id": None, "question": None}
# question.strip() -> task_id ์‚ฌ์ „.
_QUESTION_INDEX: dict = {}
def prefetch_question_index() -> dict:
"""์ฑ„์  ์„œ๋ฒ„ /questions ๋ฅผ ํ•œ ๋ฒˆ ํ˜ธ์ถœํ•ด {์งˆ๋ฌธ๋ณธ๋ฌธ: task_id} ์‚ฌ์ „์„ ๋นŒ๋“œํ•œ๋‹ค.
์‹คํŒจํ•ด๋„ ๋นˆ dict๋ฅผ ๋ฐ˜ํ™˜ํ•ด ์—์ด์ „ํŠธ๊ฐ€ ์ฒจ๋ถ€ ์—†๋Š” ๋ฌธ์ œ๋งŒ์ด๋ผ๋„ ํ’€ ์ˆ˜ ์žˆ๊ฒŒ ํ•œ๋‹ค."""
try:
r = requests.get(f"{_DEFAULT_API_URL}/questions", timeout=15)
r.raise_for_status()
idx = {}
for item in r.json():
qt = (item.get("question") or "").strip()
tid = item.get("task_id")
if qt and tid:
if qt in idx and idx[qt] != tid:
print(
"Warning: duplicate question text in prefetch index โ€” "
f"task_id {idx[qt]!r} will be overwritten by {tid!r}"
)
idx[qt] = tid
return idx
except Exception as e:
print(f"Warning: could not prefetch question index: {e}")
return {}
def set_question_index(idx: dict) -> None:
"""BasicAgent.__init__์—์„œ prefetch ๊ฒฐ๊ณผ๋ฅผ ๋ชจ๋“ˆ ์ „์—ญ์— ๋ฐ•์•„์ฃผ๋Š” ์„ธํ„ฐ."""
global _QUESTION_INDEX
_QUESTION_INDEX = idx
def set_current_task(question: str):
"""BasicAgent.__call__ ์ง„์ž… ์‹œ ํ˜„์žฌ ๋ฌธ์ œ์˜ task_id์™€ ์งˆ๋ฌธ ๋ณธ๋ฌธ์„ ๋ชจ๋“ˆ ์ „์—ญ์— ์„ธํŒ….
์งˆ๋ฌธ ๋ณธ๋ฌธ์€ ์ด๋ฏธ์ง€ ์ฒจ๋ถ€์˜ VLM ํ˜ธ์ถœ์— prompt ์ปจํ…์ŠคํŠธ๋กœ ์ „๋‹ฌ๋œ๋‹ค.
๋งค์นญ ์‹คํŒจ ์‹œ task_id๋กœ None์ด ๋“ค์–ด๊ฐ€์ง€๋งŒ question์€ ๊ทธ๋Œ€๋กœ ์ €์žฅ๋œ๋‹ค."""
tid = _QUESTION_INDEX.get(question.strip())
_CURRENT_TASK["id"] = tid
_CURRENT_TASK["question"] = question
return tid
# --- ํŒŒ์ผ ํƒ€์ž… ๋ถ„๊ธฐ ํ—ฌํผ ---
def _extract_filename(headers, url: str) -> str:
"""Content-Disposition ํ—ค๋”์—์„œ filename์„ ๋ฝ‘๊ฑฐ๋‚˜, URL ๋๋ถ€๋ถ„์œผ๋กœ ํด๋ฐฑ.
์ฑ„์  ์„œ๋ฒ„๊ฐ€ Content-Type์„ octet-stream์œผ๋กœ ์ค„ ๋•Œ ํ™•์žฅ์ž๋กœ ๋ณด๊ฐ•ํ•˜๊ธฐ ์œ„ํ•จ."""
cd = headers.get("Content-Disposition", "")
# filename* (RFC 5987) ์™€ filename= ์–‘์ชฝ ๋‹ค ์ฒ˜๋ฆฌ.
m = re.search(r'filename\*?=(?:UTF-8\'\')?"?([^";\r\n]+)"?', cd)
if m:
return m.group(1).strip().strip('"')
return url.rsplit("/", 1)[-1]
def _is_excel(content_type: str, ext: str) -> bool:
if ext in ("xlsx", "xls"):
return True
ct = content_type.lower()
return "spreadsheet" in ct or ct.endswith("xlsx") or ct.endswith("xls") or "excel" in ct
def _is_pdf(content_type: str, ext: str) -> bool:
return ext == "pdf" or "pdf" in content_type.lower()
def _is_image(content_type: str, ext: str) -> bool:
return ext in ("png", "jpg", "jpeg", "webp", "gif", "bmp") \
or content_type.lower().startswith("image/")
def _is_audio(content_type: str, ext: str) -> bool:
return ext in ("mp3", "wav", "m4a", "ogg", "flac") \
or content_type.lower().startswith("audio/")
# --- ํƒ€์ž…๋ณ„ ํ•ธ๋“ค๋Ÿฌ ---
def _handle_excel(content: bytes, content_type: str) -> str:
"""xlsx โ†’ ์‹œํŠธ๋ณ„ CSV๋กœ ์ง๋ ฌํ™”. GAIA์— ๋งค์ถœ/ํŒ๋งค ๋ฐ์ดํ„ฐ ๋ฌธ์ œ๊ฐ€ ์ž์ฃผ ๋‚˜์˜จ๋‹ค."""
try:
import pandas as _pd
bio = io.BytesIO(content)
sheets = _pd.read_excel(bio, sheet_name=None)
parts = []
for name, df in sheets.items():
parts.append(f"--- Sheet: {name} ---\n{df.to_csv(index=False)}")
combined = "\n\n".join(parts)
if len(combined) > 12000:
combined = combined[:12000] + "\n...[truncated]"
return f"[Content-Type: {content_type}]\n{combined}"
except Exception as e:
return f"Excel parse error: {e}"
def _handle_pdf(content: bytes, content_type: str) -> str:
"""pypdf๋กœ PDF ๋ณธ๋ฌธ ํ…์ŠคํŠธ ์ถ”์ถœ. ํŽ˜์ด์ง€๋ณ„๋กœ ๊ตฌ๋ถ„ํ•ด์„œ ๋ฐ˜ํ™˜.
์Šค์บ” PDF(์ด๋ฏธ์ง€๋กœ ๋œ)๋Š” ํ…์ŠคํŠธ๊ฐ€ ๋น„๊ฑฐ๋‚˜ ๊นจ์งˆ ์ˆ˜ ์žˆ๋Š”๋ฐ, ๊ทธ ๊ฒฝ์šฐ๋Š”
LLM์ด ์œ„ํ‚ค/์›น๊ฒ€์ƒ‰์œผ๋กœ ํด๋ฐฑํ•˜๋„๋ก ์‹œ์Šคํ…œ ํ”„๋กฌํ”„ํŠธ๊ฐ€ ์œ ๋„ํ•œ๋‹ค."""
try:
from pypdf import PdfReader
bio = io.BytesIO(content)
reader = PdfReader(bio)
parts = []
for i, page in enumerate(reader.pages):
try:
txt = page.extract_text() or ""
except Exception as pe:
txt = f"(extraction failed: {pe})"
parts.append(f"--- Page {i+1} ---\n{txt}")
combined = "\n\n".join(parts)
if len(combined) > 12000:
combined = combined[:12000] + "\n...[truncated]"
return f"[PDF, {len(reader.pages)} pages, Content-Type: {content_type}]\n{combined}"
except Exception as e:
return f"PDF parse error: {e}"
def _handle_image(content: bytes, content_type: str) -> str:
"""VLM(Qwen2.5-VL-7B)์œผ๋กœ ํ˜„์žฌ ์งˆ๋ฌธ ์ปจํ…์ŠคํŠธ์— ๋งž์ถฐ ์ด๋ฏธ์ง€๋ฅผ ๋ถ„์„ํ•œ๋‹ค.
HF Inference API์˜ OpenAI ํ˜ธํ™˜ chat_completion์œผ๋กœ base64 data URL์„ ์ „์†กํ•œ๋‹ค.
์งˆ๋ฌธ ์ปจํ…์ŠคํŠธ๊ฐ€ ์žˆ์œผ๋ฉด ๊ทธ๊ฑธ ๊ทธ๋Œ€๋กœ prompt์— ๋ฐ•์•„ ์ •๋‹ต์— ์ง์ ‘ ๋„์›€์ด ๋˜๋Š”
๋ถ€๋ถ„๋งŒ ๋ฝ‘์•„๋‚ด๋„๋ก ์œ ๋„(generic ์บก์…˜์€ ๋””ํ…Œ์ผ์„ ๋†“์นจ). ํ˜ธ์ถœ ์‹คํŒจ ์‹œ ์—๋Ÿฌ
๋ฌธ์ž์—ด์„ ๋ฐ˜ํ™˜ํ•ด ์—์ด์ „ํŠธ๊ฐ€ ๋‹ค๋ฅธ ์ „๋žต์œผ๋กœ ํด๋ฐฑํ•  ์ˆ˜ ์žˆ๊ฒŒ ํ•œ๋‹ค.
HF_TOKEN ํ™˜๊ฒฝ๋ณ€์ˆ˜๊ฐ€ ํ•„์š”ํ•˜๋‹ค. Space ๋ฐฐํฌ ์‹œ์—๋Š” Space secrets์— ๋“ฑ๋กํ•ด์•ผ ํ•จ.
"""
try:
import base64
from huggingface_hub import InferenceClient
question = (_CURRENT_TASK.get("question") or "").strip()
# ๋ฐ์ดํ„ฐ URL ๊ตฌ์„ฑ. content_type์ด image/* ๊ฐ€ ์•„๋‹ ์ˆ˜๋„ ์žˆ์–ด ์•ˆ์ „ํ•˜๊ฒŒ ํด๋ฐฑ.
mime = content_type.split(";")[0].strip()
if not mime.startswith("image/"):
mime = "image/png"
b64 = base64.b64encode(content).decode("utf-8")
data_url = f"data:{mime};base64,{b64}"
if question:
prompt = (
"Analyze the attached image and answer the following question. "
"Read any text, numbers, or labels visible in the image. "
"If it is a chart or table, extract the relevant data values precisely.\n\n"
f"Question: {question}"
)
else:
prompt = (
"Describe the attached image in detail, including any visible text, "
"numbers, or labels."
)
client = InferenceClient(provider="auto") # HF_TOKEN ํ™˜๊ฒฝ๋ณ€์ˆ˜ ์‚ฌ์šฉ
resp = client.chat_completion(
model="Qwen/Qwen2.5-VL-7B-Instruct",
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": prompt},
{"type": "image_url", "image_url": {"url": data_url}},
],
}
],
max_tokens=1024,
)
analysis = resp.choices[0].message.content
return (
f"[Image analysis (Content-Type: {content_type}, {len(content)} bytes)]\n"
f"{analysis}"
)
except Exception as e:
return (
f"Image attached (Content-Type: {content_type}, {len(content)} bytes). "
f"VLM analysis failed: {e}"
)
def _handle_audio(content: bytes, content_type: str) -> str:
"""Whisper(large-v3)๋กœ ์˜ค๋””์˜ค ์ „์‚ฌ. GAIA ์˜ค๋””์˜ค๋Š” ๋ณดํ†ต ์งง์€ ๋ฐœํ™”๋ผ ํ•œ ๋ฒˆ ํ˜ธ์ถœ๋กœ ์ถฉ๋ถ„.
HF_TOKEN ํ™˜๊ฒฝ๋ณ€์ˆ˜๊ฐ€ ํ•„์š”ํ•˜๋‹ค. Space ๋ฐฐํฌ ์‹œ์—๋Š” Space secrets์— ๋“ฑ๋กํ•ด์•ผ ํ•จ.
"""
try:
from huggingface_hub import InferenceClient
client = InferenceClient(provider="auto")
result = client.automatic_speech_recognition(
audio=content,
model="openai/whisper-large-v3",
)
# huggingface_hub ๋ฒ„์ „์— ๋”ฐ๋ผ dict ๋˜๋Š” dataclass-like ๊ฐ์ฒด๋กœ ๋ฐ˜ํ™˜๋˜๋ฏ€๋กœ
# ์–‘์ชฝ ๋ชจ๋‘ ์ฒ˜๋ฆฌํ•œ๋‹ค.
if hasattr(result, "text"):
transcription = result.text
elif isinstance(result, dict):
transcription = result.get("text", str(result))
else:
transcription = str(result)
return (
f"[Audio transcription (Content-Type: {content_type}, {len(content)} bytes)]\n"
f"{transcription}"
)
except Exception as e:
return (
f"Audio attached (Content-Type: {content_type}, {len(content)} bytes). "
f"Transcription failed: {e}"
)
@tool
def get_attached_file() -> str:
"""Download the file attached to the CURRENT GAIA task and return its content.
Takes no arguments โ€” the current task_id is auto-resolved from the question.
Use this whenever the question references a file, spreadsheet, image, audio, PDF, code listing,
CSV, or any external resource. Returns:
- Text/CSV/JSON/code: the decoded text (truncated to ~12k chars).
- Excel (.xlsx): each sheet rendered as CSV (truncated).
- PDF: extracted text per page (truncated).
- Image (PNG/JPEG/WEBP/GIF/BMP): a vision-language model analysis focused on the current question.
- Audio (MP3/WAV/M4A/OGG/FLAC): a Whisper transcription.
- Other binary: a metadata description (size + content-type).
"""
# ์‹œ๊ทธ๋‹ˆ์ฒ˜ ์ œ์•ฝ ๋•Œ๋ฌธ์— task_id ์ธ์ž๋ฅผ ๋ฐ›์ง€ ์•Š๊ณ , ๋ชจ๋“ˆ ์ „์—ญ _CURRENT_TASK ์—์„œ ๊ฐ€์ ธ์˜จ๋‹ค.
# ์ด ๊ฐ’์€ BasicAgent.__call__ ์ง„์ž… ์‹œ set_current_task()๋กœ ์„ธํŒ…๋œ๋‹ค.
task_id = _CURRENT_TASK.get("id")
if not task_id:
return "No task context available โ€” likely no file attached for this question."
try:
url = f"{_DEFAULT_API_URL}/files/{task_id}"
r = requests.get(url, timeout=30)
if r.status_code == 404:
return "No file attached to this task."
r.raise_for_status()
content_type = r.headers.get("Content-Type", "")
filename = _extract_filename(r.headers, url)
ext = filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
# 1) ๋ช…ํ™•ํ•œ ๋ฐ”์ด๋„ˆ๋ฆฌ ํƒ€์ž…์€ ๋จผ์ € ์ฒ˜๋ฆฌ.
# ์ผ๋ถ€ PDF/SVG๋Š” UTF-8 decode๊ฐ€ ๋˜์–ด๋„ ์›์‹œ ํ…์ŠคํŠธ๋กœ ๋ฐ˜ํ™˜ํ•˜๋ฉด ํ’ˆ์งˆ์ด ํฌ๊ฒŒ ๋–จ์–ด์ง„๋‹ค.
if _is_excel(content_type, ext):
return _handle_excel(r.content, content_type)
if _is_pdf(content_type, ext):
return _handle_pdf(r.content, content_type)
if _is_image(content_type, ext):
return _handle_image(r.content, content_type)
if _is_audio(content_type, ext):
return _handle_audio(r.content, content_type)
# 2) ํ…์ŠคํŠธ ๊ณ„์—ด์ด๋ฉด UTF-8๋กœ ๋ฐ˜ํ™˜.
try:
text = r.content.decode("utf-8")
if len(text) > 12000:
text = text[:12000] + "\n...[truncated]"
return f"[Content-Type: {content_type}]\n{text}"
except UnicodeDecodeError:
pass
# 3) ์•Œ ์ˆ˜ ์—†๋Š” ๋ฐ”์ด๋„ˆ๋ฆฌ โ€” ๋ฉ”ํƒ€๋ฐ์ดํ„ฐ๋งŒ ๋ฐ˜ํ™˜.
return (
f"Binary file (Content-Type: {content_type}, "
f"size: {len(r.content)} bytes). Cannot display as text. "
f"URL: {url}"
)
except Exception as e:
return f"get_attached_file error: {e}"