pygpt-net 2.4.41__py3-none-any.whl → 2.4.43__py3-none-any.whl

pygpt_net/provider/core/index/db_sqlite/__init__.py

@@ -6,11 +6,14 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2024.03.06 02:00:00                  #
+# Updated Date: 2024.12.14 22:00:00                  #
 # ================================================== #
 
+from typing import Dict, Any, Optional
+
 from packaging.version import Version
 
+from pygpt_net.item.index import IndexItem
 from pygpt_net.provider.core.index.base import BaseProvider
 from .patch import Patch
 from .storage import Storage
@@ -42,7 +45,10 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.patcher.execute(version)
 
-    def load(self, store_id: str) -> dict:
+    def load(
+            self,
+            store_id: str
+    ) -> Dict[str, IndexItem]:
         """
         Load items for index provider
 
@@ -51,7 +57,12 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.get_items(store_id)
 
-    def append_file(self, store_id: str, idx: str, data: dict) -> int:
+    def append_file(
+            self,
+            store_id: str,
+            idx: str,
+            data: Dict[str, Any]
+    ) -> int:
         """
         Append file to index
 
@@ -62,7 +73,13 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.insert_file(store_id, idx, data)
 
-    def append_ctx_meta(self, store_id: str, idx: str, meta_id: int, doc_id: str) -> int:
+    def append_ctx_meta(
+            self,
+            store_id: str,
+            idx: str,
+            meta_id: int,
+            doc_id: str
+    ) -> int:
         """
         Append context meta to index
 
@@ -74,7 +91,12 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.insert_ctx_meta(store_id, idx, meta_id, doc_id)
 
-    def append_external(self, store_id: str, idx: str, data: dict) -> int:
+    def append_external(
+            self,
+            store_id: str,
+            idx: str,
+            data: Dict[str, Any]
+    ) -> int:
         """
         Append external data to index
 
@@ -85,7 +107,12 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.insert_external(store_id, idx, data)
 
-    def is_meta_indexed(self, store_id: str, idx: str, meta_id: int) -> bool:
+    def is_meta_indexed(
+            self,
+            store_id: str,
+            idx: str,
+            meta_id: int
+    ) -> bool:
         """
         Check if context meta is indexed
 
@@ -96,7 +123,12 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.is_meta_indexed(store_id, idx, meta_id)
 
-    def is_file_indexed(self, store_id: str, idx: str, file_id: str) -> bool:
+    def is_file_indexed(
+            self,
+            store_id: str,
+            idx: str,
+            file_id: str
+    ) -> bool:
         """
         Check if file is indexed
 
@@ -107,7 +139,13 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.is_file_indexed(store_id, idx, file_id)
 
-    def is_external_indexed(self, store_id: str, idx: str, content: str, type: str) -> bool:
+    def is_external_indexed(
+            self,
+            store_id: str,
+            idx: str,
+            content: str,
+            type: str
+    ) -> bool:
         """
         Check if external data is indexed
 
@@ -119,7 +157,12 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.is_external_indexed(store_id, idx, content, type)
 
-    def get_meta_doc_id(self, store_id: str, idx: str, meta_id: int) -> str:
+    def get_meta_doc_id(
+            self,
+            store_id: str,
+            idx: str,
+            meta_id: int
+    ) -> str:
         """
         Get indexed document id by meta id
 
@@ -130,7 +173,12 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.get_meta_doc_id(store_id, idx, meta_id)
 
-    def get_file_doc_id(self, store_id: str, idx: str, file_id: str) -> str:
+    def get_file_doc_id(
+            self,
+            store_id: str,
+            idx: str,
+            file_id: str
+    ) -> str:
         """
         Get indexed document id by file id
 
@@ -141,7 +189,13 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.get_file_doc_id(store_id, idx, file_id)
 
-    def get_external_doc_id(self, store_id: str, idx: str, content: str, type: str) -> str:
+    def get_external_doc_id(
+            self,
+            store_id: str,
+            idx: str,
+            content: str,
+            type: str
+    ) -> str:
         """
         Get indexed document id by file id
 
@@ -153,7 +207,12 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.get_external_doc_id(store_id, idx, content, type)
 
-    def update_file(self, id: int, doc_id: str, ts: int) -> bool:
+    def update_file(
+            self,
+            id: int,
+            doc_id: str,
+            ts: int
+    ) -> bool:
         """
         Update indexed timestamp of indexed file
 
@@ -163,7 +222,11 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.update_file(id, doc_id, ts)
 
-    def update_ctx_meta(self, meta_id: int, doc_id: str) -> bool:
+    def update_ctx_meta(
+            self,
+            meta_id: int,
+            doc_id: str
+    ) -> bool:
         """
         Update indexed timestamp of indexed context meta
 
