refactoring get_vectorstore

src/app.py

@@ -12,7 +12,7 @@ if str(_project_root) not in sys.path:
 
 from src.core.ingest import load_documents, get_chunks, ingest_documents
 from src.core.retrieval import generate, retrieval
-from src.core.index import MetaData
+from src.core.index import MetaData, get_vectorstore
 from src.core.synthetic_data import EVAL_QUERIES, SYNTHETIC_DOCUMENTS
 from src.core.eval import run_full_evaluation, save_results
 from src.core.eval import generate_summary_report, setup_test_data
@@ -49,9 +49,10 @@ def ingest_files(files:List[str], index_name:str, lang:Literal["en", "ja"], doma
         language=lang, domain=domain, section=section, topic=topic, doc_type=doc_type
     )
     try:
+        vectorstore = get_vectorstore(index_name)
         docs = load_documents(files)
         chunks = get_chunks(docs, filter_data)
-        message = ingest_documents(chunks, index_name)
+        message = ingest_documents(chunks, vectorstore)
     except Exception as e:
         message = f"Error during ingestion: {str(e)}"
         print(message)
@@ -77,8 +78,8 @@ def _rag_query(
     print(f"Active Filters: {active_filters.model_dump()}")
 
     ret_start_time = time.time()
-    
-    docs = retrieval(question, index_name, active_filters)
+    vectorstore = get_vectorstore(index_name)
+    docs = retrieval(question, active_filters, vectorstore)
     retrieval_results = [doc.page_content + _add_metric(doc) for doc in docs]
     snippets_md = "\n\n---\n\n".join(retrieval_results)
 

src/core/eval.py

@@ -17,7 +17,7 @@ import numpy as np
 from langchain_core.documents import Document
 from langchain_openai import OpenAIEmbeddings
 from dotenv import load_dotenv, find_dotenv
-from .index import MetaData
+from .index import MetaData, get_vectorstore
 from .retrieval import retrieval, generate
 from .ingest import ingest_documents, get_chunks
 from .synthetic_data import SYNTHETIC_DOCUMENTS, EVAL_QUERIES, EvalQuery
@@ -28,7 +28,6 @@ load_dotenv()
 # Embedding model for semantic similarity
 emb_model = OpenAIEmbeddings(model="text-embedding-3-small", dimensions=1536)
 
-
 @dataclass
 class EvalResult:
     """Evaluation result for a single query"""
@@ -131,7 +130,8 @@ def evaluate_single_query(
     
     # Retrieval
     ret_start = time.time()
-    docs = retrieval(eval_query.query, eval_query.collection, filters)
+    vectorstore = get_vectorstore("eval_"+eval_query.collection)
+    docs = retrieval(eval_query.query, filters, vectorstore)
     ret_end = time.time()
     ret_latency = (ret_end - ret_start) * 1000  # Convert to ms
     
@@ -443,11 +443,9 @@ def setup_test_data(collections: List[str] = None):
             metadata = MetaData(**metadata)
             chunks = get_chunks([doc], metadata)
             documents.extend(chunks)
-
-        # vectorstore = get_vectorstore(collection_name)
-        # ids = [str(uuid.uuid4()) for _ in range(len(documents))]
-        # vectorstore.add_documents(documents, ids=ids)
-        ingest_documents(documents, collection_name)
+            
+        vectorstore = get_vectorstore("eval_"+collection_name, drop_old=True)
+        ingest_documents(documents, vectorstore)
         tot_docs += len(docs)
         print(f"✓ Completed '{collection_name}' collection")
     

src/core/index.py

@@ -26,12 +26,13 @@ emb_model = OpenAIEmbeddings(model="text-embedding-3-small", dimensions=1536)
 MILVUS_URI = os.getenv("MILVUS_URI","./data/rag_task.db")
 MILVUS_API_KEY = os.getenv("MILVUS_API_KEY","")
 
-def get_vectorstore(collection_name: str) -> Milvus:
+def get_vectorstore(collection_name: str, drop_old=False) -> Milvus:
     vectorstore = Milvus(
         embedding_function=emb_model,
         collection_name=collection_name,
         connection_args={"uri": MILVUS_URI,"token": MILVUS_API_KEY},
         index_params={"index_type": "FLAT", "metric_type": "L2"},
+        drop_old=drop_old,
     )
     # builtin_function=BM25BuiltInFunction(output_field_names="sparse"),
     # text_field="text",

src/core/ingest.py

@@ -1,12 +1,12 @@
 from langchain_community.document_loaders import PDFMinerLoader,TextLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_milvus import Milvus
 from langchain_core.documents import Document
 from langchain_openai import ChatOpenAI
 from dotenv import load_dotenv, find_dotenv
 from typing import List
 import uuid
 
-from src.core.index import get_vectorstore
 from .index import MetaData
 from .utils import mask_pii
 
@@ -59,14 +59,12 @@ def get_chunks(documents: List[Document], metadata: MetaData):
     return chunks
 
 
-def ingest_documents(docs: List[Document], collection_name: str, vectorstore = None):
+def ingest_documents(docs: List[Document], vectorstore:Milvus):
     """Ingest documents into the specified vectorstore collection."""
-    if vectorstore is None:
-        vectorstore = get_vectorstore(collection_name)
-        
+    
     ids = [str(uuid.uuid4()) for _ in range(len(docs))]
     vectorstore.add_documents(docs, ids=ids)
-    success_message = f"Ingested {len(docs)} documents into {collection_name} index."
+    success_message = f"Ingested {len(docs)} documents into {vectorstore.collection_name} index."
     print(success_message)
     return success_message
 

src/core/retrieval.py

@@ -1,10 +1,10 @@
 from langchain_core.documents import Document
-from langchain_openai import ChatOpenAI
 from langchain_community.retrievers import BM25Retriever
+from langchain_openai import ChatOpenAI
+from langchain_milvus import Milvus
 from dotenv import load_dotenv, find_dotenv
 from typing import List
-from .index import get_vectorstore, MetaData
-
+from .index import MetaData
 find_dotenv()
 load_dotenv()
 
@@ -23,12 +23,11 @@ def reranker(query: str, docs: List[Document]) -> List[Document]:
 
 
 def retrieval(
-    query: str, collection_name: str, filter_data: MetaData
+    query: str, filter_data: MetaData, vectorstore: Milvus
 ) -> List[tuple[Document, float]]:
     """Retrieve relevant documents from the vector store based on the query and filters."""
-    vectorstore = get_vectorstore(collection_name)
     print(
-        f"RETRIEVAL query: {query[:40]}, for {collection_name} collection, with filters: {filter_data}"
+        f"RETRIEVAL query: {query[:40]}, for {vectorstore.collection_name} collection, with filters: {filter_data}"
     )
 
     filters = [f'language == "{filter_data.language}"']