from langchain_google_bigtable import ( BigtableByteStore, BigtableEngine, init_key_value_store_table,)# Ensure the table and column family exist.init_key_value_store_table( project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID,)
all_keys = [key async for key in store.ayield_keys()]print(f"All keys: {all_keys}")prefixed_keys = [key async for key in store.ayield_keys(prefix="key1")]print(f"Prefixed keys: {prefixed_keys}")
from langchain_classic.embeddings import CacheBackedEmbeddingsfrom langchain_google_vertexai.embeddings import VertexAIEmbeddingsunderlying_embeddings = VertexAIEmbeddings( project=PROJECT_ID, model_name="textembedding-gecko@003")# Use a namespace to avoid key collisions with other data.cached_embedder = CacheBackedEmbeddings.from_bytes_store( underlying_embeddings, store, namespace="text-embeddings")
from langchain_core.retrievers import BaseRetrieverfrom langchain_core.documents import Documentfrom langchain_core.callbacks import CallbackManagerForRetrieverRunfrom typing import List, Optional, Any, Unionimport jsonclass SimpleKVStoreRetriever(BaseRetriever): """A simple retriever that retrieves documents based on a prefix match in the key-value store.""" store: BigtableByteStore documents: List[Union[Document, str]] k: int def set_up_store(self): kv_pairs_to_set = [] for i, doc in enumerate(self.documents): if isinstance(doc, str): doc = Document(page_content=doc) if not doc.id: doc.id = str(i) value = ( "Page Content\n" + doc.page_content + "\nMetadata" + json.dumps(doc.metadata) ) kv_pairs_to_set.append((doc.id, value.encode("utf-8"))) self.store.mset(kv_pairs_to_set) async def _aget_relevant_documents( self, query: str, *, run_manager: Optional[CallbackManagerForRetrieverRun] = None, ) -> List[Document]: keys = [key async for key in self.store.ayield_keys(prefix=query)][: self.k] documents_retrieved = [] async for document in await self.store.amget(keys): if document: document_str = document.decode("utf-8") page_content = document_str.split("Content\n")[1].split("\nMetadata")[0] metadata = json.loads(document_str.split("\nMetadata")[1]) documents_retrieved.append( Document(page_content=page_content, metadata=metadata) ) return documents_retrieved def _get_relevant_documents( self, query: str, *, run_manager: Optional[CallbackManagerForRetrieverRun] = None, ) -> list[Document]: keys = [key for key in self.store.yield_keys(prefix=query)][: self.k] documents_retrieved = [] for document in self.store.mget(keys): if document: document_str = document.decode("utf-8") page_content = document_str.split("Content\n")[1].split("\nMetadata")[0] metadata = json.loads(document_str.split("\nMetadata")[1]) documents_retrieved.append( Document(page_content=page_content, metadata=metadata) ) return documents_retrieved
documents = [ Document( page_content="Goldfish are popular pets for beginners, requiring relatively simple care.", metadata={"type": "fish", "trait": "low maintenance"}, id="fish#Goldfish", ), Document( page_content="Cats are independent pets that often enjoy their own space.", metadata={"type": "cat", "trait": "independence"}, id="mammals#Cats", ), Document( page_content="Rabbits are social animals that need plenty of space to hop around.", metadata={"type": "rabbit", "trait": "social"}, id="mammals#Rabbits", ),]