@@ -172,7 +235,13 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.update_ctx_meta(meta_id, doc_id)
 
-    def update_external(self, content: str, type: str, doc_id: str, ts: int) -> bool:
+    def update_external(
+            self,
+            content: str,
+            type: str,
+            doc_id: str,
+            ts: int
+    ) -> bool:
         """
         Update indexed timestamp of indexed external
 
@@ -183,7 +252,12 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.update_external(content, type, doc_id, ts)
 
-    def remove_file(self, store_id: str, idx: str, doc_id: str):
+    def remove_file(
+            self,
+            store_id: str,
+            idx: str,
+            doc_id: str
+    ):
         """
         Remove file from index
 
@@ -193,7 +267,12 @@ class DbSqliteProvider(BaseProvider):
         """
         self.storage.remove_file(store_id, idx, doc_id)
 
-    def remove_ctx_meta(self, store_id: str, idx: str, meta_id: str):
+    def remove_ctx_meta(
+            self,
+            store_id: str,
+            idx: str,
+            meta_id: str
+    ):
         """
         Remove file from index
 
@@ -203,7 +282,12 @@ class DbSqliteProvider(BaseProvider):
         """
         self.storage.remove_ctx_meta(store_id, idx, meta_id)
 
-    def remove_external(self, store_id: str, idx: str, doc_id: str):
+    def remove_external(
+            self,
+            store_id: str,
+            idx: str,
+            doc_id: str
+    ):
         """
         Remove external from index
 
@@ -213,7 +297,11 @@ class DbSqliteProvider(BaseProvider):
         """
         self.storage.remove_external(store_id, idx, doc_id)
 
-    def truncate(self, store_id: str, idx: str) -> bool:
+    def truncate(
+            self,
+            store_id: str,
+            idx: str
+    ) -> bool:
         """
         Truncate idx (remove all items)
 
@@ -223,7 +311,11 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.truncate_all(store_id, idx)
 
-    def truncate_all(self, store_id: str = None, idx: str = None) -> bool:
+    def truncate_all(
+            self,
+            store_id: Optional[str] = None,
+            idx: Optional[str] = None
+    ) -> bool:
         """
         Truncate idx (remove all items) - all store and indexes
 
@@ -233,7 +325,11 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.truncate_all(store_id, idx)
 
-    def truncate_files(self, store_id: str = None, idx: str = None) -> bool:
+    def truncate_files(
+            self,
+            store_id: Optional[str] = None,
+            idx: Optional[str] = None
+    ) -> bool:
         """
         Truncate files idx (remove all items) - all store and indexes
 
@@ -243,7 +339,11 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.truncate_files(store_id, idx)
 
-    def truncate_ctx(self, store_id: str = None, idx: str = None) -> bool:
+    def truncate_ctx(
+            self,
+            store_id: Optional[str] = None,
+            idx: Optional[str] = None
+    ) -> bool:
         """
         Truncate context meta idx (remove all items) - all store and indexes
 
@@ -253,7 +353,11 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.truncate_ctx(store_id, idx)
 
-    def truncate_external(self, store_id: str = None, idx: str = None) -> bool:
+    def truncate_external(
+            self,
+            store_id: Optional[str] = None,
+            idx: Optional[str] = None
+    ) -> bool:
         """
         Truncate external idx (remove all items) - all store and indexes
 
@@ -263,7 +367,10 @@ class DbSqliteProvider(BaseProvider):
         """
         return self.storage.truncate_external(store_id, idx)
 
-    def get_counters(self, type: str) -> dict:
+    def get_counters(
+            self,
+            type: str
+    ) -> Dict[str, Dict[str, int]]:
         """
         Get counters (stats, count items by type [file, ctx, external])
 

pygpt_net/provider/core/index/db_sqlite/storage.py

@@ -6,13 +6,15 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2024.03.06 02:00:00                  #
+# Updated Date: 2024.12.14 22:00:00                  #
 # ================================================== #
 
 import uuid
 import time
+from typing import Dict, Optional
 
 from sqlalchemy import text
+from traitlets import Any
 
 from pygpt_net.item.index import IndexItem
 from .utils import unpack_file_item
@@ -35,7 +37,10 @@ class Storage:
         """
         self.window = window
 
-    def get_items(self, store_id: str) -> dict:
+    def get_items(
+            self,
+            store_id: str
+    ) -> Dict[str, IndexItem]:
         """
         Return idx items by store ID
 
