From 1f9aa7ed9c7825551ad410a91f6b8200b4f2c926 Mon Sep 17 00:00:00 2001 From: Prakhar-Sethi012 Date: Sun, 14 Jun 2026 04:17:05 +0530 Subject: [PATCH] feat: implement pgvector semantic and hybrid rrf search with background indexing --- app/database.py | 10 ++++++ app/main.py | 5 +++ app/processing/search.py | 10 +++--- app/routes/search.py | 73 ++++++++++++++++++++++++++++++++++++++++ app/services/indexer.py | 56 ++++++++++++++++++++++++++++++ 5 files changed, 149 insertions(+), 5 deletions(-) create mode 100644 app/routes/search.py create mode 100644 app/services/indexer.py diff --git a/app/database.py b/app/database.py index 89f8575..d1aefa4 100644 --- a/app/database.py +++ b/app/database.py @@ -18,6 +18,16 @@ autocommit=False, ) +def get_db(): + """ + Dependency that creates a new SQLAlchemy session per request + and ensures it is safely closed after the request is finished. + """ + db = SessionLocal() + try: + yield db + finally: + db.close() class Base(DeclarativeBase): pass \ No newline at end of file diff --git a/app/main.py b/app/main.py index 3a30552..9210162 100644 --- a/app/main.py +++ b/app/main.py @@ -1,7 +1,12 @@ from fastapi import FastAPI +from app.routes import search + app = FastAPI( title="WhereTF Backend" ) + +app.include_router(search.router) + @app.get("/health", tags=["System"]) def health_check(): """ diff --git a/app/processing/search.py b/app/processing/search.py index d1350b5..942150d 100644 --- a/app/processing/search.py +++ b/app/processing/search.py @@ -118,7 +118,7 @@ def _sql_vector(top_k: int, file_filter: bool) -> str: fc.content_text, 1 - (fc.embedding <=> CAST(:query_vec AS vector)) AS score FROM file_content fc -JOIN file f ON f.id = fc.file_id +JOIN files f ON f.id = fc.file_id WHERE 1=1 {where} ORDER BY fc.embedding <=> CAST(:query_vec AS vector) ASC LIMIT :top_k; @@ -146,7 +146,7 @@ def _sql_keyword(top_k: int, file_filter: bool) -> str: fc.content_text, ts_rank_cd(fc.keyword_tokens, plainto_tsquery('english', :query_text)) AS score FROM file_content fc -JOIN file f ON f.id = fc.file_id +JOIN files f ON f.id = fc.file_id WHERE fc.keyword_tokens @@ plainto_tsquery('english', :query_text) {where} ORDER BY score DESC LIMIT :top_k; @@ -182,7 +182,7 @@ def _sql_hybrid(top_k: int, file_filter: bool, rrf_k: int = 60) -> str: ORDER BY fc.embedding <=> CAST(:query_vec AS vector) ASC ) AS rank FROM file_content fc - JOIN file f ON f.id = fc.file_id + JOIN files f ON f.id = fc.file_id WHERE 1=1 {where} LIMIT {pool} ), @@ -196,7 +196,7 @@ def _sql_hybrid(top_k: int, file_filter: bool, rrf_k: int = 60) -> str: ) DESC ) AS rank FROM file_content fc - JOIN file f ON f.id = fc.file_id + JOIN files f ON f.id = fc.file_id WHERE fc.keyword_tokens @@ plainto_tsquery('english', :query_text) {where} LIMIT {pool} ), @@ -217,7 +217,7 @@ def _sql_hybrid(top_k: int, file_filter: bool, rrf_k: int = 60) -> str: fused.rrf_score AS score FROM fused JOIN file_content fc ON fc.id = fused.chunk_id -JOIN file f ON f.id = fc.file_id +JOIN files f ON f.id = fc.file_id ORDER BY fused.rrf_score DESC LIMIT :top_k; """.strip() diff --git a/app/routes/search.py b/app/routes/search.py new file mode 100644 index 0000000..35f0a2b --- /dev/null +++ b/app/routes/search.py @@ -0,0 +1,73 @@ +from fastapi import APIRouter, Depends, Query, HTTPException +from sqlalchemy.orm import Session +from sqlalchemy import text +from app.database import get_db +from app.processing.search import build_query +router = APIRouter(tags=["Search Engine"]) + +@router.post("/search/") +def search_documents( + query: str = Query(..., description="The search text from the user"), + mode: str = Query("hybrid", description="Search strategy: 'vector', 'keyword', or 'hybrid'"), + top_k: int = Query(5, description="Number of results to return"), + db: Session = Depends(get_db) +): + """ + Executes a high-performance search against the pgvector database. + Defaults to Hybrid Search (Reciprocal Rank Fusion). + """ + # Validation + valid_modes = ["vector", "keyword", "hybrid"] + if mode not in valid_modes: + raise HTTPException(status_code=400, detail=f"Mode must be one of {valid_modes}") + + try: + # embedding generation + payload = build_query(query, mode=mode, top_k=top_k) + + # Exact sql string extraction + sql_string = text(payload["sql"][mode]) + + # Extracting parameters + params = payload["params"] + + raw_results = db.execute(sql_string, params).mappings().all() + + # Formatting and returning the results + return { + "status": "success", + "query": query, + "mode": mode, + "results": [dict(row) for row in raw_results] + } + + except Exception as e: + raise HTTPException(status_code=500, detail=f"Search Engine Error: {str(e)}") + + +''' + + +# THIS WAS DONE TO TEST WHETHER EVERYTHING WORKS PERFECTLY OR NOT + +from fastapi import BackgroundTasks +from app.services.indexer import background_index_file + +@router.post("/test-indexer/") +def trigger_dummy_index(background_tasks: BackgroundTasks): + """ + TEMPORARY ROUTE: Creates a fake syllabus file and feeds it to the background indexer. + Delete this once Aryan finishes the real upload route! + """ + dummy_path = "/tmp/dummy_syllabus.txt" + + # 1. Create a fake document on the server + with open(dummy_path, "w") as f: + f.write("Welcome to the Computer Science program. The core syllabus for first-year students includes Python programming, Data Structures, and Algorithms. The exact course code for Data Structures is CSE1001. Good luck with your studies!") + + # 2. Hand it to your engine + background_tasks.add_task(background_index_file, dummy_path) + + return {"status": "Cheat code activated. Dummy file sent to indexer!"} + +''' \ No newline at end of file diff --git a/app/services/indexer.py b/app/services/indexer.py new file mode 100644 index 0000000..8e24576 --- /dev/null +++ b/app/services/indexer.py @@ -0,0 +1,56 @@ +import os +import logging +from app.database import SessionLocal +from app.models import File, FileContent +from app.processing import process + +logger = logging.getLogger(__name__) + +def background_index_file(temp_file_path: str): + """ + Background task that processes a document, extracts AI embeddings, + saves them to pgvector, and cleans up the temporary file. + """ + db = SessionLocal() + + try: + logger.info(f"Starting background indexing for: {temp_file_path}") + + file_rows, content_rows = process(temp_file_path) + for fr in file_rows: + existing = db.query(File).filter_by(file_path=fr["file_path"]).first() + if existing: + for k, v in fr.items(): + setattr(existing, k, v) + else: + db.add(File(**fr)) + + db.flush() + + for cr in content_rows: + file_path = cr.pop("file_path") + file_obj = db.query(File).filter_by(file_path=file_path).one() + + db.add(FileContent( + file_id=file_obj.id, + chunk_index=cr["chunk_index"], + content_text=cr["content_text"], + embedding=cr["embedding"] + )) + + # permanently saving the vectors + db.commit() + logger.info(f"Successfully indexed {len(file_rows)} file(s) and {len(content_rows)} chunks.") + + except Exception as e: + # undo the changes if anything fails + db.rollback() + logger.error(f"Failed to index file {temp_file_path}: {str(e)}") + + finally: + db.close() + + # Server Cleanup,deleting the temp files + if os.path.exists(temp_file_path): + os.remove(temp_file_path) + logger.info(f"Cleaned up temporary file: {temp_file_path}") \ No newline at end of file