fda-drug-lookup / sql_retriever.py
GPUburnout's picture
Upload sql_retriever.py with huggingface_hub
c554cbe verified
Raw
History Blame Contribute Delete
9.55 kB
"""
SQL-based FDA label retriever (Option D).
Replaces ChromaDB / vector search. Exact-match key lookup via SQLite.
Resolution order:
1. Brand-keyed query (if brand provided) -> returns brand-specific label
2. Generic-keyed query -> returns label by generic name
3. Aliases table fallback -> resolves salt forms, common misspellings, etc.
4. Returns None if not found
Usage:
from sql_retriever import FDALabelRetriever
r = FDALabelRetriever("fda_labels.db")
label = r.fetch_label(name="atorvastatin", brand="Lipitor")
section_text = r.get_section(label["set_id"], "drug_interactions")
"""
import sqlite3
from pathlib import Path
from typing import Optional
# Standard FDA SPL section column names (must match schema.sql)
SPL_SECTIONS = [
"indications_and_usage",
"dosage_and_administration",
"dosage_forms_and_strengths",
"contraindications",
"warnings_and_precautions",
"boxed_warning",
"adverse_reactions",
"drug_interactions",
"use_in_specific_populations",
"drug_abuse_and_dependence",
"overdosage",
"description",
"clinical_pharmacology",
"clinical_studies",
]
class FDALabelRetriever:
"""SQL-based exact-match FDA label retrieval."""
def __init__(self, db_path: str | Path):
self.db_path = Path(db_path)
if not self.db_path.exists():
raise FileNotFoundError(f"FDA label database not found: {self.db_path}")
self.conn = sqlite3.connect(str(self.db_path))
self.conn.row_factory = sqlite3.Row
def close(self):
self.conn.close()
def __enter__(self):
return self
def __exit__(self, *args):
self.close()
def fetch_label(self, name: str, brand: Optional[str] = None) -> Optional[dict]:
"""
Fetch a single FDA label by drug name and optional brand.
Resolution order:
1. Brand-keyed exact match (if brand provided)
2. Generic-keyed exact match
3. Aliases table fallback
Returns the full label row as dict, or None if not found.
"""
# 1. Brand-keyed exact match
if brand:
row = self.conn.execute(
"SELECT * FROM fda_labels WHERE LOWER(brand_name) = LOWER(?)",
(brand,),
).fetchone()
if row:
return dict(row)
# 1b. Brand-keyed prefix match (handles "Wellbutrin" -> "WELLBUTRIN XL",
# "Tylenol" -> "TYLENOL Extra Strength", etc.)
# Order by section completeness to prefer the most fully-populated label.
row = self.conn.execute(
"""
SELECT *
FROM fda_labels
WHERE LOWER(brand_name) LIKE LOWER(?) || '%'
ORDER BY
-- prefer shortest brand_name match (closer to user's query)
LENGTH(brand_name),
-- then prefer most-populated label
(CASE WHEN indications_and_usage IS NOT NULL THEN 1 ELSE 0 END +
CASE WHEN dosage_and_administration IS NOT NULL THEN 1 ELSE 0 END +
CASE WHEN adverse_reactions IS NOT NULL THEN 1 ELSE 0 END +
CASE WHEN drug_interactions IS NOT NULL THEN 1 ELSE 0 END) DESC
LIMIT 1
""",
(brand,),
).fetchone()
if row:
return dict(row)
# 2. Generic-keyed lookup
row = self.conn.execute(
"SELECT * FROM fda_labels WHERE LOWER(generic_name) = LOWER(?)",
(name,),
).fetchone()
if row:
return dict(row)
# 3. Aliases table fallback (resolves salt forms, etc.)
# Pick the alias with the most metadata-complete label
row = self.conn.execute(
"""
SELECT l.*
FROM fda_labels l
JOIN drug_aliases a ON l.set_id = a.canonical_set_id
WHERE LOWER(a.alias) = LOWER(?)
ORDER BY
-- Prefer labels with brand match if user gave brand
CASE WHEN ? IS NOT NULL AND LOWER(l.brand_name) = LOWER(?) THEN 0 ELSE 1 END,
-- Then prefer labels with more populated sections
(CASE WHEN l.indications_and_usage IS NOT NULL THEN 1 ELSE 0 END +
CASE WHEN l.dosage_and_administration IS NOT NULL THEN 1 ELSE 0 END +
CASE WHEN l.adverse_reactions IS NOT NULL THEN 1 ELSE 0 END +
CASE WHEN l.drug_interactions IS NOT NULL THEN 1 ELSE 0 END +
CASE WHEN l.warnings_and_precautions IS NOT NULL THEN 1 ELSE 0 END) DESC
LIMIT 1
""",
(name, brand, brand),
).fetchone()
return dict(row) if row else None
def get_label_by_set_id(self, set_id: str) -> Optional[dict]:
"""Fetch a full label row by its canonical set_id.
Phase 11: the autocomplete matcher hands the UI an exact set_id for the
confirmed candidate, so the confirm->fetch step retrieves by that key
directly instead of re-resolving (name, brand). No ambiguity, no fallback.
"""
row = self.conn.execute(
"SELECT * FROM fda_labels WHERE set_id = ?",
(set_id,),
).fetchone()
return dict(row) if row else None
def get_section(self, set_id: str, section_name: str) -> Optional[str]:
"""
Fetch a single section's verbatim text by set_id + section name.
section_name must be one of SPL_SECTIONS.
"""
if section_name not in SPL_SECTIONS:
raise ValueError(f"Unknown section: {section_name}. Must be one of {SPL_SECTIONS}")
row = self.conn.execute(
f"SELECT {section_name} FROM fda_labels WHERE set_id = ?",
(set_id,),
).fetchone()
return row[section_name] if row and row[section_name] else None
def list_available_sections(self, set_id: str) -> list[str]:
"""Return list of section names that have content for this set_id."""
cols = ", ".join(SPL_SECTIONS)
row = self.conn.execute(
f"SELECT {cols} FROM fda_labels WHERE set_id = ?",
(set_id,),
).fetchone()
if not row:
return []
return [s for s in SPL_SECTIONS if row[s]]
def search_aliases(self, query: str) -> list[dict]:
"""Return all aliases matching a query string (case-insensitive)."""
rows = self.conn.execute(
"""SELECT a.alias, a.alias_type, a.source, l.generic_name, l.brand_name, a.canonical_set_id
FROM drug_aliases a JOIN fda_labels l ON a.canonical_set_id = l.set_id
WHERE LOWER(a.alias) = LOWER(?)""",
(query,),
).fetchall()
return [dict(r) for r in rows]
def get_dual_brand_audit(self, generic: Optional[str] = None) -> list[dict]:
"""Return dual-brand audit entries, optionally filtered by generic."""
if generic:
rows = self.conn.execute(
"SELECT * FROM dual_brand_audit WHERE LOWER(generic_name) = LOWER(?)",
(generic,),
).fetchall()
else:
rows = self.conn.execute(
"SELECT * FROM dual_brand_audit ORDER BY generic_name, brand_name"
).fetchall()
return [dict(r) for r in rows]
def stats(self) -> dict:
"""Return database statistics."""
return {
"label_count": self.conn.execute("SELECT COUNT(*) FROM fda_labels").fetchone()[0],
"alias_count": self.conn.execute("SELECT COUNT(*) FROM drug_aliases").fetchone()[0],
"dual_brand_count": self.conn.execute("SELECT COUNT(*) FROM dual_brand_audit").fetchone()[0],
"dual_brand_matched": self.conn.execute(
"SELECT COUNT(*) FROM dual_brand_audit WHERE resolution = 'distinct_label'"
).fetchone()[0],
"dual_brand_missing": self.conn.execute(
"SELECT COUNT(*) FROM dual_brand_audit WHERE resolution = 'needs_fetch'"
).fetchone()[0],
}
if __name__ == "__main__":
# Quick smoke test
r = FDALabelRetriever("fda_labels.db")
print("Stats:", r.stats())
print("\n=== Test: dual-brand resolution (Ozempic vs Wegovy) ===")
for brand in ["Ozempic", "Wegovy"]:
label = r.fetch_label(name="semaglutide", brand=brand)
if label:
ind = label.get("indications_and_usage") or ""
print(f" {brand}: brand_name={label['brand_name']}")
print(f" indications excerpt: {ind[:150]}...")
else:
print(f" {brand}: NOT FOUND")
print("\n=== Test: salt-form alias resolution (BUPROPION) ===")
label = r.fetch_label(name="BUPROPION")
if label:
print(f" generic_name={label['generic_name']}, brand={label['brand_name']}")
else:
print(" NOT FOUND")
print("\n=== Test: brand-only via alias (Lipitor) ===")
label = r.fetch_label(name="atorvastatin") # generic name
if label:
print(f" generic_name={label['generic_name']}, brand={label['brand_name']}, set_id={label['set_id'][:8]}")
else:
print(" Not found by generic name 'atorvastatin'")
# Try with calcium salt form
label = r.fetch_label(name="atorvastatin calcium")
if label:
print(f" Found via 'atorvastatin calcium': generic={label['generic_name']}")
r.close()