@@ -65,7 +70,12 @@ class Storage:
 
         return indexes
 
-    def insert_file(self, store_id: str, idx: str, data: dict) -> int:
+    def insert_file(
+            self,
+            store_id: str,
+            idx: str,
+            data: Dict[str, Any]
+    ) -> int:
         """
         Insert file to index
 
@@ -114,7 +124,13 @@ class Storage:
 
         return id
 
-    def insert_ctx_meta(self, store_id: str, idx: str, meta_id: int, doc_id: str) -> int:
+    def insert_ctx_meta(
+            self,
+            store_id: str,
+            idx: str,
+            meta_id: int,
+            doc_id: str
+    ) -> int:
         """
         Insert ctx meta to index
 
@@ -162,7 +178,12 @@ class Storage:
 
         return id
 
-    def insert_external(self, store_id: str, idx: str, data: dict) -> int:
+    def insert_external(
+            self,
+            store_id: str,
+            idx: str,
+            data: Dict[str, Any]
+    ) -> int:
         """
         Insert external data to index
 
@@ -211,7 +232,12 @@ class Storage:
 
         return id
 
-    def is_meta_indexed(self, store_id: str, idx: str, meta_id: int) -> bool:
+    def is_meta_indexed(
+            self,
+            store_id: str,
+            idx: str,
+            meta_id: int
+    ) -> bool:
         """
         Check if context meta is indexed
 
@@ -238,7 +264,12 @@ class Storage:
             data = row._asdict()
             return int(data['count']) > 0
 
-    def is_file_indexed(self, store_id: str, idx: str, file_id: str) -> bool:
+    def is_file_indexed(
+            self,
+            store_id: str,
+            idx: str,
+            file_id: str
+    ) -> bool:
         """
         Check if file is indexed
 
@@ -265,7 +296,13 @@ class Storage:
             data = row._asdict()
             return int(data['count']) > 0
 
-    def is_external_indexed(self, store_id: str, idx: str, content: str, type: str) -> bool:
+    def is_external_indexed(
+            self,
+            store_id: str,
+            idx: str,
+            content: str,
+            type: str
+    ) -> bool:
         """
         Check if external is indexed
 
@@ -295,7 +332,12 @@ class Storage:
             data = row._asdict()
             return int(data['count']) > 0
 
-    def get_meta_doc_id(self, store_id: str, idx: str, meta_id: int) -> str:
+    def get_meta_doc_id(
+            self,
+            store_id: str,
+            idx: str,
+            meta_id: int
+    ) -> str:
         """
         Get indexed document id by meta id
 
@@ -322,7 +364,12 @@ class Storage:
             data = row._asdict()
             return data['doc_id']
 
-    def get_file_doc_id(self, store_id: str, idx: str, file_id: str) -> str:
+    def get_file_doc_id(
+            self,
+            store_id: str,
+            idx: str,
+            file_id: str
+    ) -> str:
         """
         Get indexed document id by file id
 
@@ -349,7 +396,13 @@ class Storage:
             data = row._asdict()
             return data['doc_id']
 
-    def get_external_doc_id(self, store_id: str, idx: str, content: str, type: str) -> str:
+    def get_external_doc_id(
+            self,
+            store_id: str,
+            idx: str,
+            content: str,
+            type: str
+    ) -> str:
         """
         Get indexed document id by external
 
@@ -379,7 +432,12 @@ class Storage:
             data = row._asdict()
             return data['doc_id']
 
-    def update_file(self, id: int, doc_id: str, ts: int) -> bool:
+    def update_file(
+            self,
+            id: int,
+            doc_id: str,
+            ts: int
+    ) -> bool:
         """
         Update timestamp of file in index
 
@@ -403,7 +461,11 @@ class Storage:
             conn.execute(stmt)
         return True
 
-    def update_ctx_meta(self, meta_id: int, doc_id: str) -> bool:
+    def update_ctx_meta(
+            self,
+            meta_id: int,
+            doc_id: str
+    ) -> bool:
         """
         Update timestamp of ctx meta in index
 
@@ -427,7 +489,13 @@ class Storage:
             conn.execute(stmt)
         return True
 
-    def update_external(self, content: str, type: str, doc_id: str, ts: int) -> bool:
+    def update_external(
+            self,
+            content: str,
+            type: str,
+            doc_id: str,
+            ts: int
+    ) -> bool:
         """
         Update timestamp of external data in index
 
