engine = create_engine(your_tidb_url)
class Entity(Base):
__tablename__ = "entity"
id = Column(Integer, primary_key=True)
content = Column(Text)
content_vec = Column(
VectorType(dim=dim_of_embedding_model),
comment="hnsw(distance=l2)"
)
Base.metadata.create_all(engine)
content = 'TiDB is an open-source distributed SQL database'
open_ai_client = openai.OpenAI(api_key=open_ai_api_key)
embedding = open_ai_client.embeddings.create(
input=[content], model=model_name).data[0].embedding
with Session(engine) as session:
session.add(Entity(content = content, content_vec = embedding))
session.commit()
query = 'What is TiDB?'
embedding_query = open_ai_client.embeddings.create(
input=[query], model=model_name).data[0].embedding
with Session(engine) as session:
entity = session.query(Entity).order_by(
Entity.content_vec.cosine_distance(embedding_query)
).limit(1).first()
print(entity.content)