sango07/Test-Set-Generator

1import os2import json3import pandas as pd4from langchain_openai import ChatOpenAI5from langchain_core.prompts import PromptTemplate6from langchain_community.document_loaders import PyPDFLoader7from langchain_text_splitters import RecursiveCharacterTextSplitter8from prompts import *9 10class TestCaseGenerator:11    def __init__(self, api_key=None):12        # Allow API key to be passed in or read from environment13        if api_key:14            os.environ["OPENAI_API_KEY"] = os.getenv('OPENAI_API_KEY')15        16        # Predefined question types17        self.available_question_types = [18            'hallucination', 19            'conflicting_instructions', 20            'cause_and_effect_reasoning',21            'factually_incorrect_agreement_sycophancy',22            'toxicity'23        ]24 25    def load_and_split_document(self, doc, chunk_size=1000, chunk_overlap=100):26        """Load and split the document into manageable chunks."""27        # Support both file path and uploaded file28        if isinstance(doc, str):29            loader = PyPDFLoader(doc)30            docs = loader.load()31        else:32            # Assume it's a BytesIO object from Streamlit upload33            with open('temp_uploaded_file.pdf', 'wb') as f:34                f.write(doc.getvalue())35            loader = PyPDFLoader('temp_uploaded_file.pdf')36            docs = loader.load()37        38        text_splitter = RecursiveCharacterTextSplitter(39            chunk_size=chunk_size,40            chunk_overlap=chunk_overlap,41            length_function=len,42            is_separator_regex=False43        )44        return text_splitter.split_documents(docs)45 46    def get_prompt_template(self, question_type):47        """Get the prompt template for the given question type."""48        prompts = {49            "hallucination": hallucination,50            "conflicting_instructions": conflicting_instructions,51            "cause_and_effect_reasoning": cause_and_effect_reasoning,52            "factually_incorrect_agreement_sycophancy":factually_incorrect_agreement_sycophancy,53            "toxicity":toxicity54            # Add other prompts as needed55        }56        return prompts.get(question_type, None)57 58    def extract_json_from_response(self, llm_response):59        """Clean and extract JSON from LLM response."""60        llm = ChatOpenAI(temperature=0.25, model="gpt-3.5-turbo")61        clean_prompt = """62        You're a highly skilled JSON validator and formatter. 63        Convert the following text into a valid JSON format:64        {input_json}65        66        Ensure the output follows this structure:67        {{68            "questions": [69                {{70                    "id": 1,71                    "question": "...",72                    "answer": "..."73                }}74            ]75        }}76        """77        78        prompt_template = PromptTemplate.from_template(clean_prompt)79        final = prompt_template.format(input_json=llm_response)80        return llm.invoke(final).content81 82    def convert_qa_to_df(self, llm_response):83        """Convert LLM response to a pandas DataFrame."""84        try:85            if isinstance(llm_response, str):86                data = json.loads(llm_response)87            else:88                data = llm_response89 90            questions_data = data.get('questions', [])91            return pd.DataFrame(questions_data)[['question', 'answer']]92        except Exception as e:93            print(f"Error processing response: {e}")94            return pd.DataFrame()95 96    def generate_testcases(self, doc, question_type, num_testcases=10, temperature=0.7):97        """Generate test cases for a specific question type."""98        docs = self.load_and_split_document(doc)99        model = ChatOpenAI(temperature=temperature, model="gpt-3.5-turbo")100        prompt = self.get_prompt_template(question_type)101        102        if prompt is None:103            raise ValueError(f"Invalid question type: {question_type}")104 105        prompt_template = PromptTemplate.from_template(prompt)106        testset_df = pd.DataFrame(columns=['question', 'answer', 'question_type'])107        question_count = 0108 109        for doc_chunk in docs:110            if question_count >= num_testcases:111                break112                113            final_formatted_prompt = prompt_template.format(context=doc_chunk.page_content)114 115            response = model.invoke(final_formatted_prompt).content116 117            try:118                cleaned_json = self.extract_json_from_response(response)119                df = self.convert_qa_to_df(cleaned_json)120                df['question_type'] = question_type121                testset_df = pd.concat([testset_df, df], ignore_index=True)122                question_count += len(df)123            except Exception as e:124                print(f"Error generating questions: {e}")125 126        return testset_df.head(num_testcases)