import chainlit as cl
import os

from langflow.load import load_flow_from_json
from langflow.processing.process import run_graph

#google_api_key = "AIzaSyAyGDo-S86WPrMce_vIfePN_mh2_68rJao"
#composio_api_key ="s6q0gx5v78ve5290dpdj8"

google_api_key=os.getenv("GOOGLE_API_KEY")
composio_api_key=os.getenv("COMPOSIO_API_KEY")

TWEAKS = {
  "ChatInput-Gbt07": {},
  "ChatOutput-XrgG9": {},
  "ComposioAPI-gmZ8j": {"api_key": composio_api_key},
  "Agent-NB8hJ": {},
  "GoogleGenerativeAIModel-SG7kV": {"api_key": google_api_key},
  "Prompt-oe6ZL": {}
}

@cl.on_message
async def main(message: cl.Message):
    flow = load_flow_from_json("SimpleAgentFlow.json",tweaks=TWEAKS)
    flow.session_id="session_1"
    flow.user_id="user_1"
    response = await run_graph(flow, input_value=message.content, input_type="chat", output_type="chat")
    await cl.Message(response[0].outputs[0].results["message"].text).send()