first commit

data_handling.py

@@ -0,0 +1,81 @@
+# 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