WildDex/dev_scripts/dex_info_writer.py

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}")