From 3f26a5269b0f5f7702fdd29efec307c03d479bee Mon Sep 17 00:00:00 2001 From: Amnon Catav Date: Sun, 29 Oct 2023 12:19:43 +0200 Subject: [PATCH 1/3] support list canopy indexes --- src/canopy/knowledge_base/__init__.py | 1 + src/canopy/knowledge_base/knowledge_base.py | 12 ++++++++++++ tests/system/knowledge_base/test_knowledge_base.py | 11 ++++++++++- 3 files changed, 23 insertions(+), 1 deletion(-) diff --git a/src/canopy/knowledge_base/__init__.py b/src/canopy/knowledge_base/__init__.py index 615a5736..fc0bb83a 100644 --- a/src/canopy/knowledge_base/__init__.py +++ b/src/canopy/knowledge_base/__init__.py @@ -1 +1,2 @@ from .knowledge_base import KnowledgeBase +from .knowledge_base import list_canopy_indexes diff --git a/src/canopy/knowledge_base/knowledge_base.py b/src/canopy/knowledge_base/knowledge_base.py index c51f491f..7cb5f6dc 100644 --- a/src/canopy/knowledge_base/knowledge_base.py +++ b/src/canopy/knowledge_base/knowledge_base.py @@ -37,6 +37,18 @@ DELETE_STARTER_CHUNKS_PER_DOC = 32 +def list_canopy_indexes() -> List[str]: + + try: + pinecone_init() + pinecone_whoami() + except Exception as e: + raise RuntimeError("Failed to connect to Pinecone. " + "Please check your credentials and try again") from e + + return [index for index in list_indexes() if index.startswith(INDEX_NAME_PREFIX)] + + class KnowledgeBase(BaseKnowledgeBase): """ diff --git a/tests/system/knowledge_base/test_knowledge_base.py b/tests/system/knowledge_base/test_knowledge_base.py index 643b81b7..f8ef1216 100644 --- a/tests/system/knowledge_base/test_knowledge_base.py +++ b/tests/system/knowledge_base/test_knowledge_base.py @@ -12,7 +12,7 @@ ) from dotenv import load_dotenv from datetime import datetime -from canopy.knowledge_base import KnowledgeBase +from canopy.knowledge_base import KnowledgeBase, list_canopy_indexes from canopy.knowledge_base.chunker import Chunker from canopy.knowledge_base.knowledge_base import INDEX_NAME_PREFIX from canopy.knowledge_base.models import DocumentWithScore @@ -226,6 +226,15 @@ def test_create_index(index_full_name, knowledge_base): assert knowledge_base._index.describe_index_stats() +def test_list_indexes(index_full_name): + actual = list_canopy_indexes() + expected = [index for index in pinecone.list_indexes() + if index.startswith(INDEX_NAME_PREFIX)] + + assert index_full_name in actual + assert sorted(actual) == sorted(expected) + + def test_is_verify_index_connection_happy_path(knowledge_base): knowledge_base.verify_index_connection() From e83ca8dabfb60649e4f12fd32c8022df174531b3 Mon Sep 17 00:00:00 2001 From: Amnon Catav Date: Sun, 29 Oct 2023 15:01:46 +0200 Subject: [PATCH 2/3] add docstring --- src/canopy/knowledge_base/knowledge_base.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/canopy/knowledge_base/knowledge_base.py b/src/canopy/knowledge_base/knowledge_base.py index 7cb5f6dc..eb4cfbd5 100644 --- a/src/canopy/knowledge_base/knowledge_base.py +++ b/src/canopy/knowledge_base/knowledge_base.py @@ -38,6 +38,17 @@ def list_canopy_indexes() -> List[str]: + """ + List all Canopy indexes in the current Pinecone account. + + Example: + >>> from canopy.knowledge_base import list_canopy_indexes + >>> list_canopy_indexes() + ['canopy--my_index', 'canopy--my_index2'] + + Returns: + A list of Canopy index names. + """ try: pinecone_init() From d34379235fae386374b84676ecb9640579a9be88 Mon Sep 17 00:00:00 2001 From: Amnon Catav Date: Mon, 30 Oct 2023 09:02:33 +0200 Subject: [PATCH 3/3] extract connect to pinecone to an external method --- src/canopy/knowledge_base/__init__.py | 3 +- src/canopy/knowledge_base/knowledge_base.py | 35 ++++++++++----------- src/canopy_cli/cli.py | 3 +- 3 files changed, 21 insertions(+), 20 deletions(-) diff --git a/src/canopy/knowledge_base/__init__.py b/src/canopy/knowledge_base/__init__.py index fc0bb83a..e2c8ccac 100644 --- a/src/canopy/knowledge_base/__init__.py +++ b/src/canopy/knowledge_base/__init__.py @@ -1,2 +1,3 @@ -from .knowledge_base import KnowledgeBase +from .knowledge_base import connect_to_pinecone from .knowledge_base import list_canopy_indexes +from .knowledge_base import KnowledgeBase diff --git a/src/canopy/knowledge_base/knowledge_base.py b/src/canopy/knowledge_base/knowledge_base.py index eb4cfbd5..85687685 100644 --- a/src/canopy/knowledge_base/knowledge_base.py +++ b/src/canopy/knowledge_base/knowledge_base.py @@ -37,6 +37,20 @@ DELETE_STARTER_CHUNKS_PER_DOC = 32 +def connect_to_pinecone(): + """ + Connect to Pinecone. + This method is called automatically when creating a new KnowledgeBase object. + Or when calling `list_canopy_indexes()`. + """ + try: + pinecone_init() + pinecone_whoami() + except Exception as e: + raise RuntimeError("Failed to connect to Pinecone. " + "Please check your credentials and try again") from e + + def list_canopy_indexes() -> List[str]: """ List all Canopy indexes in the current Pinecone account. @@ -50,13 +64,7 @@ def list_canopy_indexes() -> List[str]: A list of Canopy index names. """ - try: - pinecone_init() - pinecone_whoami() - except Exception as e: - raise RuntimeError("Failed to connect to Pinecone. " - "Please check your credentials and try again") from e - + connect_to_pinecone() return [index for index in list_indexes() if index.startswith(INDEX_NAME_PREFIX)] @@ -183,20 +191,11 @@ def __init__(self, # `create_canopy_index()` self._index: Optional[Index] = None - @staticmethod - def _connect_pinecone(): - try: - pinecone_init() - pinecone_whoami() - except Exception as e: - raise RuntimeError("Failed to connect to Pinecone. " - "Please check your credentials and try again") from e - def _connect_index(self, connect_pinecone: bool = True ) -> None: if connect_pinecone: - self._connect_pinecone() + connect_to_pinecone() if self.index_name not in list_indexes(): raise RuntimeError( @@ -320,7 +319,7 @@ def create_canopy_index(self, "Please provide the vectors' dimension") # connect to pinecone and create index - self._connect_pinecone() + connect_to_pinecone() if self.index_name in list_indexes(): raise RuntimeError( diff --git a/src/canopy_cli/cli.py b/src/canopy_cli/cli.py index a9e473de..dbdf1ddd 100644 --- a/src/canopy_cli/cli.py +++ b/src/canopy_cli/cli.py @@ -18,6 +18,7 @@ from urllib.parse import urljoin from canopy.knowledge_base import KnowledgeBase +from canopy.knowledge_base import connect_to_pinecone from canopy.models.data_models import Document from canopy.tokenizer import Tokenizer from canopy_cli.data_loader import ( @@ -71,7 +72,7 @@ def wait_for_service(chat_service_url: str): def validate_connection(): try: - KnowledgeBase._connect_pinecone() + connect_to_pinecone() except RuntimeError as e: msg = ( f"{str(e)}\n"