Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 23 additions & 7 deletions datastew/repository/weaviate.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,15 @@
import json
import logging
import socket
import time
from typing import Any, Dict, List, Literal, Optional, Sequence, Union

import weaviate
from typing_extensions import deprecated
from weaviate import WeaviateClient
from weaviate.classes.query import Filter, MetadataQuery, QueryReference
from weaviate.collections import Collection
from weaviate.exceptions import WeaviateStartUpError
from weaviate.util import generate_uuid5

from datastew.embedding import Vectorizer
Expand Down Expand Up @@ -772,19 +774,33 @@ def _is_port_in_use(port) -> bool:
return s.connect_ex(("localhost", port)) == 0

def _connect_to_memory(self, path: str, http_port: int, grpc_port: int) -> WeaviateClient:
try:
if path is None:
raise ValueError("Path must be provided for disk mode.")
if self._is_port_in_use(http_port) and self._is_port_in_use(grpc_port):
if path is None:
raise ValueError("Path must be provided for disk mode.")

# If there's already a Weaviate instance listening; just connect to it.
if self._is_port_in_use(http_port) and self._is_port_in_use(grpc_port):
try:
return weaviate.connect_to_local(port=http_port, grpc_port=grpc_port, headers=self.headers)
else:
except Exception as e:
raise ConnectionError(f"Failed to initialize Weaviate client (embedded): {e}") from e

last_exc: Exception | None = None
for _ in range(3):
try:
return weaviate.connect_to_embedded(
persistence_data_path=path,
headers=self.headers,
environment_variables={"ENABLE_MODULES": "text2vec-ollama"},
)
except Exception as e:
raise ConnectionError(f"Failed to initialize Weaviate client: {e}")
except WeaviateStartUpError as e:
last_exc = e
time.sleep(2.0)
except Exception as e:
raise ConnectionError(f"Failed to initialize Weaviate client (embedded): {e}") from e

raise ConnectionError(
f"Failed to initalize Weaviate client (embedded) after retries: {last_exc}"
) from last_exc

def _connect_to_remote(self, path: str, port: int) -> WeaviateClient:
try:
Expand Down
Loading