Create test_case_generator.py

test_case_generator.py

@@ -0,0 +1,125 @@
+import os
+import json
+import pandas as pd
+from langchain_openai import ChatOpenAI
+from langchain_core.prompts import PromptTemplate
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+class TestCaseGenerator:
+    def __init__(self, api_key=None):
+        # Allow API key to be passed in or read from environment
+        if api_key:
+            os.environ["OPENAI_API_KEY"] = api_key
+        
+        # Predefined question types
+        self.available_question_types = [
+            'cause_and_effect_reasoning', 
+            'temporal_reasoning', 
+            'object_affordance',
+            'adversarial_tasks',
+            'common_sense_reasoning',
+            'hallucination',
+            'sycophancy'
+        ]
+
+    def load_and_split_document(self, doc, chunk_size=1000, chunk_overlap=100):
+        """Load and split the document into manageable chunks."""
+        # Support both file path and uploaded file
+        if isinstance(doc, str):
+            loader = PyPDFLoader(doc)
+            docs = loader.load()
+        else:
+            # Assume it's a BytesIO object from Streamlit upload
+            with open('temp_uploaded_file.pdf', 'wb') as f:
+                f.write(doc.getvalue())
+            loader = PyPDFLoader('temp_uploaded_file.pdf')
+            docs = loader.load()
+        
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+            length_function=len,
+            is_separator_regex=False
+        )
+        return text_splitter.split_documents(docs)
+
+    def get_prompt_template(self, question_type):
+        """Get the prompt template for the given question type."""
+        prompts = {
+            "cause_and_effect_reasoning": cause_and_effect_reasoning,
+            "temporal_reasoning": temporal_reasoning,
+            "object_affordance": object_affordance,
+            # Add other prompts as needed
+        }
+        return prompts.get(question_type, None)
+
+    def extract_json_from_response(self, llm_response):
+        """Clean and extract JSON from LLM response."""
+        llm = ChatOpenAI(temperature=0.25, model="gpt-3.5-turbo")
+        clean_prompt = """
+        You're a highly skilled JSON validator and formatter. 
+        Convert the following text into a valid JSON format:
+        {input_json}
+        
+        Ensure the output follows this structure:
+        {{
+            "questions": [
+                {{
+                    "id": 1,
+                    "question": "...",
+                    "answer": "..."
+                }}
+            ]
+        }}
+        """
+        
+        prompt_template = PromptTemplate.from_template(clean_prompt)
+        final = prompt_template.format(input_json=llm_response)
+        return llm.invoke(final).content
+
+    def convert_qa_to_df(self, llm_response):
+        """Convert LLM response to a pandas DataFrame."""
+        try:
+            if isinstance(llm_response, str):
+                data = json.loads(llm_response)
+            else:
+                data = llm_response
+
+            questions_data = data.get('questions', [])
+            return pd.DataFrame(questions_data)[['question', 'answer']]
+        except Exception as e:
+            print(f"Error processing response: {e}")
+            return pd.DataFrame()
+
+    def generate_testcases(self, doc, question_type, num_testcases=10, temperature=0.7):
+        """Generate test cases for a specific question type."""
+        docs = self.load_and_split_document(doc)
+        model = ChatOpenAI(temperature=temperature, model="gpt-3.5-turbo")
+        prompt = self.get_prompt_template(question_type)
+        
+        if prompt is None:
+            raise ValueError(f"Invalid question type: {question_type}")
+
+        prompt_template = PromptTemplate.from_template(prompt)
+        testset_df = pd.DataFrame(columns=['question', 'answer', 'question_type'])
+        question_count = 0
+
+        for doc_chunk in docs:
+            if question_count >= num_testcases:
+                break
+                
+            final_formatted_prompt = prompt_template.format(context=doc_chunk.page_content)
+
+            response = model.invoke(final_formatted_prompt).content
+
+            try:
+                cleaned_json = self.extract_json_from_response(response)
+                df = self.convert_qa_to_df(cleaned_json)
+                df['question_type'] = question_type
+                testset_df = pd.concat([testset_df, df], ignore_index=True)
+                question_count += len(df)
+            except Exception as e:
+                print(f"Error generating questions: {e}")
+
+        return testset_df.head(num_testcases)