drugs/core/clients/sources/psychwiki.py

from core.clients.base import BaseClient
from core.clients.graphql import GraphQLClient
from core.models import SEI, Dosage, Drug, Effect, Entry, Timing


class PsychWikiClient(GraphQLClient, BaseClient):
    # url = "https://api.psychonautwiki.org"
    search_limit = 5000
    source_name = "Psychonaut Wiki GraphQL API"
    source_type = "DWIKI"
    source_endpoint = "https://api.psychonautwiki.org"
    source_score = 75

    async def update_drugs(self):
        data = await self.get_all_data()
        self.store_data(data)
        return len(data)

    def store_data(self, data):
        """
        Store the data in the database.
        """
        for drug in data["substances"]:
            print("DRUG ITER", drug)
            try:
                drug_obj = Drug.objects.get(name=drug["name"])
            except Drug.DoesNotExist:
                drug_obj = Drug(name=drug["name"])

            if "commonNames" in drug:
                if drug["commonNames"]:
                    print("common names", drug["commonNames"])
                    drug_obj.common_names = ",".join(drug["commonNames"])
            if "class" in drug:
                if drug["class"]:
                    if "psychoactive" in drug["class"]:
                        if drug["class"]["psychoactive"]:
                            drug_obj.drug_class = ",".join(
                                drug["class"]["psychoactive"]
                            )
            try:
                entry = Entry.objects.get(source=self.source, url=drug["url"])
            except Entry.DoesNotExist:
                entry = Entry.objects.create(source=self.source, url=drug["url"])
            if not drug_obj.pk:
                drug_obj.save()
            if entry not in drug_obj.links.all():
                drug_obj.links.add(entry)
            if "roas" in drug:
                for roa in drug["roas"]:
                    if "name" in roa:
                        roa_name = roa["name"]

                        # Parsing dosage information
                        dose = roa["dose"]

                        if dose:
                            dosage_data = {
                                "entry": entry,
                                "roa": roa_name,
                                "unit": dose["units"],
                            }
                            # Check and assign dosage values
                            for dose_type in [
                                "threshold",
                                "light",
                                "common",
                                "strong",
                                "heavy",
                            ]:
                                if dose_type in dose:
                                    if isinstance(dose[dose_type], dict):
                                        dosage_data[f"{dose_type}_lower"] = dose[
                                            dose_type
                                        ]["min"]
                                        dosage_data[f"{dose_type}_upper"] = dose[
                                            dose_type
                                        ]["max"]
                                    else:
                                        dosage_data[f"{dose_type}_lower"] = dose[
                                            dose_type
                                        ]
                                        dosage_data[f"{dose_type}_upper"] = dose[
                                            dose_type
                                        ]

                            # Check if Dosage already exists and is linked to the Drug
                            dosage, created = Dosage.objects.get_or_create(
                                **dosage_data
                            )
                            if created or dosage not in drug_obj.dosages.all():
                                drug_obj.dosages.add(dosage)
                                print("YES DOSAGE", drug_obj.dosages)

                        # Parsing timing information
                        timing = roa["duration"]

                        print("TIMING", timing)
                        # Check and assign timing values
                        if timing:
                            timing_data = {"entry": entry, "roa": roa_name}
                            for time_type in [
                                "onset",
                                "comeup",
                                "peak",
                                "offset",
                                "total",
                            ]:
                                if (
                                    time_type in timing
                                    and timing[time_type] is not None
                                ):
                                    unit = timing[time_type].get("units", "hours")

                                    # Handle case where timing value is a single integer
                                    if isinstance(timing[time_type], int):
                                        lower = timing[time_type]
                                        upper = None
                                    else:
                                        lower = timing[time_type].get("min")
                                        upper = timing[time_type].get("max")

                                    if unit == "minutes" and lower is not None:
                                        lower = lower / 60.0
                                        if upper is not None:
                                            upper = upper / 60.0

                                    timing_data[f"{time_type}_lower"] = lower
                                    timing_data[f"{time_type}_upper"] = upper
                                    timing_data[
                                        "unit"
                                    ] = "HOURS"  # Store all times in hours

                            # Check if Timing already exists and is linked to the Drug
                            timing_obj, created = Timing.objects.get_or_create(
                                **timing_data
                            )
                            if created or timing_obj not in drug_obj.timings.all():
                                drug_obj.timings.add(timing_obj)
            if "effects" in drug:
                # Create or retrieve Effect object linked to the Entry
                effect_obj, effect_created = Effect.objects.get_or_create(entry=entry)

                for effect in drug["effects"]:
                    # Create or retrieve SEI object
                    sei_obj, sei_created = SEI.objects.get_or_create(
                        name=effect["name"],
                        url=effect.get(
                            "url", ""
                        ),  # Using .get() to handle missing urls
                    )

                    # Link SEI object to Effect if not already linked
                    if (
                        sei_created
                        or sei_obj not in effect_obj.subjective_effects.all()
                    ):
                        effect_obj.subjective_effects.add(sei_obj)
                # Link Effect object to Drug if not already linked
                if effect_created or effect_obj not in drug_obj.effects.all():
                    drug_obj.effects.add(effect_obj)

            print("SAVING DRUG", drug_obj)
            drug_obj.save()

    async def get_drugs_list(self):
        """
        Get all drug names from PsychWiki
        """
        body = {"query": "{substances(limit: %i) {name}}" % self.search_limit}
        result = await self.call(
            "?",
            http_method="post",
            data=body,
            schema="GetDrugsList",
            append_slash=False,
        )
        print("RESULT", result)

        return result["data"]

    async def get_all_data(self):
        """
        Get all the drug data from PsychWiki (warning - intensive)
        """
        body = {
            "query": """
{
  substances(limit: %i) {
    name
    url
    featured
    summary
    addictionPotential
    toxicity
    crossTolerances
    commonNames
    class {
      chemical
      psychoactive
    }
    tolerance {
      full
      half
      zero
    }
    roas {
      name
      dose {
        units
        threshold
        heavy
        common {
          min
          max
        }
        light {
          min
          max
        }
        strong {
          min
          max
        }
      }
      duration {
        afterglow {
          min
          max
          units
        }
        comeup {
          min
          max
          units
        }
        duration {
          min
          max
          units
        }
        offset {
          min
          max
          units
        }
        onset {
          min
          max
          units
        }
        peak {
          min
          max
          units
        }
        total {
          min
          max
          units
        }
      }
      bioavailability {
        min
        max
      }
    }
    effects {
      name
      url
    }
    images {
      thumb
      image
    }
    uncertainInteractions {
      name
    }
    unsafeInteractions {
      name
    }
    dangerousInteractions {
      name
    }
  }
}

"""
            % self.search_limit
        }
        result = await self.call(
            "?",
            http_method="post",
            data=body,
            schema="GetDrugsList",
            append_slash=False,
        )

        return result["data"]