import gradio as gr
from langchain_core.prompts import ChatPromptTemplate
#from langchain_openai import ChatOpenAI
from langchain_google_genai import ChatGoogleGenerativeAI
from langgraph.graph import StateGraph, END
from typing import Dict, List

# Assume you have your LangGraph setup here
def agent_node(state):
    prompt = ChatPromptTemplate.from_messages(
        [("user", "{user_input}")]
    )
    #model = ChatOpenAI()
    model = ChatGoogleGenerativeAI(model="gemini-2.0-flash",
        temperature=0,
        max_tokens=None,
        timeout=None,
        max_retries=2,
    # other params...
    )
    response = model.invoke(prompt.format_messages(**state))
    return {"response": response.content}

def create_workflow():
    workflow = StateGraph(Dict)
    workflow.add_node("agent", agent_node)
    workflow.set_entry_point("agent")
    workflow.add_edge("agent", END)
    return workflow.compile()

workflow = create_workflow()

def respond(message, history):
    result = workflow.invoke({"user_input": message})
    return result["response"]

iface = gr.ChatInterface(respond)
iface.launch()

gr.close_all()