Support to config the cache storage data size (#383)

Signed-off-by: SimFG <bang.fu@zilliz.com>

Author: SimFG

docs/_exts/docgen.py

@@ -114,7 +114,6 @@ def generate(self, lib_name):
             f.write(t)
 
         # Iterate the modules, render the function templates and write rendered output to files
-        print("modules:", modules)
         for module in modules:
             module_name = module[0]
 

docs/requirements.txt

@@ -14,3 +14,4 @@ sphinx_copybutton
 pydata-sphinx-theme
 m2r2
 sphinx_toolbox
+protobuf==3.20.2

gptcache/manager/scalar_data/__init__.py

@@ -14,17 +14,26 @@ def CacheBase(name: str, **kwargs):
     :param name: the name of the cache storage, it is support 'sqlite', 'postgresql', 'mysql', 'mariadb', 'sqlserver' and  'oracle' now.
     :type name: str
     :param sql_url: the url of the sql database for cache, such as '<db_type>+<db_driver>://<username>:<password>@<host>:<port>/<database>',
-                    and the default value is related to the `cache_store` parameter,
-                    'sqlite:///./sqlite.db' for 'sqlite',
-                    'duckdb:///./duck.db' for 'duckdb',
-                    'postgresql+psycopg2://postgres:123456@127.0.0.1:5432/postgres' for 'postgresql',
-                    'mysql+pymysql://root:123456@127.0.0.1:3306/mysql' for 'mysql',
-                    'mariadb+pymysql://root:123456@127.0.0.1:3307/mysql' for 'mariadb',
-                    'mssql+pyodbc://sa:Strongpsw_123@127.0.0.1:1434/msdb?driver=ODBC+Driver+17+for+SQL+Server' for 'sqlserver',
-                    'oracle+cx_oracle://oracle:123456@127.0.0.1:1521/?service_name=helowin&encoding=UTF-8&nencoding=UTF-8' for 'oracle'.
+        and the default value is related to the `cache_store` parameter,
+
+        - 'sqlite:///./sqlite.db' for 'sqlite',
+        - 'duckdb:///./duck.db' for 'duckdb',
+        - 'postgresql+psycopg2://postgres:123456@127.0.0.1:5432/postgres' for 'postgresql',
+        - 'mysql+pymysql://root:123456@127.0.0.1:3306/mysql' for 'mysql',
+        - 'mariadb+pymysql://root:123456@127.0.0.1:3307/mysql' for 'mariadb',
+        - 'mssql+pyodbc://sa:Strongpsw_123@127.0.0.1:1434/msdb?driver=ODBC+Driver+17+for+SQL+Server' for 'sqlserver',
+        - 'oracle+cx_oracle://oracle:123456@127.0.0.1:1521/?service_name=helowin&encoding=UTF-8&nencoding=UTF-8' for 'oracle'.
     :type sql_url: str
     :param table_name: the table name for sql database, defaults to 'gptcache'.
     :type table_name: str
+    :param table_len_config: the table length config for sql database, defaults to {}. the key includes:
+
+        - 'question_question': the question column size in the question table, default to 3000.
+        - 'answer_answer': the answer column size in the answer table, default to 3000.
+        - 'session_id': the session id column size in the session table, default to 1000.
+        - 'dep_name': the name column size in the dep table, default to 1000.
+        - 'dep_data': the data column size in the dep table, default to 3000.
+    :type table_len_config: dict
 
     :return: CacheStorage.
 

gptcache/manager/scalar_data/manager.py

@@ -27,13 +27,27 @@ def __init__(self):
 
     @staticmethod
     def get(name, **kwargs):
-        if name in ["sqlite", "postgresql", "mysql", "mariadb", "sqlserver", "oracle"]:
+        if name in [
+            "sqlite",
+            "duckdb",
+            "postgresql",
+            "mysql",
+            "mariadb",
+            "sqlserver",
+            "oracle",
+        ]:
             from gptcache.manager.scalar_data.sql_storage import SQLStorage
 
             sql_url = kwargs.get("sql_url", SQL_URL[name])
             table_name = kwargs.get("table_name", TABLE_NAME)
+            table_len_config = kwargs.get("table_len_config", {})
             import_sql_client(name)
-            cache_base = SQLStorage(db_type=name, url=sql_url, table_name=table_name)
+            cache_base = SQLStorage(
+                db_type=name,
+                url=sql_url,
+                table_name=table_name,
+                table_len_config=table_len_config,
+            )
         else:
             raise NotFoundError("cache store", name)
         return cache_base

gptcache/manager/scalar_data/sql_storage.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import List, Optional
+from typing import List, Optional, Dict
 
 import numpy as np
 
@@ -24,8 +24,22 @@
 from sqlalchemy.orm import sessionmaker  # pylint: disable=C0413
 from sqlalchemy.ext.declarative import declarative_base  # pylint: disable=C0413
 
+DEFAULT_LEN_DOCT = {
+    "question_question": 3000,
+    "answer_answer": 3000,
+    "session_id": 1000,
+    "dep_name": 1000,
+    "dep_data": 3000,
+}
 
-def get_models(table_prefix, db_type):
+
+def _get_table_len(config: Dict, column_alias: str) -> int:
+    if config and column_alias in config and config[column_alias] > 0:
+        return config[column_alias]
+    return DEFAULT_LEN_DOCT.get(column_alias, 1000)
+
+
+def get_models(table_prefix, db_type, table_len_config):
     DynamicBase = declarative_base(class_registry={})  # pylint: disable=C0103
 
     class QuestionTable(DynamicBase):
@@ -41,7 +55,10 @@ class QuestionTable(DynamicBase):
             id = Column(Integer, question_id_seq, primary_key=True, autoincrement=True)
         else:
             id = Column(Integer, primary_key=True, autoincrement=True)
-        question = Column(String(1000), nullable=False)
+        question = Column(
+            String(_get_table_len(table_len_config, "question_question")),
+            nullable=False,
+        )
         create_on = Column(DateTime, default=datetime.now)
         last_access = Column(DateTime, default=datetime.now)
         embedding_data = Column(LargeBinary, nullable=True)
@@ -61,7 +78,9 @@ class AnswerTable(DynamicBase):
         else:
             id = Column(Integer, primary_key=True, autoincrement=True)
         question_id = Column(Integer, nullable=False)
-        answer = Column(String(2000), nullable=False)
+        answer = Column(
+            String(_get_table_len(table_len_config, "answer_answer")), nullable=False
+        )
         answer_type = Column(Integer, nullable=False)
 
     class SessionTable(DynamicBase):
@@ -83,8 +102,13 @@ class SessionTable(DynamicBase):
         else:
             id = Column(Integer, primary_key=True, autoincrement=True)
         question_id = Column(Integer, nullable=False)
-        session_id = Column(String(500), nullable=False)
-        session_question = Column(String(2000), nullable=False)
+        session_id = Column(
+            String(_get_table_len(table_len_config, "session_id")), nullable=False
+        )
+        session_question = Column(
+            String(_get_table_len(table_len_config, "question_question")),
+            nullable=False,
+        )
 
     class QuestionDepTable(DynamicBase):
         """
@@ -102,8 +126,12 @@ class QuestionDepTable(DynamicBase):
         else:
             id = Column(Integer, primary_key=True, autoincrement=True)
         question_id = Column(Integer, nullable=False)
-        dep_name = Column(String(255), nullable=False)
-        dep_data = Column(String(1000), nullable=False)
+        dep_name = Column(
+            String(_get_table_len(table_len_config, "dep_name")), nullable=False
+        )
+        dep_data = Column(
+            String(_get_table_len(table_len_config, "dep_data")), nullable=False
+        )
         dep_type = Column(Integer, nullable=False)
 
     return QuestionTable, AnswerTable, QuestionDepTable, SessionTable
@@ -134,10 +162,13 @@ def __init__(
         db_type: str = "sqlite",
         url: str = "sqlite:///./sqlite.db",
         table_name: str = "gptcache",
+        table_len_config=None,
     ):
+        if table_len_config is None:
+            table_len_config = {}
         self._url = url
         self._ques, self._answer, self._ques_dep, self._session = get_models(
-            table_name, db_type
+            table_name, db_type, table_len_config
         )
         self._engine = create_engine(self._url)
         self.Session = sessionmaker(bind=self._engine)  # pylint: disable=invalid-name

tests/requirements.txt

@@ -20,5 +20,5 @@ mock
 pexpect
 spacy
 safetensors
-protobuf==3.20.0
 typing_extensions<4.6.0
+protobuf==3.20.0

tests/unit_tests/manager/test_sql_scalar.py

@@ -10,7 +10,6 @@
 
 
 class TestSQLStore(unittest.TestCase):
-
     def test_sqlite(self):
         self._inner_test_normal("sqlite")
         self._inner_test_with_deps("sqlite")
@@ -21,26 +20,38 @@ def test_duckdb(self):
         self._inner_test_with_deps("duckdb")
 
     def _inner_test_normal(self, db_name: str):
-        with TemporaryDirectory(dir='./') as root:
-            db_path = Path(root) / f'{db_name}1.db'
-            db = SQLStorage(db_type=db_name, url=f"{db_name}:///" + str(db_path))
+        with TemporaryDirectory(dir="./") as root:
+            db_path = Path(root) / f"{db_name}1.db"
+            db = SQLStorage(
+                db_type=db_name,
+                url=f"{db_name}:///" + str(db_path),
+                table_len_config={"question_question": 500},
+            )
             db.create()
             data = []
             for i in range(1, 10):
-                data.append(CacheData('question_' + str(i), ['answer_' + str(i)] * i, np.random.rand(5)))
+                data.append(
+                    CacheData(
+                        "question_" + str(i),
+                        ["answer_" + str(i)] * i,
+                        np.random.rand(5),
+                    )
+                )
 
             db.batch_insert(data)
             data = db.get_data_by_id(1)
-            self.assertEqual(data.question, 'question_1')
-            self.assertEqual(data.answers[0].answer, 'answer_1')
+            self.assertEqual(data.question, "question_1")
+            self.assertEqual(data.answers[0].answer, "answer_1")
             data = db.get_data_by_id(2)
-            self.assertEqual(data.question, 'question_2')
-            self.assertEqual(data.answers[0].answer, 'answer_2')
-            self.assertEqual(data.answers[1].answer, 'answer_2')
-            q_id = db.batch_insert([CacheData('question_single', 'answer_singel', np.random.rand(5))])[0]
+            self.assertEqual(data.question, "question_2")
+            self.assertEqual(data.answers[0].answer, "answer_2")
+            self.assertEqual(data.answers[1].answer, "answer_2")
+            q_id = db.batch_insert(
+                [CacheData("question_single", "answer_singel", np.random.rand(5))]
+            )[0]
             data = db.get_data_by_id(q_id)
-            self.assertEqual(data.question, 'question_single')
-            self.assertEqual(data.answers[0].answer, 'answer_singel')
+            self.assertEqual(data.question, "question_single")
+            self.assertEqual(data.answers[0].answer, "answer_singel")
 
             # test deleted
             self.assertEqual(len(db.get_ids(True)), 0)
@@ -54,26 +65,42 @@ def _inner_test_normal(self, db_name: str):
             self.assertEqual(db.count(is_all=True), 7)
 
     def _inner_test_with_deps(self, db_name: str):
-        with TemporaryDirectory(dir='./') as root:
-            db_path = Path(root) / f'{db_name}2.db'
+        with TemporaryDirectory(dir="./") as root:
+            db_path = Path(root) / f"{db_name}2.db"
             db = SQLStorage(db_type=db_name, url=f"{db_name}:///" + str(db_path))
             db.create()
-            data_id = db.batch_insert([
-                CacheData(
-                    Question.from_dict({
-                        "content": "test_question",
-                        "deps": [
-                            {"name": "text", "data": "how many people in this picture", "dep_type": 0},
-                            {"name": "image", "data": "object_name", "dep_type": 1}
-                        ]
-                    }),
-                    'test_answer', np.random.rand(5))
-            ])[0]
+            data_id = db.batch_insert(
+                [
+                    CacheData(
+                        Question.from_dict(
+                            {
+                                "content": "test_question",
+                                "deps": [
+                                    {
+                                        "name": "text",
+                                        "data": "how many people in this picture",
+                                        "dep_type": 0,
+                                    },
+                                    {
+                                        "name": "image",
+                                        "data": "object_name",
+                                        "dep_type": 1,
+                                    },
+                                ],
+                            }
+                        ),
+                        "test_answer",
+                        np.random.rand(5),
+                    )
+                ]
+            )[0]
 
             ret = db.get_data_by_id(data_id)
             self.assertEqual(ret.question.content, "test_question")
             self.assertEqual(ret.question.deps[0].name, "text")
-            self.assertEqual(ret.question.deps[0].data, "how many people in this picture")
+            self.assertEqual(
+                ret.question.deps[0].data, "how many people in this picture"
+            )
             self.assertEqual(ret.question.deps[0].dep_type, 0)
             self.assertEqual(ret.question.deps[1].name, "image")
             self.assertEqual(ret.question.deps[1].data, "object_name")