| import os |
| import json |
| import pandas as pd |
| from typing import List, Optional, Dict, Any |
| from pydantic import BaseModel, Field |
| from langchain_openai import ChatOpenAI |
| from langchain_core.prompts import ChatPromptTemplate |
| from datetime import datetime |
|
|
| |
| class TagUpdate(BaseModel): |
| """Schema for hierarchical tag generation.""" |
| tags: List[str] = Field(description="Full list of tags including inherited parent tags and 1-2 new complex tags.") |
|
|
| class QAResponse(BaseModel): |
| """Schema for MCQs.""" |
| question: str = Field(description="The scientific/complex question based on tags.") |
| choices: List[str] = Field(description="4 distinct multiple choice answers numbered using digits (0, 1, 2, 3).") |
| correct_choice_idx: int = Field(description="Index (0, 1, 2, or 3) of the correct answer.") |
|
|
| |
| class DatasetAgent: |
| def __init__(self, model_name: str = "gpt-4o", api_key: Optional[str] = None): |
| self.model_name = model_name |
| self.api_key = api_key |
| self._init_llm() |
|
|
| def _init_llm(self): |
| try: |
| self.llm = ChatOpenAI(model=self.model_name, temperature=0.7, openai_api_key=self.api_key) |
| self.tag_generator = self.llm.with_structured_output(TagUpdate) |
| self.qa_generator = self.llm.with_structured_output(QAResponse) |
| except Exception as e: |
| print(f"Warning: Could not initialize OpenAI LLM: {e}") |
| self.llm = None |
| self.tag_generator = None |
| self.qa_generator = None |
|
|
| def update_model(self, model_name: str, api_key: Optional[str] = None): |
| if model_name != self.model_name or api_key != self.api_key: |
| self.model_name = model_name |
| self.api_key = api_key |
| self._init_llm() |
|
|
| def generate_blueprint(self, topic: str, depth: int, breadth: int, model_name: str = None, api_key: str = None) -> pd.DataFrame: |
| """ |
| Phase 1: Build the graph structure with tags. |
| """ |
| if model_name or api_key: |
| self.update_model(model_name or self.model_name, api_key) |
| |
| if not self.tag_generator: |
| raise ValueError("OpenAI API key is missing or invalid. Please provide it in the configuration.") |
| |
| blueprint = [] |
| |
| |
| prompt_root = f"Generate a list of 2-3 core conceptual tags for the topic: {topic}." |
| root_res = self.tag_generator.invoke(prompt_root) |
| root_id = "r1" |
| blueprint.append({"id": root_id, "parent_id": None, "tags": root_res.tags}) |
| |
| |
| self._build_depth(blueprint, root_id, root_res.tags, 1, depth, breadth, topic) |
| |
| |
| df = pd.DataFrame(blueprint) |
| df['tags'] = df['tags'].apply(lambda x: json.dumps(x)) |
| return df |
|
|
| def _build_depth(self, blueprint, parent_id, parent_tags, current_depth, max_depth, breadth, topic): |
| if current_depth >= max_depth: |
| return |
| |
| for i in range(breadth): |
| child_id = f"{parent_id}_c{i+1}" |
| prompt = ( |
| f"Topic: {topic}\n" |
| f"Parent Tags: {parent_tags}\n" |
| "Task: Generate a new list of tags that inherits ALL parent tags but adds 1-2 NEW tags " |
| "to make the conceptual complexity slightly higher for a sub-question." |
| ) |
| res = self.tag_generator.invoke(prompt) |
| blueprint.append({"id": child_id, "parent_id": parent_id, "tags": res.tags}) |
| self._build_depth(blueprint, child_id, res.tags, current_depth + 1, max_depth, breadth, topic) |
|
|
| def generate_qa_batch(self, df: pd.DataFrame, model_name: str = None, api_key: str = None) -> pd.DataFrame: |
| """ |
| Phase 2: Generate Questions and Answers based on the blueprint tags. |
| """ |
| if model_name or api_key: |
| self.update_model(model_name or self.model_name, api_key) |
| |
| if not self.qa_generator: |
| raise ValueError("OpenAI API key is missing or invalid. Please provide it in the configuration.") |
| |
| questions = [] |
| choices_list = [] |
| correct_indices = [] |
| |
| for _, row in df.iterrows(): |
| tags = row['tags'] if isinstance(row['tags'], list) else json.loads(row['tags']) |
| prompt = ( |
| f"Create a multiple choice question based on these technical tags: {tags}.\n" |
| "Ensure the question is not too challenging but requires understanding of the specific sub-topics." |
| ) |
| res = self.qa_generator.invoke(prompt) |
| questions.append(res.question) |
| choices_list.append(json.dumps(res.choices)) |
| correct_indices.append(res.correct_choice_idx) |
| |
| df['question'] = questions |
| df['choices'] = choices_list |
| df['correct_idx'] = correct_indices |
| return df |
|
|
| import tempfile |
|
|
| |
| def save_to_json(df: pd.DataFrame, prefix: str) -> str: |
| """ |
| Saves a DataFrame to a temporary JSON file. |
| NamedTemporaryFile(delete=False) ensures the file persists for Gradio to handle the download. |
| """ |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".json", prefix=f"{prefix}_") as tmp: |
| path = tmp.name |
| df.to_json(path, orient="records", indent=4) |
| return path |
|
|
| def load_from_json(file_path) -> pd.DataFrame: |
| try: |
| if hasattr(file_path, 'name'): |
| file_path = file_path.name |
| return pd.read_json(file_path) |
| except Exception as e: |
| print(f"Error loading JSON: {e}") |
| return pd.DataFrame() |
|
|
