29 lines
885 B
Python
29 lines
885 B
Python
|
from typing import Any, Protocol
|
||
|
|
|
||
|
|
|
||
|
|
class VectorClient(Protocol):
|
||
|
|
def search(self, collection_name: str, data: list[list[float]], limit: int) -> Any: ...
|
||
|
|
|
||
|
|
|
||
|
|
class VectorSearchResult:
|
||
|
|
def __init__(self, hits: Any, degraded: bool = False) -> None:
|
||
|
|
self.hits = hits
|
||
|
|
self.degraded = degraded
|
||
|
|
|
||
|
|
|
||
|
|
class VectorMemoryAdapter:
|
||
|
|
def __init__(self, client: VectorClient, collection: str) -> None:
|
||
|
|
self.client = client
|
||
|
|
self.collection = collection
|
||
|
|
|
||
|
|
def search(self, embedding: list[float], limit: int = 10) -> VectorSearchResult:
|
||
|
|
try:
|
||
|
|
hits = self.client.search(
|
||
|
|
collection_name=self.collection,
|
||
|
|
data=[embedding],
|
||
|
|
limit=max(1, min(limit, 100)),
|
||
|
|
)
|
||
|
|
return VectorSearchResult(hits)
|
||
|
|
except Exception:
|
||
|
|
return VectorSearchResult([], degraded=True)
|