09_llamaindex_sql.py

from llama_index.core.settings import Settings
from llama_index.readers.database import DatabaseReader
from llama_index.core.utilities.sql_wrapper import SQLDatabase
from llama_index.core.storage.docstore import SimpleDocumentStore
from llama_index.vector_stores.faiss.base import FaissVectorStore
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.core import VectorStoreIndex
import os
import faiss
from sqlalchemy.types import UserDefinedType
from sqlalchemy.dialects.postgresql.base import ischema_names


class PGVector(UserDefinedType):
    def get_col_spec(self):
        return "vector"

    def bind_processor(self, dialect):
        return lambda value: value

    def result_processor(self, dialect, coltype):
        return lambda value: value


ischema_names["vector"] = PGVector

os.environ["OPENAI_API_KEY"] = "xxx"

supabase_uri = "postgresql://postgres:postgres@127.0.0.1:54322/postgres"
db = SQLDatabase.from_uri(supabase_uri)

query = "SELECT * FROM chat_history"
sql_reader = DatabaseReader(sql_database=db)
docs = sql_reader.load_data(query=query)

embed_model = OpenAIEmbedding(model="text-embedding-ada-002")

dimension = 1536
faiss_index = faiss.IndexFlatL2(dimension)
vector_store = FaissVectorStore(faiss_index=faiss_index)
doc_store = SimpleDocumentStore()

index = VectorStoreIndex.from_documents(
    docs, embed_model=embed_model, vector_store=vector_store
)
index.storage_context.persist(persist_dir="./storage")

llm = OpenAI(model="gpt-4o", temperature=0.7)
Settings.llm = llm

query_engine = index.as_query_engine()

while True:
    user_query = input("You: ")
    if user_query.lower() == "exit":
        break
    response = query_engine.query(user_query)
    print(f"\nResponse: {response}\n")