Merge branch 'main' of https://github.com/video-db/videodb-python int…

…o ar/add-image-support
video-db · Feb 22, 2024 · 63290e4 · 63290e4
2 parents 30da85a + 7588bc4
commit 63290e4
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 3 deletions.
diff --git a/videodb/__init__.py b/videodb/__init__.py
@@ -8,6 +8,7 @@
 from videodb._constants import (
     VIDEO_DB_API,
     MediaType,
+    SearchType,
     SubtitleAlignment,
     SubtitleBorderStyle,
     SubtitleStyle,
@@ -32,6 +33,7 @@
     "SearchError",
     "play_stream",
     "MediaType",
+    "SearchType",
     "SubtitleAlignment",
     "SubtitleBorderStyle",
     "SubtitleStyle",

diff --git a/videodb/_constants.py b/videodb/_constants.py
@@ -13,6 +13,7 @@ class MediaType:
 
 class SearchType:
     semantic = "semantic"
+    keyword = "keyword"
 
 
 class IndexType:

diff --git a/videodb/search.py b/videodb/search.py
@@ -116,7 +116,7 @@ def search_inside_video(
         search_data = self._connection.post(
             path=f"{ApiPath.video}/{video_id}/{ApiPath.search}",
             data={
-                "type": SearchType.semantic,
+                "index_type": SearchType.semantic,
                 "query": query,
                 "score_threshold": score_threshold
                 or SemanticSearchDefaultValues.score_threshold,
@@ -137,7 +137,7 @@ def search_inside_collection(
         search_data = self._connection.post(
             path=f"{ApiPath.collection}/{collection_id}/{ApiPath.search}",
             data={
-                "type": SearchType.semantic,
+                "index_type": SearchType.semantic,
                 "query": query,
                 "score_threshold": score_threshold
                 or SemanticSearchDefaultValues.score_threshold,
@@ -148,7 +148,35 @@ def search_inside_collection(
         return SearchResult(self._connection, **search_data)
 
 
-search_type = {SearchType.semantic: SemanticSearch}
+class KeywordSearch(Search):
+    def __init__(self, _connection):
+        self._connection = _connection
+
+    def search_inside_video(
+        self,
+        video_id: str,
+        query: str,
+        result_threshold: Optional[int] = None,
+        score_threshold: Optional[int] = None,
+        dynamic_score_percentage: Optional[int] = None,
+        **kwargs,
+    ):
+        search_data = self._connection.post(
+            path=f"{ApiPath.video}/{video_id}/{ApiPath.search}",
+            data={
+                "index_type": SearchType.keyword,
+                "query": query,
+                "score_threshold": score_threshold,
+                "result_threshold": result_threshold,
+            },
+        )
+        return SearchResult(self._connection, **search_data)
+
+    def search_inside_collection(**kwargs):
+        raise NotImplementedError("Keyword search will be implemented in the future")
+
+
+search_type = {SearchType.semantic: SemanticSearch, SearchType.keyword: KeywordSearch}
 
 
 class SearchFactory: