few refinements

api/retrieve_context.py

@@ -1,28 +1,30 @@
-import numpy as np
-import faiss, logging
+import faiss
+import logging
+from pathlib import Path
 from sentence_transformers import SentenceTransformer
+
 from config import TRANSCRIPT_INDEX
-logging.basicConfig(format="%(asctime)s %(levelname)s:%(message)s", level=logging.INFO)
 
-embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
-logging.info("Loaded embedding model for query.")
+logger = logging.getLogger(__name__)
 
+model = SentenceTransformer("all-MiniLM-L6-v2")
 index = faiss.read_index(TRANSCRIPT_INDEX)
-logging.info(f"Loaded FAISS index from {TRANSCRIPT_INDEX}.")
-
-def retrieve_transcripts(query, file_path, transcripts, top_k=3):
-    logging.info("Starting retrieval process...")
 
-    query_embedding = embedding_model.encode([query], convert_to_tensor=False)
-    logging.info("Encoded query to embedding.")
 
-    distances, indices = index.search(np.array(query_embedding), top_k)
-    logging.info(f"Retrieved top {top_k} results from index.")
+def retrieve_transcripts(
+    query: str,
+    file_paths: list[Path],
+    transcripts: list[str],
+    top_k: int = 3,
+) -> list[str]:
+    query_embedding = model.encode([query])
+    distances, indices = index.search(query_embedding, top_k)
 
     results = []
     for idx in indices[0]:
-        results.append(transcripts[idx])
-        logging.info(f"Retrieved transcript from: {file_path[idx]}")
+        if idx != -1:
+            results.append(transcripts[idx])
+            logger.info(f"Retrieved transcript from: {file_paths[idx]}")
 
-    logging.info("Retrieval process completed.")
+    logger.info("Retrieval process completed.")
     return results