- 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:
1 parent
7a91a9bed4
commit
f01728adb7
3 files changed
+638
No files matched your search
@@ -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}")
|
||||
Reference in new issue
Block a user