Connect RAG to SQL Agent

[ravitejajasti]

I am following the SQLAgent tutorial from Langgraph and adding RAG to it.

The idea is that we use RAG to fetch relevant DB table info and make the SQL agent job easier in finding the right table as well as in creating the correct query. However, I am really not sure if I am doing it right. I get errors when I try integrating RAG with SQL Agent.

I am new to Langchain and Langgraph and I believe some of the errors I get are basic.

To retrieve documents, I am using a custom tool

`
@tool
def retrieve_documents(state):
"""Retrieve documents from the vector store based on the query."""
docsearch = OpenSearchVectorSearch(index_name='sqlstruc_slim', # TODO: use the same index-name used in the ingestion script
embedding_function=LiteLLMEmbedder(),
opensearch_url= 'https://' + opensearch_details['host'] + f":{opensearch_details['port']}",
http_auth=(opensearch_details['creds_user'], opensearch_details['creds_pass']),
use_ssl=True,
verify_certs=opensearch_details['verify_certs'],
ssl_assert_hostname=False,
ssl_show_warn=False,
)

  question = state["messages"][0]
  documents = get_context(docsearch, question)
  return {'messages':[ToolMessage(content=f"Context: {documents} \n\n Input Question: {question}", name='retrieve_documents', tool_call_id='tool_abcd123')]}

`

and my workflow looks like this.

`

# Add a node for the first tool call
def first_tool_call(state: State) -> dict[str, list[AIMessage]]:
    return {
        "messages": [
            AIMessage(
                content="",
                tool_calls=[
                    {
                        "name": "retrieve_documents",
                        "args": {},
                        "id": "tool_abcd123",
                    }
                ],
            )
        ]
    }

Define a new graph

workflow = StateGraph(State)

workflow.add_node("first_tool_call", first_tool_call)

# Add nodes for the first two tools
workflow.add_node(
    "list_tables_tool", create_tool_node_with_fallback([list_tables_tool])
)
workflow.add_node("get_schema_tool", create_tool_node_with_fallback([get_schema_tool]))

# Add a node for a model to choose the relevant tables based on the question and available tables
model_get_schema = ChatOpenAI(model="gpt-4o", temperature=0).bind_tools(
    [get_schema_tool]
)
workflow.add_node(
    "model_get_schema",
    lambda state: {
        "messages": [model_get_schema.invoke(state["messages"])],
    },
)

# Add a node for a model to generate a query based on the question and schema
query_gen_system = """You are a SQL expert with a strong attention to detail.

Given an input question, output a syntactically correct SQLite query to run, then look at the results of the query and return the answer.

DO NOT call any tool besides SubmitFinalAnswer to submit the final answer.

When generating the query:

Output the SQL query that answers the input question without a tool call.

Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 5 results.
You can order the results by a relevant column to return the most interesting examples in the database.
Never query for all the columns from a specific table, only ask for the relevant columns given the question.

If you get an error while executing a query, rewrite the query and try again.

If you get an empty result set, you should try to rewrite the query to get a non-empty result set. 
NEVER make stuff up if you don't have enough information to answer the query... just say you don't have enough information.

If you have enough information to answer the input question, simply invoke the appropriate tool to submit the final answer to the user.

DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database."""
query_gen_prompt = ChatPromptTemplate.from_messages(
    [("system", query_gen_system), ("placeholder", "{messages}")]
)
query_gen = query_gen_prompt | ChatOpenAI(model="gpt-4o", temperature=0).bind_tools(
    [SubmitFinalAnswer]
)


workflow.add_node("retrieve_documents", retrieve_documents)

workflow.add_node("query_gen", query_gen_node)

# Add a node for the model to check the query before executing it
workflow.add_node("correct_query", model_check_query)

# Add node for executing the query
workflow.add_node("execute_query", create_tool_node_with_fallback([db_query_tool]))


# Specify the edges between the nodes

# workflow.add_edge(START, "retrieve_documents")
# workflow.add_edge('retrieve_documents', "first_tool_call")

workflow.add_edge(START, "first_tool_call")
workflow.add_edge("first_tool_call", 'retrieve_documents')

workflow.add_edge("retrieve_documents", "list_tables_tool")
workflow.add_edge("list_tables_tool", "model_get_schema")
workflow.add_edge("model_get_schema", "get_schema_tool")
workflow.add_edge("get_schema_tool", "query_gen")
workflow.add_conditional_edges(
    "query_gen",
    should_continue,
)
workflow.add_edge("correct_query", "execute_query")
workflow.add_edge("execute_query", "query_gen")

# Compile the workflow into a runnable
app = workflow.compile()`



Is this even right? Requesting some guidance on passing info from one custom tool to another...