my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import pytest
#2
#3	from embedchain.config.vector_db.pinecone import PineconeDBConfig
#4	from embedchain.vectordb.pinecone import PineconeDB
#5
#6
#7	@pytest.fixture
#8	def pinecone_pod_config():
#9	return PineconeDBConfig(
#10	index_name="test_collection",
#11	api_key="test_api_key",
#12	vector_dimension=3,
#13	pod_config={"environment": "test_environment", "metadata_config": {"indexed": ["*"]}},
#14	)
#15
#16
#17	@pytest.fixture
#18	def pinecone_serverless_config():
#19	return PineconeDBConfig(
#20	index_name="test_collection",
#21	api_key="test_api_key",
#22	vector_dimension=3,
#23	serverless_config={
#24	"cloud": "test_cloud",
#25	"region": "test_region",
#26	},
#27	)
#28
#29
#30	def test_pinecone_init_without_config(monkeypatch):
#31	monkeypatch.setenv("PINECONE_API_KEY", "test_api_key")
#32	monkeypatch.setattr("embedchain.vectordb.pinecone.PineconeDB._setup_pinecone_index", lambda x: x)
#33	monkeypatch.setattr("embedchain.vectordb.pinecone.PineconeDB._get_or_create_db", lambda x: x)
#34	pinecone_db = PineconeDB()
#35
#36	assert isinstance(pinecone_db, PineconeDB)
#37	assert isinstance(pinecone_db.config, PineconeDBConfig)
#38	assert pinecone_db.config.pod_config == {"environment": "gcp-starter", "metadata_config": {"indexed": ["*"]}}
#39	monkeypatch.delenv("PINECONE_API_KEY")
#40
#41
#42	def test_pinecone_init_with_config(pinecone_pod_config, monkeypatch):
#43	monkeypatch.setattr("embedchain.vectordb.pinecone.PineconeDB._setup_pinecone_index", lambda x: x)
#44	monkeypatch.setattr("embedchain.vectordb.pinecone.PineconeDB._get_or_create_db", lambda x: x)
#45	pinecone_db = PineconeDB(config=pinecone_pod_config)
#46
#47	assert isinstance(pinecone_db, PineconeDB)
#48	assert isinstance(pinecone_db.config, PineconeDBConfig)
#49
#50	assert pinecone_db.config.pod_config == pinecone_pod_config.pod_config
#51
#52	pinecone_db = PineconeDB(config=pinecone_pod_config)
#53
#54	assert isinstance(pinecone_db, PineconeDB)
#55	assert isinstance(pinecone_db.config, PineconeDBConfig)
#56
#57	assert pinecone_db.config.serverless_config == pinecone_pod_config.serverless_config
#58
#59
#60	class MockListIndexes:
#61	def names(self):
#62	return ["test_collection"]
#63
#64
#65	class MockPineconeIndex:
#66	db = []
#67
#68	def __init__(args, *kwargs):
#69	pass
#70
#71	def upsert(self, chunk, **kwargs):
#72	self.db.extend([c for c in chunk])
#73	return
#74
#75	def delete(self, args, *kwargs):
#76	pass
#77
#78	def query(self, args, *kwargs):
#79	return {
#80	"matches": [
#81	{
#82	"metadata": {
#83	"key": "value",
#84	"text": "text_1",
#85	},
#86	"score": 0.1,
#87	},
#88	{
#89	"metadata": {
#90	"key": "value",
#91	"text": "text_2",
#92	},
#93	"score": 0.2,
#94	},
#95	]
#96	}
#97
#98	def fetch(self, args, *kwargs):
#99	return {
#100	"vectors": {
#101	"key_1": {
#102	"metadata": {
#103	"source": "1",
#104	}
#105	},
#106	"key_2": {
#107	"metadata": {
#108	"source": "2",
#109	}
#110	},
#111	}
#112	}
#113
#114	def describe_index_stats(self, args, *kwargs):
#115	return {"total_vector_count": len(self.db)}
#116
#117
#118	class MockPineconeClient:
#119	def __init__(args, *kwargs):
#120	pass
#121
#122	def list_indexes(self):
#123	return MockListIndexes()
#124
#125	def create_index(self, args, *kwargs):
#126	pass
#127
#128	def Index(self, args, *kwargs):
#129	return MockPineconeIndex()
#130
#131	def delete_index(self, args, *kwargs):
#132	pass
#133
#134
#135	class MockPinecone:
#136	def __init__(args, *kwargs):
#137	pass
#138
#139	def Pinecone(args, *kwargs):
#140	return MockPineconeClient()
#141
#142	def PodSpec(args, *kwargs):
#143	pass
#144
#145	def ServerlessSpec(args, *kwargs):
#146	pass
#147
#148
#149	class MockEmbedder:
#150	def embedding_fn(self, documents):
#151	return [[1, 1, 1] for d in documents]
#152
#153
#154	def test_setup_pinecone_index(pinecone_pod_config, pinecone_serverless_config, monkeypatch):
#155	monkeypatch.setattr("embedchain.vectordb.pinecone.pinecone", MockPinecone)
#156	monkeypatch.setenv("PINECONE_API_KEY", "test_api_key")
#157	pinecone_db = PineconeDB(config=pinecone_pod_config)
#158	pinecone_db._setup_pinecone_index()
#159
#160	assert pinecone_db.client is not None
#161	assert pinecone_db.config.index_name == "test_collection"
#162	assert pinecone_db.client.list_indexes().names() == ["test_collection"]
#163	assert pinecone_db.pinecone_index is not None
#164
#165	pinecone_db = PineconeDB(config=pinecone_serverless_config)
#166	pinecone_db._setup_pinecone_index()
#167
#168	assert pinecone_db.client is not None
#169	assert pinecone_db.config.index_name == "test_collection"
#170	assert pinecone_db.client.list_indexes().names() == ["test_collection"]
#171	assert pinecone_db.pinecone_index is not None
#172
#173
#174	def test_get(monkeypatch):
#175	def mock_pinecone_db():
#176	monkeypatch.setenv("PINECONE_API_KEY", "test_api_key")
#177	monkeypatch.setattr("embedchain.vectordb.pinecone.PineconeDB._setup_pinecone_index", lambda x: x)
#178	monkeypatch.setattr("embedchain.vectordb.pinecone.PineconeDB._get_or_create_db", lambda x: x)
#179	db = PineconeDB()
#180	db.pinecone_index = MockPineconeIndex()
#181	return db
#182
#183	pinecone_db = mock_pinecone_db()
#184	ids = pinecone_db.get(["key_1", "key_2"])
#185	assert ids == {"ids": ["key_1", "key_2"], "metadatas": [{"source": "1"}, {"source": "2"}]}
#186
#187
#188	def test_add(monkeypatch):
#189	def mock_pinecone_db():
#190	monkeypatch.setenv("PINECONE_API_KEY", "test_api_key")
#191	monkeypatch.setattr("embedchain.vectordb.pinecone.PineconeDB._setup_pinecone_index", lambda x: x)
#192	monkeypatch.setattr("embedchain.vectordb.pinecone.PineconeDB._get_or_create_db", lambda x: x)
#193	db = PineconeDB()
#194	db.pinecone_index = MockPineconeIndex()
#195	db._set_embedder(MockEmbedder())
#196	return db
#197
#198	pinecone_db = mock_pinecone_db()
#199	pinecone_db.add(["text_1", "text_2"], [{"key_1": "value_1"}, {"key_2": "value_2"}], ["key_1", "key_2"])
#200	assert pinecone_db.count() == 2
#201
#202	pinecone_db.add(["text_3", "text_4"], [{"key_3": "value_3"}, {"key_4": "value_4"}], ["key_3", "key_4"])
#203	assert pinecone_db.count() == 4
#204
#205
#206	def test_query(monkeypatch):
#207	def mock_pinecone_db():
#208	monkeypatch.setenv("PINECONE_API_KEY", "test_api_key")
#209	monkeypatch.setattr("embedchain.vectordb.pinecone.PineconeDB._setup_pinecone_index", lambda x: x)
#210	monkeypatch.setattr("embedchain.vectordb.pinecone.PineconeDB._get_or_create_db", lambda x: x)
#211	db = PineconeDB()
#212	db.pinecone_index = MockPineconeIndex()
#213	db._set_embedder(MockEmbedder())
#214	return db
#215
#216	pinecone_db = mock_pinecone_db()
#217	# without citations
#218	results = pinecone_db.query(["text_1", "text_2"], n_results=2, where={})
#219	assert results == ["text_1", "text_2"]
#220	# with citations
#221	results = pinecone_db.query(["text_1", "text_2"], n_results=2, where={}, citations=True)
#222	assert results == [
#223	("text_1", {"key": "value", "text": "text_1", "score": 0.1}),
#224	("text_2", {"key": "value", "text": "text_2", "score": 0.2}),
#225	]
#226

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public