qcm_creator / agent /graph.py
dav74's picture
Update agent/graph.py
2c0f32b verified
import os
from typing import List, Optional, TypedDict
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from langgraph.graph import StateGraph, END
from pydantic import BaseModel, Field
# Define the output structure for the LLM
class QuestionOutput(BaseModel):
question_text: str = Field(description="The text of the question")
options: List[str] = Field(description="A list of options including distractors and the correct answer")
correct_option_index: int = Field(description="The index of the correct option in the options list")
# Define the state of the graph
class AgentState(TypedDict):
qcm_description: str
existing_questions: List[str]
document_content: Optional[str]
generated_question: Optional[dict]
# Initialize the LLM
# Ensure OPENAI_API_KEY is set in the environment
#GOOGLE_API_KEY = ""
try:
llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash-lite", temperature=0.7)
except Exception as e:
print(f"Warning: Failed to initialize ChatOpenAI: {e}")
llm = None
def generate_question_node(state: AgentState):
"""
Generates a new question based on the description and existing questions.
"""
if not llm:
return {"generated_question": None}
description = state["qcm_description"]
existing = "\n".join([f"- {q}" for q in state["existing_questions"]])
parser = JsonOutputParser(pydantic_object=QuestionOutput)
document_content = state.get("document_content", "")
prompt_text = "QCM Description: {description}\n\n"
if document_content:
prompt_text += "Document Content:\n{document_content}\n\n"
prompt_text += "Existing Questions:\n{existing}\n\n"
prompt_text += "Generate a new question:\n{format_instructions}"
prompt = ChatPromptTemplate.from_messages([
("system", "You are an expert teacher assistant helping to create a Multiple Choice Question (QCM). "
"Your goal is to generate a NEW, unique question based on the QCM description and the provided document content (if any). "
"Avoid duplicating any of the existing questions. "
"Provide the output in JSON format with 'question_text', 'options' (list of strings), and 'correct_option_index' (int)."),
("user", prompt_text)
])
chain = prompt | llm | parser
try:
result = chain.invoke({
"description": description,
"existing": existing,
"document_content": document_content,
"format_instructions": parser.get_format_instructions()
})
return {"generated_question": result}
except Exception as e:
# Fallback or error handling
print(f"Error generating question: {e}")
return {"generated_question": None}
# Build the graph
workflow = StateGraph(AgentState)
workflow.add_node("generate_question", generate_question_node)
workflow.set_entry_point("generate_question")
workflow.add_edge("generate_question", END)
app = workflow.compile()