pages/Custom.py

import os
from langchain_google_genai import ChatGoogleGenerativeAI
from typing import List, Dict
from typing_extensions import TypedDict
from dotenv import load_dotenv
import streamlit as st
from langgraph.graph import START, StateGraph, END

st.set_page_config(
    page_title="General Information Blog"
)

load_dotenv(override=True)
gemini_api_key = os.getenv("GEMINI_API_KEY")


from blog_content_generator import blog_content_agent
from image_fomatted_blog_generator import image_formatted_blog
from image_prompts_generator import image_prompts


llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash-exp", api_key = gemini_api_key) # type:ignore

class BlogState(TypedDict):   
    final_topic: str
    blog_content: str
    image_formated_blog: str  
    prompts: Dict[str, str]
    
# Graph


# Graph 2: 
builder = StateGraph(BlogState)
builder.add_node("blog_agent", lambda state: blog_content_agent(llm, state))
builder.add_node("image_formatted_blog", lambda state: image_formatted_blog(llm, state))
builder.add_node("image_prompts", lambda state: image_prompts(llm, state))

builder.add_edge(START, "blog_agent")
builder.add_edge("blog_agent", "image_formatted_blog")
builder.add_edge("image_formatted_blog", "image_prompts")
builder.add_edge("image_prompts", END)
graph = builder.compile()



st.title("Generate Custom Title Blog")
input_data = st.text_input("Enter a Title for your blog:")

if input_data:  # Ensure input_data is not empty
    if st.button("Generate Blog"):
        try:
            content = graph.invoke({"final_topic": input_data})
            st.markdown(content.get('image_formated_blog', "No content generated."))  # Avoid KeyError
        except Exception as e:
            st.error(str(e))