Superman37891/PostgreSQL_Documentation_AI_Assistant
0
1import spaces2import os3import torch4import gradio as gr5from huggingface_hub import InferenceClient6 7import time # For measuring latency8import traceback # For debugging9from sentence_transformers import SentenceTransformer10import psycopg211from pgvector.psycopg2 import register_vector12from pgvector import Vector13 14embedding_model = None15 16'''17def satisfies_zerogpu():18 _ = torch.tensor([1.0]).cuda()19 return True20'''21 22def get_embedding_model():23 global embedding_model24 if embedding_model is None:25 embedding_model = SentenceTransformer("BAAI/bge-small-en-v1.5", device="cpu")26 return embedding_model27 28model_name = "meta-llama/Llama-3.1-8B-Instruct"29 30HF_TOKEN = os.environ.get("HF_READ_TOKEN", "")31 32llm_client = InferenceClient(33 model=model_name,34 token=HF_TOKEN35)36 37embedding_cache = {}38 39@spaces.GPU(duration=7)40def get_embedding(text):41 try:42 if text in embedding_cache:43 return embedding_cache[text]44 model = get_embedding_model()45 model.to("cuda")46 embedding = model.encode(text)47 embedding_list = embedding.tolist() if hasattr(embedding, 'tolist') else list(embedding)48 embedding_cache[text] = embedding_list49 return embedding_list50 except Exception as e:51 print(f"Error getting embedding for text: {text[:50]}...")52 traceback.print_exc()53 return []54 55GLOBAL_TOP_K=556GLOBAL_SIMILARITY_THRESHOLD=0.7057 58 59# --- Database Connection Details ---60PG_HOST = os.environ.get("PGSQL_AI_CHATBOT_NEONDB_HOST")61PG_PORT = 543262PG_USER = 'neondb_owner'63PG_PASSWORD = os.environ.get("PGSQL_AI_CHATBOT_NEONDB_PASSWORD")64PG_DBNAME = 'neondb'65 66def get_db_connection():67 conn = psycopg2.connect(68 host=PG_HOST,69 port=PG_PORT,70 user=PG_USER,71 password=PG_PASSWORD,72 dbname=PG_DBNAME,73 sslmode='require' # Add this line to enforce SSL74 )75 register_vector(conn) # Register the pgvector type with psycopg276 return conn77 78def format_chunk(row):79 return {80 "chunk_id": row[0],81 "text": row[1],82 "page_numbers": row[2],83 "source_header": row[3],84 "source_subheader": row[4],85 "block_type": row[5],86 "similarity": row[6]87 }88 89def format_keyword_chunk(row):90 return {91 'chunk_id': row[0],92 'text': row[1],93 'page_numbers': row[2],94 'source_header': row[3],95 'source_subheader': row[4],96 'block_type': row[5],97 'rank': row[6]98 }99 100def vector_search(query_embedding, top_k=GLOBAL_TOP_K, similarity_threshold=GLOBAL_SIMILARITY_THRESHOLD, filter_block_types=None):101 if not query_embedding:102 print("vector_search skipped: Empty query embedding provided.")103 return []104 conn = None105 try:106 conn = get_db_connection()107 cur = conn.cursor()108 109 # Build the WHERE clause for filtering by block_type110 where_clause = ""111 if filter_block_types:112 # Ensure filter_block_types is a tuple or list for the IN clause113 if isinstance(filter_block_types, str):114 filter_block_types = (filter_block_types,)115 placeholders = ', '.join(['%s'] * len(filter_block_types))116 where_clause = f"WHERE block_type IN ({placeholders})"117 118 # Perform a vector similarity search using cosine similarity (1 - cosine distance).119 # The <=> operator computes cosine distance. 1 - (cosine distance) gives cosine similarity.120 # ORDER BY DESC for closest matches (highest similarity).121 # The LIMIT k clause restricts the number of results.122 query_sql = f"""123 SELECT124 chunk_id,125 text,126 page_numbers,127 source_header,128 source_subheader,129 block_type,130 1 - (embedding <=> %s) AS similarity131 FROM132 document_chunks133 {where_clause}134 ORDER BY135 similarity DESC136 LIMIT %s;137 """138 139 # Prepare parameters for the query140 # Explicitly cast the query_embedding to a Vector object141 142 params = [Vector(query_embedding)]143 144 if filter_block_types:145 params.extend(filter_block_types)146 params.append(top_k)147 148 cur.execute(query_sql, params)149 results = cur.fetchall()150 151 # Convert results to a more readable format (list of dictionaries)152 search_results = []153 for row in results:154 if row[6] >= similarity_threshold:155 search_results.append(format_chunk(row))156 return search_results157 158 except Exception as e:159 print("Error during vector_search")160 traceback.print_exc()161 return []162 finally:163 if conn:164 conn.close()165 166def keyword_search(keyword_query, top_k=GLOBAL_TOP_K):167 if not keyword_query:168 return []169 conn = None170 try:171 conn = get_db_connection()172 cur = conn.cursor()173 # Perform keyword search174 cur.execute("""175 SELECT176 chunk_id,177 text,178 page_numbers,179 source_header,180 source_subheader,181 block_type,182 ts_rank(search_vector, plainto_tsquery('english', %s)) AS rank183 FROM document_chunks184 WHERE search_vector @@ plainto_tsquery('english', %s)185 ORDER BY rank DESC186 LIMIT %s;187 """, (keyword_query, keyword_query, top_k))188 results = cur.fetchall()189 190 keyword_search_results = []191 for row in results:192 keyword_search_results.append(format_keyword_chunk(row))193 return keyword_search_results194 195 except Exception as e:196 print("Error during keyword_search")197 traceback.print_exc()198 return []199 finally:200 if conn:201 conn.close()202 203def get_relevant_pages(relevant_chunks):204 # Collect all page numbers, maintaining relevance order and ensuring uniqueness.205 ordered_unique_pages = []206 seen_pages = set()207 208 for chunk in relevant_chunks:209 if chunk.get('page_numbers'):210 for page_num in chunk['page_numbers']:211 if page_num not in seen_pages:212 ordered_unique_pages.append(page_num)213 seen_pages.add(page_num)214 215 return ordered_unique_pages216 217def setup_prompt(query_text, relevant_chunks):218 try:219 parts = [f"""220 You are an expert on PostgreSQL documentation.221 222 Use the Retrieved Documentation as your primary source of truth.223 If the documentation fully answers the question, answer only from it.224 If the documentation is incomplete, explicitly state what information was missing before using general PostgreSQL knowledge.225 226 Always clearly state the source of your answer at the end, choosing *one* of the following two options:227 1. 'Answer Source: PostgreSQL Documentation.' (If the answer is derived primarily or entirely from the provided documentation.)228 2. 'Answer Source: General PostgreSQL Knowledge.' (If the answer relies significantly on knowledge outside the provided documentation.)229 230 Never invent documentation that was not retrieved.231 When answering, synthesize information from *all* relevant retrieved chunks to provide a comprehensive answer, especially if the query has multiple parts.232 233 """]234 if relevant_chunks:235 parts.append("\nRetrieved Documentation:\n\n")236 for i, chunk in enumerate(relevant_chunks, 1):237 parts.append(238 f"[Chunk {i}]\n"239 f"Header: {chunk['source_header']}\n"240 f"Subheader: {chunk['source_subheader']}\n"241 f"Page Numbers: {chunk['page_numbers']}\n"242 f"Text:\n{chunk['text']}\n\n"243 )244 parts.append(f"\nQuestion: {query_text}\n")245 prompt = "\n".join(parts)246 return prompt247 except Exception as e:248 print("Error in setup_prompt")249 traceback.print_exc()250 return []251 252def rrf_sort_chunks(vector_results, keyword_results, k=60, top_k = GLOBAL_TOP_K):253 rrf_scores = {}254 chunk_map = {}255 256 for rank, chunk in enumerate(vector_results, start=1):257 chunk_id = chunk['chunk_id']258 chunk_map[chunk_id] = chunk259 rrf_scores[chunk_id] = rrf_scores.get(chunk_id, 0.0) + (1.0/(k+rank))260 for rank, chunk in enumerate(keyword_results, start=1):261 chunk_id = chunk['chunk_id']262 263 if chunk_id not in chunk_map:264 chunk_map[chunk_id] = chunk265 else:266 chunk_map[chunk_id]['rank'] = chunk.get('rank')267 rrf_scores[chunk_id] = rrf_scores.get(chunk_id, 0.0) + (1.0/(k+rank))268 269 for chunk_id, score in rrf_scores.items():270 chunk_map[chunk_id][('rrf_score')] = score271 272 sorted_chunks = sorted(273 chunk_map.values(),274 key=lambda item: item['rrf_score'],275 reverse=True276 )277 278 return sorted_chunks[:top_k]279 280def get_answer(query_text, top_k=GLOBAL_TOP_K, use_hybrid_search=True, debug=False):281 # Call the LLM to answer the question from setup_prompt282 try:283 #print("Generating embedding...")284 start_time = time.time()285 query_embedding = get_embedding(query_text)286 query_embedding_end_time = time.time()287 query_embedding_latency = query_embedding_end_time - start_time288 289 relevant_chunks = []290 if use_hybrid_search:291 # Fetch slightly more results than top_k for a better RRF pool292 fetch_k = max(top_k*2, 20)293 vector_results = vector_search(query_embedding, top_k=top_k) if query_embedding else []294 keyword_results = keyword_search(query_text, top_k=top_k) if query_text else []295 296 relevant_chunks = rrf_sort_chunks(vector_results, keyword_results, k=60, top_k=top_k)297 # Combine and de-duplicate results298 elif query_embedding is not None:299 # Fallback to pure vector search if hybrid search is off or keyword query is empty300 relevant_chunks = vector_search(query_embedding, top_k=top_k)301 elif query_text is not None: # Fallback to pure keyword search if no embedding or no hybrid search302 relevant_chunks = keyword_search(query_text, top_k=top_k)303 304 retrieval_end_time = time.time()305 retrieval_latency = retrieval_end_time - query_embedding_end_time306 307 if not relevant_chunks:308 print("No relevant chunks found.")309 310 # Get the complete prompt that will be sent to the LLM311 full_prompt_sent_to_llm = setup_prompt(query_text, relevant_chunks)312 # Pass prompt to an LLM313 try:314 completion = llm_client.chat_completion(315 messages=[316 {317 "role": "user",318 "content": full_prompt_sent_to_llm319 }320 ],321 max_tokens=500,322 temperature=0.2323 )324 llm_answer = completion.choices[0].message.content.strip()325 except Exception as e:326 #print(f"Exception occurred: {e}")327 traceback.print_exc()328 #return "The language model could not generate a response."329 return f"Error in get_answer when calling llm client: {e}"330 331 relevant_page_numbers = get_relevant_pages(relevant_chunks)332 if not relevant_page_numbers:333 llm_answer += "\nNo relevant pages found."334 else:335 llm_answer += f'\n\nRelevant page numbers of PostgreSQL 18 Official Documentation: {relevant_page_numbers}\n'336 llm_answer_end_time = time.time()337 llm_answer_latency = llm_answer_end_time - retrieval_end_time338 if debug==True:339 llm_answer += f"\nQuery embedding latency: {query_embedding_latency:.3f} seconds"340 llm_answer += f"\nRetrieval latency: {retrieval_latency:.3f} seconds"341 llm_answer += f"\nLLM latency: {llm_answer_latency:.3f} seconds"342 llm_answer += f"\nTotal latency: {llm_answer_end_time - start_time:.3f} seconds\n"343 llm_answer += '\n' + '=' * 30344 llm_answer += '\n' + '=' * 30345 llm_answer += "\nRELEVANT CHUNKS"346 for i, chunk in enumerate(relevant_chunks):347 llm_answer += f"\nChunk {i+1}:"348 llm_answer += f"\nText (first 150 chars): {chunk.get('text', '')[:150]}..."349 if 'similarity' in chunk: llm_answer += f"\nCosine Similarity: {chunk['similarity']}"350 if 'rank' in chunk: llm_answer += f"\nRank: {chunk['rank']}"351 print('\n\n')352 return llm_answer353 except Exception as e:354 return f"Error in get_answer: {e}"355 traceback.print_exc()356 return None357 358def answer_question_gradio(question):359 """360 Wrapper function for Gradio interface to call the RAG model.361 """362 #satisfies_zerogpu()363 print(f"Received question: {question}")364 try:365 response = get_answer(question, debug=False) # Set debug=False for deployment366 if response: # Ensure response is not None367 return response368 else:369 #return "I could not generate an answer for that question."370 return "Unable to generate an answer for this question"371 except Exception as e:372 print(f"Error in answer_question_gradio: {e}")373 traceback.print_exc()374 return "An error occurred while processing your request: {str(e)}."375 376# Create the Gradio interface377iface = gr.Interface(378 fn=answer_question_gradio,379 inputs=gr.Textbox(lines=2, placeholder="Enter your question about PostgreSQL..."),380 outputs="text",381 title="PostgreSQL AI Documentation Assistant",382 description="Ask any question about PostgreSQL 18 and get an answer backed by RAG on the official PostgreSQL 18 Documentation."383)384 385if __name__ == "__main__":386 iface.launch()