Superman37891/PostgreSQL_Documentation_AI_Assistant

PostgreSQL Documentation AI Assistant

Summary

This is a Retrieval-Augmented Generation (RAG) system that utilizes the official PostgreSQL-18 documentation to answer user queries related to PostgreSQL 18

Process

In my Setup notebook, I extracted the text as well as important metadata using the pdfminer library. I then analyzed the metadata and documentation to classify text (ex: header, suheader, body text, query, etc.) and then group lines into blocks that I would chunk.

I then chose a custom chunking strategy of keeping complete paragraphs intact while keeping chunks at or below a certain number of characters whenever possible with a decent size overlap so that each chunk could be meaningful.

After grouping the text into chunks with relevant metadata (ex: header, subheader, page number, etc.), I embedded those chunks using Sentence Transformers before uploading them to a free-tier NeonDB PostgreSQL Database for storage.

My RAG system then used cosine similarity to retrieve relevant chunks in text form along with metadata from the database, and I used these chunks as context in the prompt given to the LLM.

I used RRF to create a combined score for vector/keyword similarity so that chunks that were relevant in both meaning and keywords got more priority than chunks that were relevant in only one category.

Model Selection

Model selection was very difficult because I had to find a model that followed the instructions and gave satisfactory answers, was not too expensive, and hardest of all, was compatible with HuggingFace's Inference Client API.

Gemma-2-2b-it does not attempt to reasonably answer query 8 even when the prompt says to use the general knowledge if the provided context can't give a reasonable answer. Gemma-2 also sometimes treats query 3's start_date as a separate column rather than a nested calculation

Gemma-7b-it took way too long to answer simple queries, such as 45 seconds to answer the relatively simple query 1

I wanted to settle on MistralAI 7B due to the quality and instruction-following of its responses but I noted that this model would not be compatible with my HuggingFace application.

I also liked Qwen2.5-7B-Instruct due to the quality of its responses and it also following instructions excellently, but it wasn't supported as an InferenceClient on HuggingFace

I then tried Qwen2.5-Coder-32B-Instruct but it was way too expensive, despite achieving significantly lower LLM latencies than my other models.

My final model I settled on was meta-llama/Llama-3.1-8B-Instruct because it was compatible with HuggingFace, I was satisfied with the responses, and the cost was not too high

Other choices

I chose GLOBALTOPK=5 because, when I tried higher values, the last retrieved chunks were often marginally or not relevant to the question being asked.

I chose to use cosine similarity instead of Euclidean distance because this was a text-based application where I did not want the magnitude of the vectors to have a big impact

Improvements

To improve this assistant further, I could focus on utilizing outside documentation containing example queries pertaining to specific use cases to help train this assistant.

I considered using objective metrics like recall\@k or precision@\k, but I realized that this would require extensive labor in manually sifting through and identifying documentation / chunks for each query, which I feel is too much for this project due to having over 15,000 chunks and the documentation being over 3300 pages. I also considered using another, more well-known LLM like ChatGPT or Gemini to evaluate my current model, but I noted that these outside LLMs could hallucinate or be biased by the context in my prompts and don't have a ground truth to refer to.

I also considered utilizing chat history to guide my assistant, but I noted that in order to do so, I would likely have to store the chat history somewhere and I did not want any free tier storage limit to be reached in case someone misused the assistant. In a production application, I would save the chat history to a database as well AND have safeguards to prevent one user's chat history from taking too much storage.

One big source of latency was the GPU acquiring process on HuggingFace's ZeroGPU. In order to combat this in a production application, I could subscribe to HuggingFace Pro to use CPU hardware to bypass this GPU process altogether, as my code runs just fine on CPU.