AI Engineer - Embedding Generation from Germini API and Storing in PostgreSQL

In this post, I first demonstrate that sample Python programs have been developed and tested to do the following two tasks: 

1. Create embeddings from sample source data using the Germini API 

2. Use PostgreSQL as a vector database and store embeddings data in it using pgvector. 

Then I demonstrate similarity search in PostgreSQL with pgvector. 

1. Source data preparation: 

 From the DBAETS application `events` table in the PostgreSQL, dumping ` event_id` and `title` data into a csv file. Total 2021 records.

 \copy ( select event_id, title from events order by 1 ) to 'events_title.csv' with csv 

 The `events` table stores data about database administration events, each entry represents a certain kind of DBA task or activity. The `title` of the `event` table is trivial, but the purpose of using it as the source of content is to make me familiar with the embedding creation and storing process. The end goal is to use the `description` column of the `events` table, which contains a detailed description of the event, therefore providing the useful domain knowledge for DBAs.

 2. Embedding generation 

The python program emb_to_csv.py is developed to generate embedding from source csv file and generate a csv file including embeddings of `title`. Gemini API is used with model: ‘text-embedding-004’, which generates vectors with 768 dimensions .

==============  screenshot run the program ==============
(.venv) C:\Users\dsun\Denis_files\python_proj\ai_apps\dbaets_pgvector>python emb_to_csv_2.py
   event_id                                              title
0         3  PREP CALL for the 2pm actual call to discuss r...
1         4                   Query Certification / 04-11-2012
2         5        [DemoApp] [Migration] -  RD175935 - Tonight
3         6  RE: myhostpd3:testprod3 ALERT_LOG Error: [SR 3-...
4         7                    space add in mrosple/mxdsacdd02
total batches : 22 batch size: 100 total records: 2120
process batch no: 0
loop: 0 event_id: 3
loop: 10 event_id: 13
loop: 20 event_id: 24
loop: 30 event_id: 34
loop: 40 event_id: 44
loop: 50 event_id: 54
loop: 60 event_id: 64
loop: 70 event_id: 74
loop: 80 event_id: 84
loop: 90 event_id: 94
process batch no: 1
loop: 100 event_id: 104
loop: 110 event_id: 114
loop: 120 event_id: 124
loop: 130 event_id: 134
loop: 140 event_id: 144
loop: 150 event_id: 155
loop: 160 event_id: 165
loop: 170 event_id: 175
loop: 180 event_id: 185
loop: 190 event_id: 195
process batch no: 2
loop: 200 event_id: 205
…
(ommitting ..)

I intentionally used some small batch size and sleep 200 seconds between batches to avoid Gemini API free tier rate limit, if I don't do that, I could end up with errors "RATE_LIMIT_EXCEEDED". 

3. Embedding loacding 

The python program emb_csv_to_pgvector.py is developed to load the embeddings from csv to a PostgreSQL table. 

4. Similarity search using SQL 

 The following query calculates vector distances using the <#> negative inner product operator to compare event titles against a target event (event_id = 2690). It retrieves the target event's self-similarity score alongside the top 10 most similar distinct events, returning their event IDs, titles, and similarity scores in a single combined result.

etsdb=> -- top event titles that are mostly similar to first title
etsdb=> (
etsdb(> select event_id, title,  embedding <#> embedding as similarity
etsdb(> from event_title_embeddings
etsdb(> where event_id=2690
etsdb(> )
etsdb-> union all
etsdb-> (
etsdb(> select t.event_id , t.title , (t.embedding <#> i.embedding ) as similarity
etsdb(> from event_title_embeddings t join event_title_embeddings i
etsdb(> on t.event_id != i.event_id
etsdb(> and i.event_id=2690
etsdb(> order by similarity
etsdb(> limit 10
etsdb(> )
etsdb-> ;
 event_id |                              title                               |     similarity
----------+------------------------------------------------------------------+---------------------
     2690 | Master note for operation support - mydb2                        | -0.9999982118606567
     2431 | Master note for operation support - mydb1 BVC                    | -0.8375328183174133
     2629 | Master note for operation support - omrtppdb OEM repostory       | -0.8193813562393188
     2444 | Master note for operation support - iexdb                        | -0.7917341589927673
     2450 | mydb4 database support master notes                              | -0.7660208344459534
     2459 | master note for operation support - mydb12sc APP1                | -0.766010582447052
     2876 | Operations Master notes for EV6V- mmypos on dbhost41/42/43       | -0.6405361294746399
     2953 | Master Notes -APP2                                               | -0.6208126544952393
     2937 | SCM master notes                                                 | -0.6061583161354065
     2275 | mydb1 database related doc                                       | -0.6051017045974731
     2879 | Master Notes - APP1 - mydb3                                      | -0.6019692420959473
(11 rows)

In summary, I successfully established an end-to-end embedding pipeline using the Gemini text-embedding-004 model and PostgreSQL with pgvector. By generating 768-dimensional vector embeddings from DBA event titles and querying them using inner product distance (<#>), the setup demonstrated accurate semantic similarity retrieval across database task records. This foundational work sets the stage for scaling up to full event descriptions and implementing a robust domain-specific vector search system.