Retriever Tool for Vertex AI Model 
Labels: name: area: agent, name: area: embeddings, name: area: vector store, name: auto:bug, name: area: models

### Issue you'd like to raise.

I am trying to add retriever tool to SQL Database Toolkit and run SQL Agent with Vertex AI Text Bison LLM, however I am facing error.
I am following below documentation for implementation
[https://python.langchain.com/docs/use_cases/qa_structured/sql?ref=blog.langchain.dev#case-3-sql-agents](url)

My Python Code
`llm = VertexAI(
    model_name='text-bison@001',
    max_output_tokens=256,
    temperature=0,
    top_p=0.8,
    top_k=40,
    verbose=True,
)

few_shots = {'a': sql a query;',
              "b '.": "sql b query",
              "c.": "sql c query;",
              "d": "sql d query; "
          }

from langchain.embeddings import VertexAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.schema import Document

embeddings = VertexAIEmbeddings()

few_shot_docs = [Document(page_content=question, metadata={'sql_query': few_shots[question]}) for question in few_shots.keys()]
vector_db = FAISS.from_documents(few_shot_docs, embeddings)
#retriever = vector_db.as_retriever(search_type="similarity_score_threshold", search_kwargs={"score_threshold": .5})
retriever = vector_db.as_retriever()

from langchain.agents import initialize_agent, Tool
from langchain.agents import load_tools
from langchain.agents import AgentType
from langchain.tools import BaseTool

from langchain.agents.agent_toolkits import create_retriever_tool

tool_description = """
This tool will help you understand similar examples to adapt them to the user question.
Input to this tool should be the user question.
"""

retriever_tool = create_retriever_tool(
        retriever,
        name='sql_get_similar_examples',
        description=tool_description
    )
custom_tool_list = [retriever_tool]

from langchain.agents import create_sql_agent
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.sql_database import SQLDatabase
from langchain.agents import AgentExecutor
from langchain.agents.agent_types import AgentType

from sqlalchemy import *
from sqlalchemy.engine import create_engine
from sqlalchemy.schema import *
import pandas as pd

table_uri = f"bigquery://{project}/{dataset_id}"
engine = create_engine(f"bigquery://{project}/{dataset_id}")
db = SQLDatabase(engine=engine,metadata=MetaData(bind=engine),include_tables=table_name,view_support=True)

toolkit = SQLDatabaseToolkit(db=db, llm=llm)

custom_suffix = """Begin!
Question: {input}
Thought: I should first get the similar examples I know.
If the examples are enough to construct the query, I can build it.
Otherwise, I can then look at the tables in the database to see what I can query.
Then I should query the schema of the most relevant tables.
{agent_scratchpad}"""

agent = create_sql_agent(llm=llm,
                         toolkit=toolkit,
                         verbose=True,
                         agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
                         extra_tools=custom_tool_list,
                         suffix=custom_suffix
                        )

langchain.debug = True

Thought:The table does not have the column 'alloted_terminal'.
Action: sql_get_similar_examples
Action Input: {Question}
Observation: {Document(fewshots}
Thought:I don't have enough information to construct the query.
Action: sql_get_similar_examples
Action Input: {Question}
Observation: {Document(fewshots}
Thought:I don't have enough information to construct the query.
 Action: sql_get_similar_examples
Action Input: {Question}
Observation: {Document(fewshots}
Thought:I don't have enough information to construct the query.
Action: sql_get_similar_examples
Action Input: {Question}
Observation: {Document(fewshots}
Thought:I don't have enough information to construct the query.
Action: sql_get_similar_examples
Action Input: {Question}
Observation: {Document(fewshots}
Thought:I don't have enough information to construct the query.
Action: sql_get_similar_examples
Action Input: {Question}
Observation: {Document(fewshots}
Thought:I don't have enough information to construct the query.
Action: sql_get_similar_examples
Action Input: {Question}
Observation: {Document(fewshots}
Thought:I don't have enough information to construct the query.
Action: sql_get_similar_examples
Action Input: {Question}
Observation: {Document(fewshots}
Thought:I don't have enough information to construct the query.
Action: sql_get_similar_examples
Action Input: {Question}
Observation: {Document(fewshots}
Thought:I don't have enough information to construct the query.
Action: sql_get_similar_examples
Action Input: {Question}
Observation: {Document(fewshots}
Thought:I don't have enough information to construct the query.
Action: sql_get_similar_examples
Action Input: {Question}
Observation: {Document(fewshots}
Thought:I don't have enough information to construct the query.
Action: sql_get_similar_examples
Action Input: {Question}
Observation: {Document(fewshots}
Thought:

> Finished chain.
Agent stopped due to iteration limit or time limit.

{
  "output": "Agent stopped due to iteration limit or time limit."
}
'
Agent stopped due to iteration limit or time limit.

### Suggestion:

_No response_