Skip to content

Commit b8e65eb

Browse files
committed
added doc-id for inserting vectors. Right now its random, must fix later
1 parent 398b64e commit b8e65eb

File tree

2 files changed

+21
-26
lines changed

2 files changed

+21
-26
lines changed

engine/base_client/search.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import functools
22
import random
3+
import itertools
34
import time
45
from multiprocessing import Process, Queue
56
from typing import Iterable, List, Optional, Tuple
@@ -19,6 +20,7 @@
1920

2021

2122
class BaseSearcher:
23+
_doc_id_counter = itertools.count(1)
2224
MP_CONTEXT = None
2325

2426
def __init__(self, host, connection_params, search_params):
@@ -43,9 +45,7 @@ def search_one(
4345
raise NotImplementedError()
4446

4547
@classmethod
46-
def insert_one(
47-
cls, vector: List[float], meta_conditions
48-
) -> List[Tuple[int, float]]:
48+
def insert_one(cls, doc_id: int, vector: List[float], meta_conditions):
4949
raise NotImplementedError()
5050

5151
@classmethod
@@ -70,7 +70,11 @@ def _search_one(cls, query, top: Optional[int] = None):
7070
@classmethod
7171
def _insert_one(cls, query):
7272
start = time.perf_counter()
73-
cls.insert_one(query.vector, query.meta_conditions)
73+
74+
# Generate unique doc_id here
75+
doc_id = next(cls._doc_id_counter)
76+
77+
cls.insert_one(doc_id, query.vector, query.meta_conditions)
7478
end = time.perf_counter()
7579
# No precision metric for inserts, so precision=1.0
7680
return 1.0, end - start

engine/clients/redis/search.py

Lines changed: 13 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -102,28 +102,19 @@ def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]:
102102

103103
return [(int(result.id), float(result.vector_score)) for result in results.docs]
104104

105-
@classmethod
106-
def insert_one(cls, doc_id: int, vector, metadata: dict = None):
107-
"""
108-
Insert a single vector and optional metadata into Redis.
109-
Designed for mixed workload support.
110-
"""
111-
if cls.client is None:
112-
raise RuntimeError("Redis client not initialized")
105+
@classmethod
106+
def insert_one(cls, doc_id: int, vector, meta_conditions):
107+
if cls.client is None:
108+
raise RuntimeError("Redis client not initialized")
113109

114-
# Convert vector to correct binary format
115-
if isinstance(vector, bytes):
116-
vec_param = vector
117-
else:
118-
vec_param = np.array(vector, dtype=cls.np_data_type).tobytes()
110+
if not isinstance(vector, bytes):
111+
vec_param = np.array(vector, dtype=cls.np_data_type).tobytes()
112+
else:
113+
vec_param = vector
119114

120-
# Prepare the document for Redis
121-
doc = {"vector": vec_param}
122-
if metadata:
123-
# Flatten metadata into string fields Redis can store
124-
for k, v in metadata.items():
125-
doc[k] = str(v)
115+
doc = {"vector": vec_param}
116+
if meta_conditions:
117+
for k, v in meta_conditions.items():
118+
doc[k] = str(v)
126119

127-
# Use Redis hash to store the document
128-
# You might use JSON.SET if your index expects JSON
129-
cls.client.hset(f"doc:{doc_id}", mapping=doc)
120+
cls.client.hset(f"doc:{doc_id}", mapping=doc)

0 commit comments

Comments
 (0)