200 lines
7.3 KiB
Python
200 lines
7.3 KiB
Python
import csv
|
|
import json
|
|
import os
|
|
import re
|
|
from pathlib import Path
|
|
from urllib.parse import urlparse
|
|
|
|
import requests
|
|
from bs4 import BeautifulSoup
|
|
from dotenv import load_dotenv
|
|
from openai import OpenAI, BadRequestError
|
|
|
|
load_dotenv(Path(__file__).parent.parent / ".env")
|
|
|
|
|
|
INPUT_FILE = Path(__file__).parent.parent / "wildlife" / "portland_wildlife.csv"
|
|
OUTPUT_FILE = Path(__file__).parent.parent / "wildlife" / "portland_wildlife_dex_entries.csv"
|
|
|
|
SYSTEM_PROMPT = """\
|
|
You are a wildlife encyclopedia writer. When given the common name of a species and its Wikipedia URL, \
|
|
use the fetch_wikipedia_article tool to retrieve the article, then write a concise, engaging dex entry \
|
|
(2-3 short short short sentences) describing the animal's appearance, behavior, or habitat. \
|
|
Write in a factual but vivid tone. Do not mention Wikipedia or cite sources.\
|
|
"""
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Wikipedia tool
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _wiki_mirror_url(wikipedia_url: str) -> str:
|
|
"""Convert a wikipedia.org URL to the local Kiwix mirror URL."""
|
|
# WIKI_LINK looks like: http://192.168.1.50:8085/viewer#wikipedia_en_all_maxi_2025-07/
|
|
wiki_link = os.environ["WIKI_LINK"]
|
|
parsed = urlparse(wiki_link)
|
|
base = f"{parsed.scheme}://{parsed.netloc}"
|
|
book = parsed.fragment.rstrip("/") # e.g. wikipedia_en_all_maxi_2025-07
|
|
|
|
article = wikipedia_url.rstrip("/").split("/wiki/")[-1] # e.g. Gray_treefrog
|
|
return f"{base}/{book}/A/{article}"
|
|
|
|
|
|
def fetch_wikipedia_article(wikipedia_url: str) -> str:
|
|
"""Fetch a Wikipedia article, using the local Kiwix mirror if WIKI_LINK is set.
|
|
Falls back to live Wikipedia if the mirror doesn't have the article.
|
|
Raises requests.HTTPError if neither source has it."""
|
|
wikipedia_url = wikipedia_url.replace(" ", "_")
|
|
|
|
urls_to_try = []
|
|
if os.environ.get("WIKI_LINK"):
|
|
urls_to_try.append(_wiki_mirror_url(wikipedia_url))
|
|
urls_to_try.append(wikipedia_url)
|
|
|
|
response = None
|
|
for url in urls_to_try:
|
|
r = requests.get(url, timeout=10)
|
|
if r.status_code == 404:
|
|
print(f" Not found at {url}, trying next source...")
|
|
continue
|
|
r.raise_for_status()
|
|
response = r
|
|
break
|
|
|
|
if response is None:
|
|
raise requests.HTTPError(f"Article not found in any source: {wikipedia_url}")
|
|
|
|
soup = BeautifulSoup(response.text, "html.parser")
|
|
|
|
# Remove nav/footer noise
|
|
for tag in soup(["nav", "footer", "table", "script", "style"]):
|
|
tag.decompose()
|
|
|
|
# Grab paragraphs from the main content
|
|
content = soup.find(id="mw-content-text") or soup.body
|
|
paragraphs = [p.get_text(" ", strip=True) for p in content.find_all("p") if p.get_text(strip=True)]
|
|
return "\n\n".join(paragraphs)
|
|
|
|
|
|
_TOOLS = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "fetch_wikipedia_article",
|
|
"description": "Fetch the text of a Wikipedia article from the local mirror.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"wikipedia_url": {
|
|
"type": "string",
|
|
"description": "A full en.wikipedia.org URL, e.g. https://en.wikipedia.org/wiki/Gray_treefrog",
|
|
}
|
|
},
|
|
"required": ["wikipedia_url"],
|
|
},
|
|
},
|
|
}
|
|
]
|
|
|
|
_TOOL_HANDLERS = {
|
|
"fetch_wikipedia_article": fetch_wikipedia_article,
|
|
}
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# LLM
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def prompt_llm(prompt: str) -> str:
|
|
"""Send a prompt to the LiteLLM instance, handling tool calls, and return the final text."""
|
|
client = OpenAI(
|
|
base_url=os.environ["LITELLM_URL"],
|
|
api_key=os.environ["LITELLM_API_KEY"],
|
|
)
|
|
messages = [
|
|
{"role": "system", "content": SYSTEM_PROMPT},
|
|
{"role": "user", "content": prompt},
|
|
]
|
|
|
|
while True:
|
|
response = client.chat.completions.create(
|
|
model=os.environ["LITELLM_MODEL"],
|
|
messages=messages,
|
|
tools=_TOOLS,
|
|
tool_choice="auto",
|
|
)
|
|
msg = response.choices[0].message
|
|
|
|
# No tool call — we have the final answer
|
|
if not msg.tool_calls:
|
|
return msg.content
|
|
|
|
# Execute each requested tool and feed results back
|
|
messages.append(msg)
|
|
for tc in msg.tool_calls:
|
|
args = json.loads(tc.function.arguments)
|
|
result = _TOOL_HANDLERS[tc.function.name](**args)
|
|
messages.append({
|
|
"role": "tool",
|
|
"tool_call_id": tc.id,
|
|
"content": result,
|
|
})
|
|
# Some models (e.g. Qwen via LM Studio) require a user message to
|
|
# follow tool results or their jinja prompt template errors out.
|
|
messages.append({"role": "user", "content": "Now write the Pokédex entry."})
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CSV
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def load_species_from_csv(csv_path: str | Path) -> list[dict]:
|
|
"""Read a wildlife CSV and return taxon_id, common_name, and wikipedia_url for each row."""
|
|
species = []
|
|
with open(csv_path, newline="", encoding="utf-8") as f:
|
|
reader = csv.DictReader(f)
|
|
for row in reader:
|
|
species.append({
|
|
"taxon_id": row["taxon_id"],
|
|
"common_name": row["common_name"],
|
|
"wikipedia_url": row["wikipedia_url"],
|
|
})
|
|
return species
|
|
|
|
if __name__ == "__main__":
|
|
species = load_species_from_csv(INPUT_FILE)
|
|
|
|
# Load already-processed taxon_ids so we can resume interrupted runs
|
|
completed_ids = set()
|
|
if OUTPUT_FILE.exists():
|
|
with open(OUTPUT_FILE, newline="", encoding="utf-8") as f:
|
|
for row in csv.DictReader(f):
|
|
completed_ids.add(row["taxon_id"])
|
|
print(f"Resuming — {len(completed_ids)} entries already written.")
|
|
|
|
# Open in append mode; write header only if the file is new
|
|
write_header = not OUTPUT_FILE.exists() or OUTPUT_FILE.stat().st_size == 0
|
|
with open(OUTPUT_FILE, "a", newline="", encoding="utf-8") as f:
|
|
writer = csv.writer(f)
|
|
if write_header:
|
|
writer.writerow(["taxon_id", "common_name", "description"])
|
|
|
|
total = len(species)
|
|
for i, specie in enumerate(species, start=1):
|
|
taxon_id = specie["taxon_id"]
|
|
if taxon_id in completed_ids:
|
|
print(f"[{i}/{total}] Skipping {specie['common_name']} (already done)")
|
|
continue
|
|
|
|
print(f"[{i}/{total}] Generating entry for {specie['common_name']}...")
|
|
try:
|
|
description = prompt_llm(specie["wikipedia_url"])
|
|
except requests.HTTPError:
|
|
print(f" No article found, skipping.")
|
|
continue
|
|
except BadRequestError as e:
|
|
print(f" Bad request error, skipping: {e}")
|
|
continue
|
|
description = re.sub(r"[^A-Za-z0-9.,!?'\-: ]", "", description)
|
|
print(f" {description}")
|
|
writer.writerow([taxon_id, specie["common_name"], description])
|
|
f.flush()
|
|
|
|
print(f"\nDone. Written to {OUTPUT_FILE}") |