82 lines
3.1 KiB
Python
82 lines
3.1 KiB
Python
# data_handling.py
|
|
import os
|
|
from pathlib import Path
|
|
from typing import List, Optional
|
|
import logging # Added logging
|
|
|
|
from haystack import Document
|
|
from milvus_haystack import MilvusDocumentStore
|
|
|
|
# Import config variables needed
|
|
from config import (
|
|
OPENAI_EMBEDDING_DIM, # Keep for logging/validation if desired, but not passed to init
|
|
USER_ID_PREFIX,
|
|
MILVUS_PERSIST_BASE_DIR,
|
|
MILVUS_INDEX_PARAMS,
|
|
MILVUS_SEARCH_PARAMS,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__) # Use logger
|
|
|
|
|
|
# get_user_milvus_path function remains the same
|
|
def get_user_milvus_path(user_id: str, base_dir: Path = MILVUS_PERSIST_BASE_DIR) -> str:
|
|
# user_db_dir = base_dir / user_id
|
|
# user_db_dir.mkdir(parents=True, exist_ok=True)
|
|
return str("milvus_lite.db")
|
|
|
|
|
|
def initialize_milvus_lite(user_id: str) -> MilvusDocumentStore:
|
|
"""
|
|
Initializes Milvus Lite DocumentStore for a user using milvus-haystack.
|
|
Dimension is inferred by Milvus upon first write, not passed here.
|
|
"""
|
|
print(f"Initializing Milvus Lite store for user: {user_id}")
|
|
milvus_uri = get_user_milvus_path(user_id)
|
|
print(f"Milvus Lite URI: {milvus_uri}")
|
|
# Log the dimension expected based on config, even if not passed directly
|
|
print(f"Expecting Embedding Dimension (for first write): {OPENAI_EMBEDDING_DIM}")
|
|
|
|
document_store = MilvusDocumentStore(
|
|
connection_args={"uri": milvus_uri},
|
|
collection_name=user_id, # Default or customize
|
|
index_params=MILVUS_INDEX_PARAMS, # Pass index config
|
|
search_params=MILVUS_SEARCH_PARAMS, # Pass search config
|
|
drop_old=False, # Keep drop_old for testing convenience
|
|
)
|
|
# Note: The actual schema dimension is set when the first document with an embedding is written.
|
|
print(f"Milvus Lite store instance created for user {user_id} at {milvus_uri}")
|
|
return document_store
|
|
|
|
|
|
# add_user_document_to_store and get_user_documents can remain if needed for other purposes,
|
|
def add_user_document_to_store(
|
|
document_store: MilvusDocumentStore, user_id: str, text: str
|
|
):
|
|
doc = Document(content=text, meta={"user_id": user_id})
|
|
print(f"Adding document for user {user_id}: '{text[:50]}...'")
|
|
document_store.write_documents([doc])
|
|
|
|
|
|
# get_user_documents function remains the same
|
|
def get_user_documents(
|
|
document_store: MilvusDocumentStore, user_id: str
|
|
) -> List[Document]:
|
|
print(f"Retrieving all documents for user {user_id}...")
|
|
all_docs = document_store.get_all_documents()
|
|
print(f"Found {len(all_docs)} documents for user {user_id}.")
|
|
return all_docs
|
|
|
|
|
|
# Optional: Test code similar to before, but now using the OpenAI dimension
|
|
if __name__ == "__main__":
|
|
test_user = "test_user_openai_data"
|
|
store = initialize_milvus_lite(test_user)
|
|
# Add dummy docs (won't be embedded here, just stored)
|
|
add_user_document_to_store(store, test_user, "第一个文档，关于 OpenAI。")
|
|
add_user_document_to_store(store, test_user, "第二个文档，使用 API。")
|
|
docs = get_user_documents(store, test_user)
|
|
for d in docs:
|
|
print(f" - {d.content} (Meta: {d.meta})")
|
|
# Cleanup code similar to before
|