feat(Vector Stores): Move the OpenAI and MCP servers to use pluggable stores.

Author: stronk7

wiki_rag/mcp_server/main.py

@@ -12,12 +12,13 @@
 from dotenv import load_dotenv
 from langfuse.langchain import CallbackHandler
 
-import wiki_rag.index as index
 import wiki_rag.mcp_server as mcp_global
+import wiki_rag.vector as vector
 
 from wiki_rag import LOG_LEVEL, ROOT_DIR, __version__, server
 from wiki_rag.search.util import ContextSchema
 from wiki_rag.util import setup_logging
+from wiki_rag.vector import load_vector_store
 
 
 def main():
@@ -68,10 +69,10 @@ def main():
         sys.exit(1)
         # TODO: Validate that only numbers, letters and underscores are used.
 
-    index.milvus_url = os.getenv("MILVUS_URL")
-    if not index.milvus_url:
-        logger.error("Milvus URL not found in environment. Exiting.")
-        sys.exit(1)
+    index_vendor = os.getenv("INDEX_VENDOR")
+    if not index_vendor:
+        logger.warning("Index vendor (INDEX_VENDOR) not found in environment. Defaulting to 'milvus'.")
+        index_vendor = "milvus"
 
     # If LangSmith tracing is enabled, put a name for the project and verify that all required env vars are set.
     if os.getenv("LANGSMITH_TRACING", "false") == "true":
@@ -163,6 +164,8 @@ def main():
     else:
         mcp_port = 8081
 
+    vector.store = load_vector_store(index_vendor)  # Set up the global wiki_rag.vector.store to be used elsewhere.
+
     # Calculate the file that we are going to use as source for the resources.
     input_candidate = ""
     for file in sorted(loader_dump_path.iterdir()):

wiki_rag/server/main.py

@@ -14,11 +14,12 @@
 from dotenv import load_dotenv
 from langfuse.langchain import CallbackHandler
 
-import wiki_rag.index as index
+import wiki_rag.vector as vector
 
 from wiki_rag import LOG_LEVEL, ROOT_DIR, __version__, server
 from wiki_rag.search.util import ContextSchema, build_graph
 from wiki_rag.util import setup_logging
+from wiki_rag.vector import load_vector_store
 
 
 def main():
@@ -69,10 +70,10 @@ def main():
         sys.exit(1)
         # TODO: Validate that only numbers, letters and underscores are used.
 
-    index.milvus_url = os.getenv("MILVUS_URL")
-    if not index.milvus_url:
-        logger.error("Milvus URL not found in environment. Exiting.")
-        sys.exit(1)
+    index_vendor = os.getenv("INDEX_VENDOR")
+    if not index_vendor:
+        logger.warning("Index vendor (INDEX_VENDOR) not found in environment. Defaulting to 'milvus'.")
+        index_vendor = "milvus"
 
     # If LangSmith tracing is enabled, put a name for the project and verify that all required env vars are set.
     if os.getenv("LANGSMITH_TRACING", "false") == "true":
@@ -141,6 +142,8 @@ def main():
 
     contextualisation_model = os.getenv("CONTEXTUALISATION_MODEL")
 
+    vector.store = load_vector_store(index_vendor)  # Set up the global wiki_rag.vector.store to be used elsewhere.
+
     wrapper_api_base = os.getenv("WRAPPER_API_BASE")
     if not wrapper_api_base:
         logger.error("Wrapper API base not found in environment. Exiting.")