Merge pull request #7 from video-db/ar/add-keyword-search

Ar/add keyword search
video-db · Feb 22, 2024 · 7588bc4 · 7588bc4
2 parents 6646cd8 + 17a195f
commit 7588bc4
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 5 deletions.
diff --git a/videodb/__init__.py b/videodb/__init__.py
@@ -5,7 +5,7 @@
 
 from typing import Optional
 from videodb._utils._video import play_stream
-from videodb._constants import VIDEO_DB_API, MediaType
+from videodb._constants import VIDEO_DB_API, MediaType, SearchType
 from videodb.client import Connection
 from videodb.exceptions import (
     VideodbError,
@@ -16,7 +16,7 @@
 
 logger: logging.Logger = logging.getLogger("videodb")
 
-__version__ = "0.0.3"
+__version__ = "0.0.4"
 __author__ = "videodb"
 
 __all__ = [
@@ -26,6 +26,7 @@
     "SearchError",
     "play_stream",
     "MediaType",
+    "SearchType",
 ]
 
 

diff --git a/videodb/_constants.py b/videodb/_constants.py
@@ -11,6 +11,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: