Update main.py

main.py

@@ -1,9 +1,7 @@
 from fastapi import FastAPI, UploadFile, File
 import json
-import pypdf
-import io
 
-from db import conn, cursor
+from db import get_db_connection
 from embeddings import semantic_chunking, embedding_model
 from retrieval import retrieve_top_chunks
 from llm import build_context_string, call_llm
@@ -16,24 +14,23 @@ def health():
 
 @app.post("/ingest")
 async def ingest(file: UploadFile = File(...)):
-    if file.filename.endswith(".pdf"):
-        pdf_reader = pypdf.PdfReader(io.BytesIO(await file.read()))
-        text = ""
-        for page in pdf_reader.pages:
-            text += page.extract_text() + "\n"
-    else:
-        text = (await file.read()).decode("utf-8")
+    text = (await file.read()).decode("utf-8")
 
     chunks = semantic_chunking(text)
     embeddings = embedding_model.encode(chunks)
 
-    for i, (chunk, emb) in enumerate(zip(chunks, embeddings)):
-        cursor.execute(
-            "INSERT INTO chunks (document, chunk_id, text, embedding) VALUES (%s, %s, %s, %s)",
-            (file.filename, i, chunk, json.dumps(emb.tolist()))
-        )
-
-    conn.commit()
+    conn = get_db_connection()
+    cursor = conn.cursor()
+    try:
+        for i, (chunk, emb) in enumerate(zip(chunks, embeddings)):
+            cursor.execute(
+                "INSERT INTO chunks (document, chunk_id, text, embedding) VALUES (%s, %s, %s, %s)",
+                (file.filename, i, chunk, json.dumps(emb.tolist()))
+            )
+        conn.commit()
+    finally:
+        cursor.close()
+        conn.close()
     return {"chunks_ingested": len(chunks)}
 
 @app.post("/ask")
@@ -80,12 +77,24 @@ def ask(question: str):
 
 @app.post("/reset")
 def reset_db():
-    cursor.execute("TRUNCATE TABLE chunks")
-    conn.commit()
+    conn = get_db_connection()
+    cursor = conn.cursor()
+    try:
+        cursor.execute("TRUNCATE TABLE chunks")
+        conn.commit()
+    finally:
+        cursor.close()
+        conn.close()
     return {"status": "Database cleared"}
 
 @app.get("/view")
 def view_db(limit: int = 50):
-    cursor.execute(f"SELECT id, document, chunk_id, left(text, 200) as text FROM chunks ORDER BY id DESC LIMIT {limit}")
-    rows = cursor.fetchall()
-    return [{"id": r[0], "document": r[1], "chunk_id": r[2], "text_snippet": r[3]} for r in rows]
+    conn = get_db_connection()
+    cursor = conn.cursor()
+    try:
+        cursor.execute(f"SELECT id, document, chunk_id, left(text, 200) as text FROM chunks ORDER BY id DESC LIMIT {limit}")
+        rows = cursor.fetchall()
+        return [{"id": r[0], "document": r[1], "chunk_id": r[2], "text_snippet": r[3]} for r in rows]
+    finally:
+        cursor.close()
+        conn.close()