Supabase SQL Agent

1
from composio import Composio
2
from composio_llamaindex import LlamaIndexProvider
3
4
from llama_index.llms.openai import OpenAI
5
from llama_index.core.agent.workflow import FunctionAgent
6
from llama_index.core.workflow import Context
7
8
9
def create_agent(user_id: str, composio_client: Composio[LlamaIndexProvider]):
10
   """
11
   Create a function agent that can perform Supabase tasks.
12
   """
13
14
   # Setup client
15
   llm = OpenAI(model="gpt-4-turbo")
16
17
   # Get All the tools
18
   tools = composio_client.tools.get(
19
      user_id=user_id,
20
      tools=[
21
         "SUPABASE_LIST_ALL_PROJECTS",
22
         "SUPABASE_BETA_RUN_SQL_QUERY",
23
      ],
24
   )
25
26
   agent = FunctionAgent(
27
      tools=tools,
28
      llm=llm,
29
      system_prompt=(
30
         "You are a helpful assistant that can help with supabase queries."
31
      ),
32
   )
33
34
   #  Since this is a continuosly running agent, we need to maintain state across
35
   # different user messages
36
   ctx = Context(workflow=agent)
37
   return agent, ctx

1
# Function to initiate a connected account
2
def create_connection(composio_client: Composio[OpenAIProvider], user_id: str):
3
    """
4
    Create a connection for a given user id and auth config id.
5
    """
6
    # Fetch or create the auth config for the supabase toolkit
7
    auth_config = fetch_auth_config(composio_client=composio_client)
8
    if not auth_config:
9
        auth_config = create_auth_config(composio_client=composio_client)
10
11
    # Create a connection for the user
12
    return composio_client.connected_accounts.initiate(
13
        user_id=user_id,
14
        auth_config_id=auth_config.id,
15
    )

1
def check_connected_account_exists(
2
    composio_client: Composio[LangchainProvider],
3
    user_id: str,
4
):
5
    """
6
    Check if a connected account exists for a given user id.
7
    """
8
    # Fetch all connected accounts for the user
9
    connected_accounts = composio_client.connected_accounts.list(
10
        user_ids=[user_id],
11
        toolkit_slugs=["SUPABASE"],
12
    )
13
14
    # Check if there's an active connected account
15
    for account in connected_accounts.items:
16
        if account.status == "ACTIVE":
17
            return True
18
19
        # Ideally you should not have inactive accounts, but if you do, you should delete them
20
        print(f"[warning] inactive account {account.id} found for user id: {user_id}")
21
    return False

1
async def run_loop(user_id: str):
2
    # Initialize composio client.
3
    composio_client = Composio(provider=LlamaIndexProvider())
4
5
    # Setup connection if required
6
    if not check_connected_account_exists(
7
        composio_client=composio_client,
8
        user_id=user_id,
9
    ):
10
        connection_request = create_connection(
11
            composio_client=composio_client,
12
            user_id=user_id,
13
        )
14
        print(
15
            f"Authenticate with the following link: {connection_request.redirect_url}"
16
        )
17
18
    # Create agent
19
    agent, ctx = create_agent(
20
        user_id=user_id,
21
        composio_client=composio_client,
22
    )
23
24
    # Run a simple REPL loop
25
    while True:
26
        user_input = input("user > ")
27
        if user_input.lower() == "exit":
28
            break
29
30
        result = await agent.run(user_msg=user_input, ctx=ctx)
31
        print("agent > ", result)
32
    print("Exiting...")