diff --git a/videodb/_constants.py b/videodb/_constants.py index 0e80158..c8d920d 100644 --- a/videodb/_constants.py +++ b/videodb/_constants.py @@ -17,11 +17,13 @@ class SearchType: keyword = "keyword" scene = "scene" llm = "llm" + custom = "custom" class IndexType: spoken_word = "spoken_word" scene = "scene" + hybrid = "hybrid" class SceneExtractionType: @@ -36,6 +38,9 @@ class Workflows: class SemanticSearchDefaultValues: result_threshold = 5 score_threshold = 0.2 + rerank = False + rerank_params = {} + sort_docs_on = None class Segmenter: @@ -58,6 +63,7 @@ class ApiPath: index = "index" search = "search" compile = "compile" + clip = "clip" workflow = "workflow" timeline = "timeline" delete = "delete" diff --git a/videodb/collection.py b/videodb/collection.py index e941cf4..4e2dcdb 100644 --- a/videodb/collection.py +++ b/videodb/collection.py @@ -8,6 +8,7 @@ ApiPath, IndexType, SearchType, + SemanticSearchDefaultValues, ) from videodb.video import Video from videodb.audio import Audio @@ -387,6 +388,9 @@ def search( result_threshold: Optional[int] = None, score_threshold: Optional[float] = None, dynamic_score_percentage: Optional[float] = None, + rerank: bool = SemanticSearchDefaultValues.rerank, + rerank_params: dict = SemanticSearchDefaultValues.rerank_params, + sort_docs_on: str = SemanticSearchDefaultValues.sort_docs_on, filter: List[Dict[str, Any]] = [], ) -> SearchResult: """Search for a query in the collection. @@ -397,6 +401,9 @@ def search( :param int result_threshold: Number of results to return (optional) :param float score_threshold: Threshold score for the search (optional) :param float dynamic_score_percentage: Percentage of dynamic score to consider (optional) + :param bool rerank: Rerank search results (optional) + :param dict rerank_params: Parameters for reranking (optional) + :param str sort_docs_on: Parameter to specify what metric to sort the docs of video on :raise SearchError: If the search fails :return: :class:`SearchResult ` object :rtype: :class:`videodb.search.SearchResult` @@ -410,7 +417,10 @@ def search( result_threshold=result_threshold, score_threshold=score_threshold, dynamic_score_percentage=dynamic_score_percentage, + rerank=rerank, + rerank_params=rerank_params, filter=filter, + sort_docs_on=sort_docs_on ) def search_title(self, query) -> List[Video]: diff --git a/videodb/search.py b/videodb/search.py index ba557be..7750a57 100644 --- a/videodb/search.py +++ b/videodb/search.py @@ -29,6 +29,7 @@ def __init__(self, _connection, **kwargs): self.player_url = None self.collection_id = "default" self._results = kwargs.get("results", []) + self.metadata = kwargs.get("meta", None) self._format_results() def _format_results(self): @@ -45,6 +46,9 @@ def _format_results(self): doc.get("end"), doc.get("text"), doc.get("score"), + scene_index_id=doc.get("scene_index_id"), + scene_index_name=doc.get("scene_index_name"), + metadata=doc.get("metadata"), ) ) @@ -123,6 +127,8 @@ def search_inside_video( result_threshold: Optional[int] = None, score_threshold: Optional[float] = None, dynamic_score_percentage: Optional[float] = None, + rerank: bool = SemanticSearchDefaultValues.rerank, + rerank_params: dict = SemanticSearchDefaultValues.rerank_params, **kwargs, ): search_data = self._connection.post( @@ -132,10 +138,14 @@ def search_inside_video( "index_type": index_type, "query": query, "score_threshold": score_threshold - or SemanticSearchDefaultValues.score_threshold, + if score_threshold is not None + else SemanticSearchDefaultValues.score_threshold, "result_threshold": result_threshold - or SemanticSearchDefaultValues.result_threshold, + if result_threshold is not None + else SemanticSearchDefaultValues.result_threshold, "dynamic_score_percentage": dynamic_score_percentage, + "rerank": rerank, + "rerank_params": rerank_params, **kwargs, }, ) @@ -150,6 +160,9 @@ def search_inside_collection( result_threshold: Optional[int] = None, score_threshold: Optional[float] = None, dynamic_score_percentage: Optional[float] = None, + rerank: bool = SemanticSearchDefaultValues.rerank, + rerank_params: dict = SemanticSearchDefaultValues.rerank_params, + sort_docs_on: str = SemanticSearchDefaultValues.sort_docs_on, **kwargs, ): search_data = self._connection.post( @@ -159,10 +172,15 @@ def search_inside_collection( "index_type": index_type, "query": query, "score_threshold": score_threshold - or SemanticSearchDefaultValues.score_threshold, + if score_threshold is not None + else SemanticSearchDefaultValues.score_threshold, "result_threshold": result_threshold - or SemanticSearchDefaultValues.result_threshold, + if result_threshold is not None + else SemanticSearchDefaultValues.result_threshold, "dynamic_score_percentage": dynamic_score_percentage, + "rerank": rerank, + "rerank_params": rerank_params, + "sort_docs_on": sort_docs_on, **kwargs, }, ) @@ -238,6 +256,7 @@ def search_inside_collection(self, **kwargs): SearchType.semantic: SemanticSearch, SearchType.keyword: KeywordSearch, SearchType.scene: SceneSearch, + SearchType.custom: SemanticSearch, } diff --git a/videodb/shot.py b/videodb/shot.py index c2fadcb..2fee6d4 100644 --- a/videodb/shot.py +++ b/videodb/shot.py @@ -19,6 +19,9 @@ class Shot: :ivar int search_score: Search relevance score :ivar str stream_url: URL to stream the shot :ivar str player_url: URL to play the shot in a player + :ivar Optional[str] scene_index_id: ID of the scene index for scene search results + :ivar Optional[str] scene_index_name: Name of the scene index for scene search results + :ivar Optional[dict] metadata: Additional metadata for the shot """ def __init__( @@ -31,6 +34,10 @@ def __init__( end: float, text: Optional[str] = None, search_score: Optional[int] = None, + scene_index_id: Optional[str] = None, + scene_index_name: Optional[str] = None, + metadata: Optional[dict] = None, + ) -> None: self._connection = _connection self.video_id = video_id @@ -40,21 +47,37 @@ def __init__( self.end = end self.text = text self.search_score = search_score + self.scene_index_id = scene_index_id + self.scene_index_name = scene_index_name + self.metadata = metadata self.stream_url = None self.player_url = None def __repr__(self) -> str: - return ( + repr_str = ( f"Shot(" f"video_id={self.video_id}, " f"video_title={self.video_title}, " f"start={self.start}, " f"end={self.end}, " f"text={self.text}, " - f"search_score={self.search_score}, " - f"stream_url={self.stream_url}, " + f"search_score={self.search_score}" + ) + if self.scene_index_id: + repr_str += f", scene_index_id={self.scene_index_id}" + + if self.scene_index_name: + repr_str += f", scene_index_name={self.scene_index_name}" + + if self.metadata: + repr_str += f", metadata={self.metadata}" + + repr_str += ( + f", stream_url={self.stream_url}, " f"player_url={self.player_url})" ) + return repr_str + def __getitem__(self, key): """Get an item from the shot object""" diff --git a/videodb/video.py b/videodb/video.py index eb4da7b..edb7a1e 100644 --- a/videodb/video.py +++ b/videodb/video.py @@ -5,6 +5,7 @@ IndexType, SceneExtractionType, SearchType, + SemanticSearchDefaultValues, Segmenter, SubtitleStyle, Workflows, @@ -69,6 +70,8 @@ def search( result_threshold: Optional[int] = None, score_threshold: Optional[float] = None, dynamic_score_percentage: Optional[float] = None, + rerank: bool = SemanticSearchDefaultValues.rerank, + rerank_params: dict = SemanticSearchDefaultValues.rerank_params, filter: List[Dict[str, Any]] = [], **kwargs, ) -> SearchResult: @@ -80,6 +83,8 @@ def search( :param int result_threshold: (optional) Number of results to return :param float score_threshold: (optional) Threshold score for the search :param float dynamic_score_percentage: (optional) Percentage of dynamic score to consider + :param bool rerank: (optional) Rerank search results + :param dict rerank_params: (optional) Parameters for reranking :raise SearchError: If the search fails :return: :class:`SearchResult ` object :rtype: :class:`videodb.search.SearchResult` @@ -93,6 +98,8 @@ def search( result_threshold=result_threshold, score_threshold=score_threshold, dynamic_score_percentage=dynamic_score_percentage, + rerank=rerank, + rerank_params=rerank_params, filter=filter, **kwargs, ) @@ -277,12 +284,14 @@ def translate_transcript( def index_spoken_words( self, language_code: Optional[str] = None, + segmentation_type: Optional[str] = None, force: bool = False, callback_url: str = None, ) -> None: """Semantic indexing of spoken words in the video. :param str language_code: (optional) Language code of the video + :param str segmentation_type: (optional) Segmentation type used for indexing :param bool force: (optional) Force to index the video :param str callback_url: (optional) URL to receive the callback :raises InvalidRequestError: If the video is already indexed @@ -294,6 +303,7 @@ def index_spoken_words( data={ "index_type": IndexType.spoken_word, "language_code": language_code, + "segmentation_type": segmentation_type, "force": force, "callback_url": callback_url, }, @@ -562,6 +572,31 @@ def add_subtitle(self, style: SubtitleStyle = SubtitleStyle()) -> str: ) return subtitle_data.get("stream_url", None) + def clip( + self, + prompt: str, + content_type: str, + model_name: str, + ) -> str: + """Generate a clip from the video using a prompt. + + :param str prompt: Prompt to generate the clip + :param str content_type: Content type for the clip + :param str model_name: Model name for generation + :return: The stream url of the generated clip + :rtype: str + """ + + clip_data = self._connection.post( + path=f"{ApiPath.video}/{self.id}/{ApiPath.clip}", + data={ + "prompt": prompt, + "content_type": content_type, + "model_name": model_name, + }, + ) + return SearchResult(self._connection, **clip_data) + def insert_video(self, video, timestamp: float) -> str: """Insert a video into another video