from maspi.agent_components.llms import OpenAILLM
import yaml
import types

class SafetyFilter:
    def __init__(self, query):
        self.llm = OpenAILLM()
        self.llm.from_config(yaml.safe_load(open("configs/judge.yaml", "r")))
        self.query = query
        self.SAFETY_FILTER_PROMPT = """
        You are a safety filter that extracts only the information from the previous answer that is directly relevant to the given user query, preserving the original wording exactly.

        Task:
        1. Read the user question and the previous answer.
        2. Identify only the sentences, phrases, or bullet points that directly address the user question.
        3. Preserve their original wording exactly — do not paraphrase, summarize, or alter meaning.
        4. Remove any content that is unrelated, speculative, or off-topic with respect to the query.
        5. Maintain the original order of the extracted content as it appears in the previous answer.

        Input:
        User question: "{query}"
        Previous answer: "{previous_answer}"

        Output:
        Return only the relevant excerpts, in their original form and order, without adding any commentary or explanation.
        """

    def __call__(self, previous_answer):
        messages = [
            {"role": "user", "content": self.SAFETY_FILTER_PROMPT.format(query=self.query, previous_answer=previous_answer)}
        ]
        return self.llm.call_llm(messages=messages)
    
    
    