import logging
from together import Together
from virtueai.models import VirtueAIModel

class TogetherSafetyModel:
    def __init__(self, safety_model: VirtueAIModel):
        self.safety_model = safety_model.value

    def __call__(self, query: str) -> bool:
        return self.model.safety_check(query)

    def safety_check(self, query: str) -> bool:
        try:
            response = self.together_client.chat.completions.create(
                model=self.safety_model,
                messages=[{"role": "user", "content": query}],
            )
            logging.debug(response)
            verdict = response.choices[0].message.content.strip().lower()
            return verdict.startswith("safe")
        except Exception as exc:  # Fallback: if safety model fails, err on safe side
            # Optionally you could log this, but for now we allow the flow to continue.
            return True
