import gradio as gr
from workflow import build_agent
from langchain.memory import ConversationBufferMemory




session={}

def manage_memory(session_id:str):
    if session_id not in session:
        session[session_id]=ConversationBufferMemory(memory_key="chat_history", return_messages=True)
    return session[session_id]



state={
    'query':'',
    'cat':'',
    'bank':'',
    'product':'',
    'search_result':'',
    'final_answer':'',
    'chat_history':[]
}

agent=build_agent(state)

def chat_history(message,history,request:gr.Request):
    state['query']=message
    state["chat_history"] = [f"User: {u}\nAssistant: {a}" for u, a in history if u and a]
    session_id=request.session_hash
    memory=manage_memory(session_id)
    memory.chat_memory.add_user_message(message)
    ans = agent.invoke(state)
    memory.chat_memory.add_ai_message(ans["final_answer"])
    return ans["final_answer"]

with gr.Blocks() as demo:
    gr.ChatInterface(
        chat_history,
        title="💬 Banking Research Agent",
        description="Ask me about banking products, and I'll fetch & compare results for you."
    )

if __name__=="__main__":
    demo.launch()