import chainlit as cl
import os
from langflow.load import load_flow_from_json
from langflow.processing.process import run_graph

google_api_key=os.environ.get("GOOGLE_API_KEY")
composio_api_key=os.environ.get("COMPOSIO_API_KEY")

TWEAKS = {
  "ChatInput-2x8vF": {},
  "Prompt-CLQNW": {},
  "Agent-6ilGB": {"api_key":google_api_key},
  "Agent-e7sUF": {"api_key":google_api_key},
  "Agent-womuY": {"api_key":google_api_key},
  "ComposioAPI-RfpkY": {"api_key":composio_api_key},
  "Agent-g9eV8": {"api_key":google_api_key},
  "ChatOutput-IIhYR": {},
  "Prompt-Od4Ay": {},
  "Agent-CjuzA": {"api_key":google_api_key}
}

@cl.on_message
async def main(message: cl.Message):
    flow = load_flow_from_json("AgentFlow2.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()