- Added dex_info_writer.py with a full pipeline for

generating Pokédex-style species descriptions
  - Loads species from portland_wildlife.csv (taxon_id,
  common_name, wikipedia_url)
  - Fetches Wikipedia articles from a self-hosted Kiwix mirror
  (WIKI_LINK env var); falls back to live Wikipedia if unset
  - Spaces in Wikipedia URLs are normalized to underscores
  before fetching
  - LiteLLM instance prompted via OpenAI-compatible client
  (LITELLM_URL, LITELLM_API_KEY, LITELLM_MODEL env vars)
  - LLM has tool-calling access to fetch_wikipedia_article with
   an agentic loop
  - Output descriptions are sanitized — only alphanumerics and
  . , ! ? ' - : allowed
  - Writes results to portland_wildlife_dex_entries.csv in
  append mode with resume support (skips already-processed
  taxon_ids)
  - Live [i/total] progress counter printed each iteration
  along with the generated description
This commit is contained in:
SuppaiNeko committed 2026-03-25 21:02:30 -07:00
1 parent 7a91a9bed4
commit f01728adb7
3 files changed
+638

No files matched your search

+173
View File
@@ -0,0 +1,173 @@
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
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."""
wikipedia_url = wikipedia_url.replace(" ", "_")
url = _wiki_mirror_url(wikipedia_url) if os.environ.get("WIKI_LINK") else wikipedia_url
response = requests.get(url, timeout=10)
response.raise_for_status()
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,
})
# ---------------------------------------------------------------------------
# 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']}...")
description = prompt_llm(specie["wikipedia_url"])
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}")