Quick tip: Using SingleStoreDB with LlamaIndex
Tags: singlestoredb, llamaindex, langchain, vectordatabase
Published: 2023-07-08 00:30:15

## Abstract
In a [previous article](https://dev.to/singlestore/quick-tip-singlestoredb-integration-with-langchain-51po), we saw an example of the integration of SingleStoreDB with LangChain. Another emerging framework is LlamaIndex. In this short article, we'll take an example from the LlamaIndex documentation and apply it to SingleStoreDB.

## Introduction
In this short article, we'll test the ability of SingleStoreDB to work with LlamaIndex using an [example](https://gpt-index.readthedocs.io/en/latest/examples/index_structs/struct_indices/SQLIndexDemo.html) from the documentation. This should be very straightforward. In later articles, we'll see if we can achieve tighter integration.

As described in a [previous article](https://thenewstack.io/building-gpt-applications-on-open-source-langchain-part-2/), we'll follow the instructions to create a SingleStoreDB Cloud account, Workspace Group, Workspace, and Notebook.

## Create the Database
From the left nav in our SingleStoreDB Cloud account, we'll select the **SQL Editor** and run the following commands:

```sql
CREATE DATABASE IF NOT EXISTS demo_db;

USE demo_db;

DROP TABLE IF EXISTS city_stats;
```

## Fill out the Notebook
We'll follow the [example](https://gpt-index.readthedocs.io/en/latest/examples/index_structs/struct_indices/SQLIndexDemo.html) described in the documentation and modify it, where required, to work with SingleStoreDB.

First, we'll suppress some warnings which result from SQLAlchemy as we'll use the version installed by LlamaIndex rather than the version installed by SingleStore:

```python
import warnings

warnings.filterwarnings("ignore")
```

Next, we'll install a recent version of LlamaIndex:

```shell
!pip uninstall sqlalchemy-singlestoredb -y --quiet
!pip install llama-index==0.6.35 --quiet
```

We'll now enter and save our `OpenAI API Key`:

```python
import os
import getpass

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
```

We'll keep the logging-level commented out:

```python
#import logging
#import sys

#logging.basicConfig(stream=sys.stdout, level=logging.INFO)
#logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
```

And directly import from `IPython.display`:

```python
from IPython.display import Markdown, display
```

## Create Database Schema
We'll get the appropriate imports from SQLAlchemy:

```python
from sqlalchemy import (
    create_engine,
    MetaData,
    Table,
    Column,
    String,
    Integer,
    select,
    column,
)
```

Next, we'll create the connection:

```python
engine = create_engine("mysql+pymysql://admin:<password>@<host>:3306/demo_db")
metadata_obj = MetaData()
```

We'll replace the `<password>` and `<host>` with the values from our SingleStoreDB Cloud account.

Now we'll create a table:

```python
# create city SQL table
table_name = "city_stats"
city_stats_table = Table(
    table_name,
    metadata_obj,
    Column("city_name", String(16), primary_key=True),
    Column("population", Integer),
    Column("country", String(16), nullable=False),
)
metadata_obj.create_all(engine)
```

## Define SQL Database
We'll continue with the LlamaIndex documentation:

```python
from llama_index import SQLDatabase, ServiceContext
from langchain import OpenAI
from llama_index import LLMPredictor
```

```python
llm = OpenAI(temperature=0, model="text-davinci-002")
service_context = ServiceContext.from_defaults(llm=llm)
```

```python
sql_database = SQLDatabase(engine, include_tables=["city_stats"])
```

```python
sql_database.table_info
```

The result should be as follows:

```
'\nCREATE TABLE city_stats (\n\tcity_name VARCHAR(16) CHARACTER SET utf8 COLLATE utf8_general_ci NOT NULL, \n\tpopulation INTEGER(11), \n\tcountry VARCHAR(16) CHARACTER SET utf8 COLLATE utf8_general_ci NOT NULL\n)\n\n/*\n3 rows from city_stats table:\ncity_name\tpopulation\tcountry\n\n*/'
```

Next, we'll continue with the code:

```python
sql_database = SQLDatabase(engine, include_tables=["city_stats"])
from sqlalchemy import insert

rows = [
    {"city_name": "Toronto", "population": 2930000, "country": "Canada"},
    {"city_name": "Tokyo", "population": 13960000, "country": "Japan"},
    {"city_name": "Chicago", "population": 2679000, "country": "United States"},
    {"city_name": "Seoul", "population": 9776000, "country": "South Korea"},
]
for row in rows:
    stmt = insert(city_stats_table).values(**row)
    with engine.connect() as connection:
        cursor = connection.execute(stmt)
        connection.commit()
```

```python
# view current table
stmt = select(city_stats_table.c["city_name", "population", "country"]).select_from(
    city_stats_table
)

with engine.connect() as connection:
    results = connection.execute(stmt).fetchall()
    print(results)
```

The result should be as follows:

```
[('Seoul', 9776000, 'South Korea'), ('Toronto', 2930000, 'Canada'), ('Tokyo', 13960000, 'Japan'), ('Chicago', 2679000, 'United States')]
```

## Query Index
We'll now run the following query:

```python
from sqlalchemy import text

with engine.connect() as con:
    rows = con.execute(text("SELECT city_name from city_stats"))
    for row in rows:
        print(row)
```

The result should be as follows:

```
('Toronto',)
('Tokyo',)
('Chicago',)
('Seoul',)
```

## Natural language SQL
We can also prepare a natural language query:

```python
from llama_index.indices.struct_store.sql_query import NLSQLTableQueryEngine

query_engine = NLSQLTableQueryEngine(
    sql_database=sql_database,
    tables=["city_stats"],
)
query_str = (
    "Which city has the highest population?"
)
response = query_engine.query(query_str)
```

## Building our Table Index

```python
import openai

openai.api_key = os.environ["OPENAI_API_KEY"]
```

```python
from llama_index.indices.struct_store.sql_query import SQLTableRetrieverQueryEngine
from llama_index.objects import SQLTableNodeMapping, ObjectIndex, SQLTableSchema
from llama_index import VectorStoreIndex

# set Logging to DEBUG for more detailed outputs
table_node_mapping = SQLTableNodeMapping(sql_database)
table_schema_objs = [(SQLTableSchema(table_name="city_stats"))] # add a SQLTableSchema for each table

obj_index = ObjectIndex.from_objects(
    table_schema_objs,
    table_node_mapping,
    VectorStoreIndex,
)
query_engine = SQLTableRetrieverQueryEngine(
    sql_database, obj_index.as_retriever(similarity_top_k=1)
)
```

Now we can use the natural language query:

```python
response = query_engine.query("Which city has the highest population?")
display(Markdown(f"<b>{response}</b>"))
```

The result should be similar to the following:

```
Tokyo has the highest population with 13,960,000 people.
```

The raw output can also be obtained using the following:

```python
# you can also fetch the raw result from SQLAlchemy! 
response.metadata["result"]
```

The result should be as follows:

```
[('Tokyo', 13960000)]
```

Additional context info can be added as follows:

```python
# manually set context text
city_stats_text = (
    "This table gives information regarding the population and country of a given city.\n"
    "The user will query with codewords, where 'foo' corresponds to population and 'bar'"
    "corresponds to city."
)

table_node_mapping = SQLTableNodeMapping(sql_database)
table_schema_objs = [(SQLTableSchema(table_name="city_stats", context_str=city_stats_text))]
```

### Using LangChain for Querying
LangChain can also be used:

```python
from langchain import OpenAI, SQLDatabase, SQLDatabaseChain
```

```python
llm = OpenAI(temperature=0)
```

```python
# set Logging to DEBUG for more detailed outputs
db_chain = SQLDatabaseChain(llm=llm, database=sql_database)
```

```python
db_chain.run("Which city has the highest population?")
```

The result should be similar to the following:

```
'Tokyo has the highest population with 13960000 people.'
```

## Summary
In this quick example, we have seen how to use LlamaIndex with SingleStoreDB. Only a few minor code modifications were required to the example from the LlamaIndex documentation.

## License
The MIT License

Copyright © Jerry Liu

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.