@@ -454,7 +522,12 @@ class Storage:
             conn.execute(stmt)
         return True
 
-    def remove_file(self, store_id: str, idx: str, doc_id: str):
+    def remove_file(
+            self,
+            store_id: str,
+            idx: str,
+            doc_id: str
+    ):
         """
         Remove file from index
 
@@ -471,7 +544,12 @@ class Storage:
                     doc_id=doc_id,
                 ))
 
-    def remove_ctx_meta(self, store_id: str, idx: str, meta_id: str):
+    def remove_ctx_meta(
+            self,
+            store_id: str,
+            idx: str,
+            meta_id: str
+    ):
         """
         Remove file from index
 
@@ -488,7 +566,12 @@ class Storage:
                     meta_id=meta_id,
                 ))
 
-    def remove_external(self, store_id: str, idx: str, doc_id: str):
+    def remove_external(
+            self,
+            store_id: str,
+            idx: str,
+            doc_id: str
+    ):
         """
         Remove file from index
 
@@ -505,7 +588,11 @@ class Storage:
                     doc_id=doc_id,
                 ))
 
-    def truncate_all(self, store_id: str = None, idx: str = None) -> bool:
+    def truncate_all(
+            self,
+            store_id: Optional[str] = None,
+            idx: Optional[str] = None
+    ) -> bool:
         """
         Truncate all idx tables in database (all stores)
 
@@ -518,7 +605,11 @@ class Storage:
         self.truncate_external(store_id, idx)
         return True
 
-    def truncate_files(self, store_id: str = None, idx: str = None) -> bool:
+    def truncate_files(
+            self,
+            store_id: Optional[str] = None,
+            idx: Optional[str] = None
+    ) -> bool:
         """
         Truncate files table in database
 
@@ -528,7 +619,11 @@ class Storage:
         """
         return self.truncate_by_db_table("idx_file", store_id, idx)
 
-    def truncate_ctx(self, store_id: str = None, idx: str = None) -> bool:
+    def truncate_ctx(
+            self,
+            store_id: Optional[str] = None,
+            idx: Optional[str] = None
+    ) -> bool:
         """
         Truncate context table in database
 
@@ -538,7 +633,11 @@ class Storage:
         """
         return self.truncate_by_db_table("idx_ctx", store_id, idx)
 
-    def truncate_external(self, store_id: str = None, idx: str = None) -> bool:
+    def truncate_external(
+            self,
+            store_id: Optional[str] = None,
+            idx: Optional[str] = None
+    ) -> bool:
         """
         Truncate external table in database
 
@@ -548,7 +647,12 @@ class Storage:
         """
         return self.truncate_by_db_table("idx_external", store_id, idx)
 
-    def truncate_by_db_table(self, tbl: str, store_id: str = None, idx: str = None) -> bool:
+    def truncate_by_db_table(
+            self,
+            tbl: str,
+            store_id: Optional[str] = None,
+            idx: Optional[str] = None
+    ) -> bool:
         """
         Truncate external table in database
 
@@ -574,7 +678,10 @@ class Storage:
                 text(query).bindparams(**params))
         return True
 
-    def get_counters(self, type: str) -> dict:
+    def get_counters(
+            self,
+            type: str
+    ) -> Dict[str, Dict[str, int]]:
         """
         Get counters (stats, count items by type [file, ctx, external])
 

pygpt_net/provider/core/index/db_sqlite/utils.py

@@ -6,13 +6,15 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2024.03.06 02:00:00                  #
+# Updated Date: 2024.12.14 22:00:00                  #
 # ================================================== #
 
+from typing import Tuple, Dict, Any
+
 from pygpt_net.utils import unpack_var
 
 
-def unpack_file_item(row: dict) -> (str, dict):
+def unpack_file_item(row: Dict[str, Any]) -> Tuple[str, dict]:
     """
     Unpack item from DB row
 

pygpt_net/provider/core/mode/base.py

@@ -6,9 +6,11 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2023.12.31 04:00:00                  #
+# Updated Date: 2024.12.14 22:00:00                  #
 # ================================================== #
 
+from typing import Dict
+
 from pygpt_net.item.mode import ModeItem
 
 
@@ -24,10 +26,10 @@ class BaseProvider:
     def create(self, mode: ModeItem) -> str:
         pass
 
-    def load(self) -> dict:
+    def load(self) -> Dict[str, str]:
         pass
 
-    def save(self, items: dict):
+    def save(self, items: Dict[str, str]):
         pass
 
     def remove(self, id: str):