main.py

from fastapi import FastAPI
from http_models import *
from sentence_transformers import SentenceTransformer, CrossEncoder
from faiss import IndexHNSWFlat, IndexFlatIP # IndexFlatIP more accurate, but slower
import numpy as np
import gc
import os
from utils import read_files_from_dir, read_dirs_from_dir, b64d, b64e
from database import insert_cache, clear_cache, get_docs_by_ids, create_connection
import pickle


# setup - Globals -
app = FastAPI(title="Mini RankBrain API for files and directories")
BATCH_SIZE = 5000
encoder = SentenceTransformer("all-MiniLM-L6-v2")

# Create FAISS index
dimension = encoder.get_sentence_embedding_dimension()

# Load cross-encoder
cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

# init cache
cx, cu = create_connection(first_time=True)
cx.close()

def index_strings(encoder, dir_path: str, table: str, read_func, index, batch_size=BATCH_SIZE, max_workers=8):
    """indexing in parallel with FAISS"""
    def save_and_clear():
        vectors = encoder.encode(batch_buffer, normalize_embeddings=True)
        index.add(np.array(vectors, dtype=np.float32))

        # cache - 100
        cx, cu = create_connection()
        insert_cache(cx, cu, b64e(pickle.dumps(batch_buffer)), table)
        cx.close()

        batch_buffer.clear()

        # run garbage collection after batch
        gc.collect()

    batch_buffer = []

    for string in read_func(dir_path, max_workers=max_workers):
        batch_buffer.append(string)

        if len(batch_buffer) >= batch_size:
            save_and_clear()

    if batch_buffer:
        save_and_clear()

#? if text files cached into database:
#batch_size = 100
#end_reached = False
#offset = 0
#while not end_reached:
    #docs = get_doc_vectors(cursor, batch_size, offset)
    #if len(docs) == 0:
    #    end_reached = True
    #else:
    #    index.add(docs)
    #    offset += batch_size



@app.post("/search/directory", response_model=QueryResponse)
def search_directories(request: QueryRequest_directory):
    if not os.path.exists(request.path):
        return QueryResponse(query=request.query, results=["Not valid path"])

    global dir_path
    dir_path = request.path
    query_vector = encoder.encode([request.query], normalize_embeddings=True)
    # search directory
    index_IP = IndexFlatIP(dimension)
    index_strings(encoder, request.path, "cache.directories", read_dirs_from_dir, index_IP)

    # distance and indicies
    _, I = index_IP.search(np.array(query_vector, dtype=np.float32), request.top_k)
    cx, cu = create_connection()

    indices = []
    for i in I[0]:
        tmp = int((i // BATCH_SIZE) + 1)
        if not tmp in indices:
            indices.append(tmp)

    candidates = get_docs_by_ids(cu, indices, "cache.directories")
    valid_candidates = []
    for c, index in zip(candidates, indices):
        for i in I[0]:
            if int((i // BATCH_SIZE) + 1) == index:
                candidate = pickle.loads(b64d(c.encode("utf-8")))[i % BATCH_SIZE]
                valid_candidates.append(candidate)

    candidates = valid_candidates

    # debug
    #print(list(map(lambda x: (x + 1 // BATCH_SIZE), I[0])))
    #print(indices)
    #print(candidates)

    # cross-encoder
    pairs = [(request.query, doc) for doc in candidates]
    scores = cross_encoder.predict(pairs)

    reranked_docs = [doc for _, doc in sorted(zip(scores, candidates), reverse=True)]

    # clear after search
    index_IP = IndexFlatIP(dimension)
    clear_cache(cu, clear_all=True)
    gc.collect()
    cx.close()

    return QueryResponse(query=request.query, results=reranked_docs)


@app.post("/search/files", response_model=QueryResponse)
def search_files(request: QueryRequest_files):
    try:
        if not os.path.exists(dir_path):
            return QueryResponse(query=request.query, results=["Not valid path"])
    except NameError:
        return QueryResponse(query=request.query, results=["Not valid path"])

    query_vector = encoder.encode([request.query], normalize_embeddings=True)
    #? search files, too slow for large directories
    index_HNSW = IndexHNSWFlat(dimension, 16)
    index_strings(encoder, dir_path, "cache.files", read_files_from_dir, index_HNSW)

    # distance and indicies
    _, I = index_HNSW.search(np.array(query_vector, dtype=np.float32), request.top_k)

    cx, cu = create_connection()
    # cross-encoder database starts from index 1 not 0
    indices = []
    for i in I[0]:
        tmp = int((i // BATCH_SIZE) + 1)
        if not tmp in indices:
            indices.append(tmp)

    candidates = get_docs_by_ids(cu, indices, "cache.directories")
    valid_candidates = []
    for c, index in zip(candidates, indices):
        for i in I[0]:
            if int((i // BATCH_SIZE) + 1) == index:
                candidate = pickle.loads(b64d(c.encode("utf-8")))[i % BATCH_SIZE]
                valid_candidates.append(candidate)

    candidates = valid_candidates

    # cross-encoder
    pairs = [(request.query, doc) for doc in candidates]
    scores = cross_encoder.predict(pairs)

    # : split to get file path .split(": ")[0]
    reranked_docs = [doc for _, doc in sorted(zip(scores, candidates), reverse=True)]

    # clear after search
    index_HNSW = IndexHNSWFlat(dimension, 16)
    clear_cache(cu)
    gc.collect()
    cx.close()

    return QueryResponse(query=request.query, results=reranked_docs)

@app.get("/")
async def root():
    return {"message": "Hello World"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=False)