import git
import unittest
from environs import Env
from langfuse import Langfuse
from langchain.chat_models import ChatOpenAI
from toxicity import load_toxicity_check_chain
from observability import trace, span, previous_observation


githash = git.Repo(search_parent_directories=True).head.object.hexsha

class TestToxicity(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        env = Env()
        env.read_env()
        llm = ChatOpenAI(
            temperature=0, 
            openai_api_key=env.str("OPENAI_API_KEY"),
        )
        self.chain = load_toxicity_check_chain(env, llm)

    @trace(
        name="test toxicity check", 
        release=githash,
    )
    @span(name="toxicity chain")
    async def toxicity_check(self, instruction: str) -> bool:
        return await self.chain.apredict_and_parse(instruction=instruction, stop=['\n'])

    async def test_check(self):
        dataset = Langfuse().get_dataset("toxicity")
        for item in dataset.items:
            has_toxicity = await self.toxicity_check(item.input)
            obs = previous_observation()
            item.link(obs, "test_toxicity_check:"+githash[:8])
            obs.score(
                name = "exact_match",
                value = has_toxicity == item.expected_output,
            )