153 lines
4.6 KiB
Python
153 lines
4.6 KiB
Python
"""
|
|
Create the ProductProfile Weaviate collection with Gemini embeddings.
|
|
|
|
Uses Google AI Studio (Gemini API) with:
|
|
- model: gemini-embedding-001
|
|
- task type: RETRIEVAL_DOCUMENT (optimized for indexing documents)
|
|
|
|
Usage:
|
|
python setup_weaviate.py # create collection if it does not exist
|
|
python setup_weaviate.py --recreate # delete and recreate (required after embedding config changes)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
from typing import Dict
|
|
|
|
try:
|
|
import weaviate
|
|
except ModuleNotFoundError:
|
|
print(
|
|
"[FATAL] Missing dependency 'weaviate'. Activate the project virtual environment "
|
|
"and install requirements:\n"
|
|
" venv\\Scripts\\activate\n"
|
|
" pip install -r requirements.txt",
|
|
file=sys.stderr,
|
|
)
|
|
sys.exit(1)
|
|
|
|
from dotenv import load_dotenv
|
|
from weaviate.classes.config import Configure, DataType, Property
|
|
|
|
COLLECTION_NAME = "ProductProfile"
|
|
EMBEDDING_MODEL = "gemini-embedding-001"
|
|
EMBEDDING_TASK_TYPE = "RETRIEVAL_DOCUMENT"
|
|
|
|
|
|
def load_config() -> Dict[str, str]:
|
|
load_dotenv()
|
|
gemini_api_key = os.getenv("GEMINI_API_KEY", "")
|
|
if not gemini_api_key:
|
|
raise EnvironmentError(
|
|
"Missing required environment variable: GEMINI_API_KEY. "
|
|
"Copy .env.example to .env and fill in the values."
|
|
)
|
|
return {
|
|
"GEMINI_API_KEY": gemini_api_key,
|
|
"WEAVIATE_HOST": os.getenv("WEAVIATE_HOST", "localhost"),
|
|
"WEAVIATE_HTTP_PORT": os.getenv("WEAVIATE_HTTP_PORT", "8080"),
|
|
"WEAVIATE_GRPC_PORT": os.getenv("WEAVIATE_GRPC_PORT", "50051"),
|
|
}
|
|
|
|
|
|
def connect_weaviate(cfg: Dict[str, str]) -> weaviate.WeaviateClient:
|
|
return weaviate.connect_to_local(
|
|
host=cfg["WEAVIATE_HOST"],
|
|
port=int(cfg["WEAVIATE_HTTP_PORT"]),
|
|
grpc_port=int(cfg["WEAVIATE_GRPC_PORT"]),
|
|
headers={"X-Goog-Studio-Api-Key": cfg["GEMINI_API_KEY"]},
|
|
)
|
|
|
|
|
|
def ensure_product_profile_collection(
|
|
client: weaviate.WeaviateClient,
|
|
*,
|
|
recreate: bool = False,
|
|
) -> bool:
|
|
"""
|
|
Ensure the ProductProfile collection exists with the expected embedding config.
|
|
|
|
Returns True if a new collection was created, False if an existing one was kept.
|
|
"""
|
|
if client.collections.exists(COLLECTION_NAME):
|
|
if recreate:
|
|
client.collections.delete(COLLECTION_NAME)
|
|
else:
|
|
return False
|
|
|
|
client.collections.create(
|
|
name=COLLECTION_NAME,
|
|
properties=[
|
|
Property(
|
|
name="product_id",
|
|
data_type=DataType.TEXT,
|
|
skip_vectorization=True,
|
|
index_filterable=True,
|
|
),
|
|
Property(
|
|
name="product_name",
|
|
data_type=DataType.TEXT,
|
|
skip_vectorization=True,
|
|
index_filterable=True,
|
|
),
|
|
Property(
|
|
name="retrieval_summary",
|
|
data_type=DataType.TEXT,
|
|
),
|
|
],
|
|
vector_config=Configure.Vectors.text2vec_google_gemini(
|
|
source_properties=["retrieval_summary"],
|
|
model=EMBEDDING_MODEL,
|
|
task_type=EMBEDDING_TASK_TYPE,
|
|
vectorize_collection_name=False,
|
|
),
|
|
)
|
|
return True
|
|
|
|
|
|
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Create the ProductProfile Weaviate collection with Gemini embeddings."
|
|
)
|
|
parser.add_argument(
|
|
"--recreate",
|
|
action="store_true",
|
|
help="Delete and recreate the collection (use after changing embedding settings).",
|
|
)
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
cfg = load_config()
|
|
client = connect_weaviate(cfg)
|
|
|
|
try:
|
|
had_collection = client.collections.exists(COLLECTION_NAME)
|
|
created = ensure_product_profile_collection(client, recreate=args.recreate)
|
|
if created and had_collection and args.recreate:
|
|
print(
|
|
f"Recreated collection '{COLLECTION_NAME}' "
|
|
f"(model={EMBEDDING_MODEL}, task_type={EMBEDDING_TASK_TYPE})."
|
|
)
|
|
elif created:
|
|
print(
|
|
f"Created collection '{COLLECTION_NAME}' "
|
|
f"(model={EMBEDDING_MODEL}, task_type={EMBEDDING_TASK_TYPE})."
|
|
)
|
|
else:
|
|
print(f"Collection '{COLLECTION_NAME}' already exists; no changes made.")
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except EnvironmentError as exc:
|
|
print(f"[FATAL] {exc}", file=sys.stderr)
|
|
sys.exit(1)
|