diff --git a/brancheneinstufung.py b/brancheneinstufung.py index 8a20ec17..fd043810 100644 --- a/brancheneinstufung.py +++ b/brancheneinstufung.py @@ -97,6 +97,17 @@ system_prompt = { } # === WIKIPEDIA DATEN LADEN === +# Positivliste für Wikipedia-Kategorien, die auf Unternehmen hinweisen können +WHITELIST_KATEGORIEN = [ + "Unternehmen", + "Hersteller", + "Produktion", + "Industrie", + "Maschinenbau", + "Technik", + "Dienstleistungsunternehmen" +] + def get_wikipedia_data(name, website_hint=""): begriffe = [name.strip(), " ".join(name.split()[:2])] if website_hint: @@ -109,6 +120,12 @@ def get_wikipedia_data(name, website_hint=""): for title in results: try: page = wikipedia.page(title) + # Titelprüfung verbessern + if any(x in page.title.lower() for x in ["krankenkasse", "versicherung"]): + continue + if website_hint: + if website_hint.lower().split(".")[0] not in html.lower(): + continue if name.lower().split()[0] not in page.title.lower(): continue url = page.url @@ -137,7 +154,10 @@ def get_wikipedia_data(name, website_hint=""): umsatz = match.group(1).replace(",", ".") if not branche: cats = page.categories - branche = cats[0] if cats else "k.A." + if any(any(w.lower() in c.lower() for w in WHITELIST_KATEGORIEN) for c in cats): + branche = cats[0] + else: + return "", "k.A.", "k.A." return url, branche or "k.A.", umsatz or "k.A." except: continue