├── example_.env ├── assets └── how-semantic-cache-works.png ├── upstash_semantic_cache ├── __init__.py └── semantic_cache.py ├── .github └── workflows │ └── ci.yml ├── LICENSE ├── examples └── llm.py ├── pyproject.toml ├── .gitignore ├── tests └── test_semantic_caching.py └── README.md /example_.env: -------------------------------------------------------------------------------- 1 | UPSTASH_VECTOR_REST_URL=your_url 2 | UPSTASH_VECTOR_REST_TOKEN=your_token 3 | OPENAI_API_KEY=your_key 4 | -------------------------------------------------------------------------------- /assets/how-semantic-cache-works.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/upstash/semantic-cache-py/HEAD/assets/how-semantic-cache-works.png -------------------------------------------------------------------------------- /upstash_semantic_cache/__init__.py: -------------------------------------------------------------------------------- 1 | from upstash_semantic_cache.semantic_cache import SemanticCache 2 | 3 | __version__ = "1.0.0" 4 | __all__ = ["SemanticCache"] -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | on: 4 | push: 5 | branches: [ master ] 6 | pull_request: 7 | branches: [ master ] 8 | 9 | jobs: 10 | test: 11 | runs-on: ubuntu-latest 12 | 13 | steps: 14 | - name: Checkout code 15 | uses: actions/checkout@v3 16 | 17 | - name: Set up Python 18 | uses: actions/setup-python@v4 19 | with: 20 | python-version: '3.8' 21 | 22 | - name: Install dependencies 23 | run: | 24 | python -m pip install --upgrade pip 25 | pip install poetry 26 | poetry install 27 | 28 | 29 | - name: Run ruff 30 | run: | 31 | poetry run ruff check . 32 | 33 | - name: Run mypy 34 | run: | 35 | poetry run mypy --show-error-codes . 36 | 37 | - name: Run tests 38 | env: 39 | UPSTASH_VECTOR_REST_URL: ${{ secrets.UPSTASH_VECTOR_REST_URL }} 40 | UPSTASH_VECTOR_REST_TOKEN: ${{ secrets.UPSTASH_VECTOR_REST_TOKEN }} 41 | OPENAI_API_KEY: ${{ secrets. OPENAI_API_KEY }} 42 | run: | 43 | poetry run python -m unittest discover tests 44 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | # MIT License 2 | 3 | Copyright (c) 2023 Upstash, Inc. 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /examples/llm.py: -------------------------------------------------------------------------------- 1 | import os 2 | from time import sleep 3 | import time 4 | from dotenv import load_dotenv 5 | from upstash_semantic_cache.semantic_cache import SemanticCache 6 | from langchain.globals import set_llm_cache 7 | from langchain_openai import OpenAI 8 | 9 | 10 | def example1(llm): 11 | prompt1 = "Why is the Moon always showing the same side?" 12 | prompt2 = "How come we always see one face of the moon?" 13 | 14 | start_time = time.time() 15 | response1 = llm.invoke(prompt1) 16 | print(response1) 17 | end_time = time.time() 18 | print("Time difference 1:", end_time - start_time, "seconds") 19 | sleep(1) 20 | 21 | start_time = time.time() 22 | response2 = llm.invoke(prompt2) 23 | print(response2) 24 | end_time = time.time() 25 | time_difference = end_time - start_time 26 | print("Time difference 2:", time_difference, "seconds") 27 | 28 | 29 | def main(): 30 | # set environment variables 31 | load_dotenv() 32 | UPSTASH_VECTOR_REST_URL = os.getenv("UPSTASH_VECTOR_REST_URL") 33 | UPSTASH_VECTOR_REST_TOKEN = os.getenv("UPSTASH_VECTOR_REST_TOKEN") 34 | 35 | # create a cache object 36 | cache = SemanticCache( 37 | url=UPSTASH_VECTOR_REST_URL, token=UPSTASH_VECTOR_REST_TOKEN, min_proximity=0.7 38 | ) 39 | cache.flush() 40 | sleep(1) 41 | 42 | llm = OpenAI(model_name="gpt-3.5-turbo-instruct", n=2, best_of=2) 43 | set_llm_cache(cache) 44 | example1(llm) 45 | 46 | 47 | main() 48 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "upstash-semantic-cache" 3 | version = "1.0.0" 4 | description = "Semantic Caching Algorithm for Upstash Vector" 5 | license = "MIT" 6 | authors = ["Upstash "] 7 | maintainers = ["Upstash "] 8 | readme = "README.md" 9 | repository = "https://github.com/upstash/semantic-cache-py" 10 | keywords = ["Upstash", "Upstash Vector", "Semantic Cache"] 11 | classifiers = [ 12 | "Development Status :: 5 - Production/Stable", 13 | "Intended Audience :: Developers", 14 | "License :: OSI Approved :: MIT License", 15 | "Operating System :: OS Independent", 16 | "Programming Language :: Python", 17 | "Programming Language :: Python :: 3", 18 | "Programming Language :: Python :: 3 :: Only", 19 | "Programming Language :: Python :: 3.8", 20 | "Programming Language :: Python :: 3.9", 21 | "Programming Language :: Python :: 3.10", 22 | "Programming Language :: Python :: 3.11", 23 | "Programming Language :: Python :: 3.12", 24 | "Programming Language :: Python :: Implementation :: CPython", 25 | "Topic :: Database", 26 | "Topic :: Software Development :: Libraries", 27 | ] 28 | 29 | packages = [{ include = "upstash_semantic_cache" }] 30 | 31 | [tool.poetry.dependencies] 32 | python = ">=3.8.1,<4.0" 33 | langchain_core = "^0.2.0" 34 | upstash_vector = ">=0.5.0, <1" 35 | 36 | 37 | [tool.poetry.group.dev.dependencies] 38 | pytest = "^8.2.2" 39 | python-dotenv = "^1.0.1" 40 | ruff = "^0.0.267" 41 | mypy = "^1.1.1" 42 | langchain_openai = "0.1.7" 43 | langchain = "0.2.0" 44 | 45 | [build-system] 46 | requires = ["poetry-core>=1.0.0"] 47 | build-backend = "poetry.core.masonry.api" -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | # PyCharm project settings 132 | .idea 133 | 134 | # VS Code settings 135 | .vscode 136 | 137 | # Poetry lock 138 | poetry.lock 139 | 140 | # Ruff 141 | .ruff_cache -------------------------------------------------------------------------------- /tests/test_semantic_caching.py: -------------------------------------------------------------------------------- 1 | from time import sleep 2 | import unittest 3 | from upstash_semantic_cache.semantic_cache import SemanticCache 4 | from dotenv import load_dotenv 5 | import os 6 | 7 | load_dotenv() 8 | 9 | UPSTASH_VECTOR_REST_URL = os.getenv("UPSTASH_VECTOR_REST_URL") 10 | UPSTASH_VECTOR_REST_TOKEN = os.getenv("UPSTASH_VECTOR_REST_TOKEN") 11 | 12 | # Initialize Upstash Vector Index 13 | 14 | 15 | class TestSemanticCache(unittest.TestCase): 16 | key1 = "key1" 17 | key2 = "key2" 18 | data1 = "value1" 19 | data2 = "value2" 20 | 21 | def setUp(self): 22 | # Initialize the SemanticCache instance 23 | self.cache = SemanticCache( 24 | UPSTASH_VECTOR_REST_URL, UPSTASH_VECTOR_REST_TOKEN, min_proximity=0.7 25 | ) 26 | self.refresh() 27 | 28 | def test_get_existing_key(self): 29 | # Set up a key-value pair in the cache 30 | 31 | self.cache.set(self.key1, self.data1) 32 | sleep(1) 33 | # Retrieve the value using the key 34 | result = self.cache.get(self.key1) 35 | sleep(1) 36 | # Assert that the retrieved value is correct 37 | self.assertEqual(result, self.data1) 38 | self.refresh() 39 | 40 | def test_get_nonexistent_key(self): 41 | # Retrieve a non-existent key 42 | result = self.cache.get("nonexistent_key") 43 | 44 | # Assert that the result is None 45 | self.assertIsNone(result) 46 | self.refresh() 47 | 48 | def test_set_multiple_key_values(self): 49 | # Set up multiple key-value pairs in the cache 50 | keys = [self.key1, self.key2] 51 | data = [self.data1, self.data2] 52 | self.cache.set(keys, data) 53 | sleep(1) 54 | # Retrieve the values using the keys 55 | result1 = self.cache.get(keys[0]) 56 | sleep(1) 57 | result2 = self.cache.get(keys[1]) 58 | sleep(1) 59 | # Assert that the retrieved values are correct 60 | self.assertEqual(result1, data[0]) 61 | self.assertEqual(result2, data[1]) 62 | self.refresh() 63 | 64 | def test_delete_existing_key(self): 65 | self.cache.set(self.key1, self.data1) 66 | sleep(1) 67 | # Delete the key 68 | self.cache.delete(self.key1) 69 | sleep(1) 70 | # Retrieve the value using the key 71 | result = self.cache.get(self.key1) 72 | sleep(1) 73 | # Assert that the result is None 74 | self.assertIsNone(result) 75 | self.refresh() 76 | 77 | def test_delete_nonexistent_key(self): 78 | # Set up a key-value pair in the cache 79 | key = self.key1 80 | data = self.data1 81 | self.cache.set(key, data) 82 | sleep(1) 83 | # Delete a non-existent key 84 | result = self.cache.delete("nonexistent_key") 85 | sleep(1) 86 | # Assert that the result is False 87 | self.assertFalse(result) 88 | self.refresh() 89 | 90 | def test_bulk_delete(self): 91 | # Set up multiple key-value pairs in the cache 92 | keys = [self.key1, self.key2, "key3"] 93 | data = [self.data1, self.data2, "value3"] 94 | self.cache.set(keys, data) 95 | sleep(1) 96 | # Delete multiple keys 97 | self.cache.delete(keys) 98 | sleep(1) 99 | # Retrieve the values using the keys 100 | result1 = self.cache.get(keys[0]) 101 | sleep(1) 102 | result2 = self.cache.get(keys[1]) 103 | sleep(1) 104 | result3 = self.cache.get(keys[2]) 105 | sleep(1) 106 | 107 | # Assert that the results are None 108 | self.assertIsNone(result1) 109 | self.assertIsNone(result2) 110 | self.assertIsNone(result3) 111 | self.refresh() 112 | 113 | def test_flush(self): 114 | # Set up a key-value pair in the cache 115 | key = self.key1 116 | data = self.data1 117 | self.cache.set(key, data) 118 | # Flush the cache 119 | self.cache.flush() 120 | # Retrieve the value using the key 121 | result = self.cache.get(key) 122 | # Assert that the result is None 123 | self.assertIsNone(result) 124 | self.refresh() 125 | 126 | def refresh(self): 127 | self.cache.flush() 128 | sleep(1) 129 | 130 | 131 | if __name__ == "__main__": 132 | unittest.main() 133 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Semantic Cache 2 | 3 | Semantic Cache is a tool for caching natural text based on semantic similarity. It's ideal for any task that involves querying or retrieving information based on meaning, such as natural language classification or caching AI responses. Two pieces of text can be similar but not identical (e.g., "great places to check out in Spain" vs. "best places to visit in Spain"). Traditional caching doesn't recognize this semantic similarity and misses opportunities for reuse. 4 | 5 | Semantic Cache allows you to: 6 | 7 | - Easily classify natural text into predefined categories 8 | - Avoid redundant LLM work by caching AI responses 9 | - Reduce API latency by responding to similar queries with already cached values 10 | 11 | 12 | 13 | ## Highlights 14 | 15 | - **Uses semantic similarity**: Stores cache entries by their meaning, not just the literal characters 16 | - **Handles synonyms**: Recognizes and handles synonyms 17 | - **Complex query support**: Understands long and nested user queries 18 | - **Customizable**: Set a custom proximity threshold to filter out less relevant results 19 | 20 | ## Getting Started 21 | 22 | ### Prerequisites 23 | 24 | - An Upstash Vector database (create one [here](https://console.upstash.com/vector)) 25 | 26 | ### Installation 27 | 28 | After creating a vector database, you should install the repository using the following command. 29 | 30 | ```bash 31 | pip install upstash-semantic-cache 32 | ``` 33 | 34 | To use it in your project, you must import it at the beginning of your file: 35 | 36 | ```python 37 | from upstash_semantic_cache import SemanticCache 38 | ``` 39 | 40 | ### Setup 41 | 42 | First, create an Upstash Vector database [here](https://console.upstash.com/vector). You'll need the `url` and `token` credentials to connect your semantic cache. Important: Choose any pre-made embedding model when creating your database. 43 | 44 | > [!NOTE] 45 | > Different embedding models are great for different use cases. For example, if low latency is a priority, choose a model with a smaller dimension size like `bge-small-en-v1.5`. If accuracy is important, choose a model with more dimensions. 46 | 47 | Create a `.env` file in the root directory of your project and add your Upstash Vector URL and token: 48 | 49 | ```plaintext 50 | UPSTASH_VECTOR_REST_URL=https://example.upstash.io 51 | UPSTASH_VECTOR_REST_TOKEN=your_secret_token_here 52 | ``` 53 | 54 | ### Using Semantic Cache 55 | 56 | After setting up environment variables and installing the repository, a basic demo can be created like this: 57 | 58 | ```python 59 | 60 | def main(): 61 | # set environment variables 62 | load_dotenv() 63 | UPSTASH_VECTOR_REST_URL = os.getenv('UPSTASH_VECTOR_REST_URL') 64 | UPSTASH_VECTOR_REST_TOKEN = os.getenv('UPSTASH_VECTOR_REST_TOKEN') 65 | 66 | # initialize Upstash database 67 | cache = SemanticCache(url=UPSTASH_VECTOR_REST_URL, token=UPSTASH_VECTOR_REST_TOKEN, min_proximity=0.7) 68 | cache.set('The most crowded city in Turkiye', 'Istanbul') 69 | sleep(1) 70 | result = cache.get('Which city has the most population in Turkiye?') 71 | sleep(1) 72 | print(result) 73 | 74 | if __name__ == '__main__': 75 | main() # outputs Istanbul 76 | ``` 77 | 78 | ### Using Tests and Examples 79 | 80 | In the root folder of the project, you will find tests and examples folders. If you have made changes to the project and want to examine the outputs, you can run the tests and examples using the following commands from the root of the project: 81 | 82 | ```bash 83 | python3 -m unittest discover tests 84 | ``` 85 | 86 | ```bash 87 | python3 -m examples.llm 88 | ``` 89 | 90 | If you added new example files, replace llm with the filename. 91 | 92 | ### The `minProximity` Parameter 93 | 94 | The `minProximity` parameter ranges from `0` to `1`. It lets you define the minimum relevance score to determine a cache hit. The higher this number, the more similar your user input must be to the cached content to be a hit. In practice, a score of 0.95 indicates a very high similarity, while a score of 0.75 already indicates a low similarity. For example, a value of 1.00, the highest possible, would only accept an _exact_ match of your user query and cache content as a cache hit. 95 | 96 | ## Examples 97 | 98 | The following examples demonstrate how you can utilize Semantic Cache in various use cases: 99 | 100 | > [!NOTE] 101 | > We add a 1-second delay after setting the data to allow time for the vector index to update. This delay is necessary to ensure that the data is available for retrieval. 102 | 103 | ### Basic Semantic Retrieval 104 | 105 | ```python 106 | cache.set('Capital of Turkiye', 'Ankara') 107 | sleep(1) 108 | result = cache.get('What is the capital of Turkiye?') 109 | sleep(1) 110 | print(result) # outputs Ankara 111 | ``` 112 | 113 | ### Handling Synonyms 114 | 115 | ```python 116 | cache.set('The last champion of European Football Championship', 'Italy') 117 | sleep(1) 118 | result = cache.get('Which country is the winner of the most recent European Football Championship?') 119 | sleep(1) 120 | print(result) # outputs Italy 121 | ``` 122 | 123 | ### Complex Queries 124 | 125 | ```python 126 | cache.set('The largest economy in the world, 'USA') 127 | sleep(1) 128 | result = cache.get('Which country has the highest GDP?') 129 | sleep(1) 130 | print(result) # outputs USA 131 | ``` 132 | 133 | ### Different Contexts 134 | 135 | ```python 136 | cache.set("New York population as of 2020 census", "8.8 million") 137 | cache.set("Major economic activities in New York", "Finance, technology, and tourism") 138 | sleep(1) 139 | result1 = cache.get("How many people lived in NYC according to the last census?") 140 | sleep(1) 141 | result2 = cache.get("What are the key industries in New York?") 142 | sleep(1) 143 | print(result1) # outputs 8.8 million 144 | print(result2) # outputs Finance, technology, and tourism 145 | ``` 146 | 147 | ## Contributing 148 | 149 | We appreciate your contributions! If you'd like to contribute to this project, please fork the repository, make changes, and submit a pull request. 150 | 151 | ## License 152 | 153 | It is distributed under the MIT License. See `LICENSE` for more information. 154 | -------------------------------------------------------------------------------- /upstash_semantic_cache/semantic_cache.py: -------------------------------------------------------------------------------- 1 | import hashlib 2 | import json 3 | from typing import List, Optional, Union 4 | from upstash_vector import Index 5 | from langchain_core.outputs.generation import Generation 6 | 7 | 8 | class SemanticCache: 9 | """ 10 | A class to represent a semantic cache using Upstash Vector Database. 11 | """ 12 | 13 | def __init__( 14 | self, 15 | url: str, 16 | token: str, 17 | min_proximity: float = 0.9, 18 | namespace: str = "", 19 | ) -> None: 20 | """ 21 | Initializes the SemanticCache with the given URL, token, and optional namespace. 22 | 23 | Args: 24 | url (str): The URL of the Upstash Vector database. 25 | token (str): The token for accessing the Upstash Vector database. 26 | min_proximity (float): The minimum proximity score to consider a cache hit. 27 | namespace (str): The namespace to logically separate data sets 28 | within the same index. 29 | """ 30 | self.min_proximity = min_proximity 31 | self.namespace = namespace 32 | self.index = Index(url=url, token=token) 33 | 34 | def get(self, key: str) -> Optional[str]: 35 | """ 36 | Searches the cache for the key and returns the value if it exists. 37 | 38 | Args: 39 | key (str): The key to search in the cache. 40 | 41 | Returns: 42 | Optional[str]: The value associated with the key 43 | if it exists and meets the proximity score; otherwise, None. 44 | """ 45 | response = self._query_key(key) 46 | if response is None or response.score <= self.min_proximity: 47 | return None 48 | return response.metadata["value"] 49 | 50 | def lookup( 51 | self, prompt: str, llm_string: Optional[str] = None 52 | ) -> Optional[List[Generation]]: 53 | """ 54 | Converts the JSON string to generations and returns them. 55 | It is called by Langchain to check if the generations are in the cache. 56 | 57 | Args: 58 | prompt (str): The prompt to lookup in the cache. 59 | llm_string (Optional[str]): Optional string for LLM context. 60 | 61 | Returns: 62 | Optional[List[Generation]]: 63 | If response is found in the cache; otherwise, None. 64 | """ 65 | result = self.get(prompt) 66 | return self._loads_generations(result) if result else None 67 | 68 | def update( 69 | self, 70 | prompt: str, 71 | llm_string: Optional[str] = None, 72 | result: Optional[str] = None, 73 | ) -> None: 74 | """ 75 | Converts the generations to a JSON string and stores it in the cache. 76 | It is called by Langchain to upload the generations to the cache. 77 | 78 | Args: 79 | prompt (str): The prompt to update in the cache. 80 | llm_string (Optional[str]): Optional string for LLM context. 81 | result (Optional[str]): The result to store in the cache. 82 | """ 83 | self.set(prompt, self._dumps_generations(result)) 84 | 85 | def set(self, key: Union[str, List[str]], value: Union[str, List[str]]) -> None: 86 | """ 87 | Sets the key and value in the cache. 88 | 89 | Args: 90 | key (Union[str, List[str]]): The key or list of keys to set in the cache. 91 | value (Union[str, List[str]]): The value or list of values. 92 | """ 93 | if isinstance(key, list) and isinstance(value, list): 94 | batch = [ 95 | ( 96 | self._hash_key(k), 97 | k, 98 | {"value": v}, 99 | ) 100 | for k, v in zip(key, value) 101 | ] 102 | self.index.upsert(batch, self.namespace) 103 | else: 104 | self.index.upsert( 105 | [ 106 | ( 107 | self._hash_key(key), 108 | key, 109 | {"value": value}, 110 | ) 111 | ], 112 | self.namespace, 113 | ) 114 | 115 | def delete(self, key: Union[str, List[str]]) -> None: 116 | """ 117 | Deletes the key from the cache. 118 | 119 | Args: 120 | key (str): The key (or keys) to delete from the cache. 121 | """ 122 | if isinstance(key, list): 123 | batch = [] 124 | for current_key in key: 125 | batch.append(self._hash_key(current_key)) 126 | self.index.delete(batch, namespace=self.namespace) 127 | else: 128 | self.index.delete([self._hash_key(key)], namespace=self.namespace) 129 | 130 | def flush(self) -> None: 131 | """ 132 | Resets the cache, removing all keys and values. 133 | """ 134 | self.index.reset(namespace=self.namespace) 135 | 136 | def _query_key(self, key: str): 137 | """ 138 | Queries the cache for the key. 139 | 140 | Args: 141 | key (str): The key to query in the cache. 142 | 143 | Returns: 144 | The response from the cache query. 145 | """ 146 | response = self.index.query( 147 | data=key, top_k=1, include_metadata=True, namespace=self.namespace 148 | ) 149 | return response[0] if response else None 150 | 151 | def _dumps_generations(self, generations): 152 | """ 153 | Converts the generations to a JSON string. 154 | """ 155 | 156 | def generation_to_dict(generation): 157 | if isinstance(generation, Generation): 158 | return { 159 | "text": generation.text, 160 | "generation_info": generation.generation_info, 161 | } 162 | else: 163 | raise TypeError( 164 | f"Object of type {generation.__class__.__name__}" 165 | + "is not JSON serializable" 166 | ) 167 | 168 | return json.dumps([generation_to_dict(g) for g in generations]) 169 | 170 | def _loads_generations(self, json_str): 171 | """ 172 | Converts the JSON string to generations. 173 | """ 174 | 175 | def dict_to_generation(d: dict) -> Generation: 176 | if isinstance(d, dict): 177 | return Generation(text=d["text"], generation_info=d["generation_info"]) 178 | else: 179 | raise TypeError( 180 | f"Object of type {d.__class__.__name__}" 181 | + "is not a valid Generation dict" 182 | ) 183 | 184 | return [dict_to_generation(d) for d in json.loads(json_str)] 185 | 186 | def _hash_key(self, key): 187 | """ 188 | Hashes the key to generate an ID. 189 | """ 190 | return hashlib.sha256(key.encode("utf-8")).hexdigest() 191 | --------------------------------------------------------------------------------