SQL Agent - Langgraph - python

[MatrixChess]

Body

Hello everyone,
I'm follwing this guide trying to adapt the code to use BigQuery and Gemini. There is a step that generate an error, The following code is an executable excerpt to reproduce the error:

from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import ChatPromptTemplate
import os
import getpass
from langchain_core.tools import tool

os.environ['GOOGLE_API_KEY'] = getpass.getpass('Gemini API Key:')

@tool
def db_query_tool(query: str) -> str:
    """
    Execute a SQL query against the database and get back the result.
    If the query is not correct, an error message will be returned.
    If an error is returned, rewrite the query, check the query, and try again.
    """
    result = query
    if not result:
        return "Error: Query failed. Please rewrite your query and try again."
    return result


query_check_system = """You are a SQL expert with a strong attention to detail.
Double check the BigQuery query for common mistakes, if there are any , rewrite the query. If there are no mistakes, just reproduce the original query.
You will call the appropriate tool to execute the query after running this check."""

query_check_prompt = ChatPromptTemplate.from_messages(
    [("system", query_check_system), ("placeholder", "{messages}")]
)

query_check = query_check_prompt | ChatGoogleGenerativeAI(temperature=0, model="gemini-pro", convert_system_message_to_human=True).bind_tools(
    [db_query_tool], tool_choice="required"
)

query_check.invoke({"messages": [("user", "SELECT COUNT(id_cliente) FROM my_project.my_dataset.cliente;")]})

This code produces the following error:

Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/proto/message.py", line 729, in __init__
    pb_value = marshal.to_proto(pb_type, value)
  File "/usr/local/lib/python3.10/dist-packages/proto/marshal/marshal.py", line 235, in to_proto
    pb_value = self.get_rule(proto_type=proto_type).to_proto(value)
  File "/usr/local/lib/python3.10/dist-packages/proto/marshal/rules/message.py", line 36, in to_proto
    return self._descriptor(**value)
ValueError: unknown enum label "any"

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/user/sqlchain/app_langgraph_base.py", line 112, in <module>
    query_check.invoke({"messages": [("user", "SELECT COUNT(id_cliente) FROM my_project.my_dataset.cliente;")]})
  File "/home/user/.local/lib/python3.10/site-packages/langchain_core/runnables/base.py", line 2507, in invoke
    input = step.invoke(input, config)
  File "/home/user/.local/lib/python3.10/site-packages/langchain_core/runnables/base.py", line 4588, in invoke
    return self.bound.invoke(
  File "/home/user/.local/lib/python3.10/site-packages/langchain_core/language_models/chat_models.py", line 248, in invoke
    self.generate_prompt(
  File "/home/user/.local/lib/python3.10/site-packages/langchain_core/language_models/chat_models.py", line 677, in generate_prompt
    return self.generate(prompt_messages, stop=stop, callbacks=callbacks, **kwargs)
  File "/home/user/.local/lib/python3.10/site-packages/langchain_core/language_models/chat_models.py", line 534, in generate
    raise e
  File "/home/user/.local/lib/python3.10/site-packages/langchain_core/language_models/chat_models.py", line 524, in generate
    self._generate_with_cache(
  File "/home/user/.local/lib/python3.10/site-packages/langchain_core/language_models/chat_models.py", line 749, in _generate_with_cache
    result = self._generate(
  File "/home/user/.local/lib/python3.10/site-packages/langchain_google_genai/chat_models.py", line 758, in _generate
    request = self._prepare_request(
  File "/home/user/.local/lib/python3.10/site-packages/langchain_google_genai/chat_models.py", line 907, in _prepare_request
    formatted_tool_config = ToolConfig(
  File "/usr/local/lib/python3.10/dist-packages/proto/message.py", line 757, in __init__
    pb_value = marshal.to_proto(pb_type, value)
  File "/usr/local/lib/python3.10/dist-packages/proto/marshal/marshal.py", line 235, in to_proto
    pb_value = self.get_rule(proto_type=proto_type).to_proto(value)
  File "/usr/local/lib/python3.10/dist-packages/proto/marshal/rules/message.py", line 36, in to_proto
    return self._descriptor(**value)
ValueError: unknown enum label "any"

Any ideas on what I'm doing wrong?

Thanks in advance for the help,
Alessio

Guidelines

• I have read and understood this category's guidelines before making this post.

[ebndev]

Hi @CodeCandyCelvin,

We’ve clarified our stance on using generative AI tools like ChatGPT within our Community via this announcement. Please review the guidelines to ensure your post meets them as failure to adhere to those rules can result in action taken by our moderator team.  You can read our updated Code of Conduct and the announcement for more details. Thank you for helping us maintain an authentic and beneficial space for everyone.

[MatrixChess]

Hi CodeCandyCelvin and thanks for the answer.
I'm not really sure of what the modifications are, the code you posted seems to be exactly the same, Can you explain what I have to try?

Thanks,
Alessio

[CodeCandyCelvin]

Hi Alessio,

Thanks for your patience. Here's a revised version of your code. I've removed the tool_choice parameter to avoid the unknown enum label "any" error. Additionally, I included proper BigQuery query execution and error handling within the db_query_tool function.

Please replace the placeholder in db_query_tool with actual BigQuery execution logic and ensure that the API key provided has the necessary permissions to access BigQuery and execute queries.
`from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import ChatPromptTemplate
import os
import getpass
from langchain_core.tools import tool
from google.cloud import bigquery
from google.api_core.exceptions import GoogleAPIError

Set the Gemini API key using environment variable

os.environ['GOOGLE_API_KEY'] = getpass.getpass('Gemini API Key:')

Initialize the BigQuery client

client = bigquery.Client()

Define the tool for executing SQL queries

@tool
def db_query_tool(query: str) -> str:
"""
Execute a SQL query against the BigQuery database and get back the result.
If the query is not correct, an error message will be returned.
If an error is returned, rewrite the query, check the query, and try again.
"""
try:
query_job = client.query(query) # Make an API request.
results = query_job.result() # Wait for the job to complete.
return str([row for row in results])
except GoogleAPIError as e:
return f"Error: Query failed. Please rewrite your query and try again. Details: {str(e)}"

Define the prompt for the query check system

query_check_system = """You are a SQL expert with a strong attention to detail.
Double check the BigQuery query for common mistakes, if there are any, rewrite the query. If there are no mistakes, just reproduce the original query.
You will call the appropriate tool to execute the query after running this check."""

Create the prompt template

query_check_prompt = ChatPromptTemplate.from_messages(
[("system", query_check_system), ("placeholder", "{messages}")]
)

Configure the query check using the ChatGoogleGenerativeAI model

query_check = query_check_prompt | ChatGoogleGenerativeAI(
temperature=0, model="gemini-pro", convert_system_message_to_human=True
).bind_tools(
[db_query_tool]
)

Invoke the query check with a sample query

response = query_check.invoke({"messages": [("user", "SELECT COUNT(id_cliente) FROM my_project.my_dataset.cliente;")]})

Print the response

print(response)
`
Let me know if you have any questions or need further assistance!
Celvin
This should help resolve the issue by ensuring the code correctly handles the BigQuery query execution and avoids the enum label error.

[MatrixChess]

Thanks Celvin, removing tool_choice="required" fixed the problem for now, I will try to understand what was exactly its purpose.
I obmitted the BigQuery connection from the question because it was not part of the issue, but thanks for proposing that also :)

Alessio

[ebndev]

@MatrixChess thanks for being a part of the GitHub Community! Glad you were able to resolve your issue. I will go ahead and close this discussion 😄