class VideoDBRetriever(BaseRetriever):
def __init__(
self,
api_key: Optional[str] = None,
collection: Optional[str] = "default",
video: Optional[str] = None,
score_threshold: Optional[float] = 0.2,
result_threshold: Optional[int] = 5,
search_type: Optional[str] = SearchType.semantic,
base_url: Optional[str] = None,
callback_manager: Optional[CallbackManager] = None,
) -> None:
"""Creates a new VideoDB Retriever."""
if api_key is None:
api_key = os.environ.get("VIDEO_DB_API_KEY")
if api_key is None:
raise Exception(
"No API key provided. Set an API key either as an environment variable (VIDEO_DB_API_KEY) or pass it as an argument."
)
self._api_key = api_key
self._base_url = base_url
self.video = video
self.collection = collection
self.score_threshold = score_threshold
self.result_threshold = result_threshold
self.search_type = search_type
super().__init__(callback_manager)
def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
"""Retrieve."""
kwargs = {"api_key": self._api_key}
if self._base_url is not None:
kwargs["base_url"] = self._base_url
conn = connect(**kwargs)
if self.video:
coll = conn.get_collection(self.collection)
video = coll.get_video(self.video)
search_res = video.search(
query_bundle.query_str,
search_type=self.search_type,
score_threshold=self.score_threshold,
result_threshold=self.result_threshold,
)
else:
coll = conn.get_collection(self.collection)
search_res = coll.search(
query_bundle.query_str,
search_type=self.search_type,
score_threshold=self.score_threshold,
result_threshold=self.result_threshold,
)
nodes = []
collection_id = search_res.collection_id
for shot in search_res.get_shots():
score = shot.search_score
textnode = TextNode(
text=shot.text,
metadata={
"collection_id": collection_id,
"video_id": shot.video_id,
"length": shot.video_length,
"title": shot.video_title,
"start": shot.start,
"end": shot.end,
},
)
nodes.append(NodeWithScore(node=textnode, score=score))
return nodes