%pip install -qU langchain-google-bigtable langchain-google-vertexai

# Automatically restart kernel after installs so that your environment can access the new packages
# import IPython

# app = IPython.Application.instance()
# app.kernel.do_shutdown(True)

# @markdown Please fill in your project, instance, and a new table name.
PROJECT_ID = "test-project"  # @param {type:"string"}
INSTANCE_ID = "test-instance"  # @param {type:"string"}
TABLE_ID = "your-vector-store-table-3"  # @param {type:"string"}

!gcloud config set project {PROJECT_ID}

from google.colab import auth

auth.authenticate_user(project_id=PROJECT_ID)

from langchain_google_vertexai import VertexAIEmbeddings

embeddings = VertexAIEmbeddings(project=PROJECT_ID, model_name="gemini-embedding-001")

from langchain_google_bigtable.vector_store import init_vector_store_table

DATA_COLUMN_FAMILY = "doc_data"

try:
    init_vector_store_table(
        project_id=PROJECT_ID,
        instance_id=INSTANCE_ID,
        table_id=TABLE_ID,
        content_column_family=DATA_COLUMN_FAMILY,
        embedding_column_family=DATA_COLUMN_FAMILY,
    )
    print(f"Table '{TABLE_ID}' is ready.")
except ValueError as e:
    print(e)

from langchain_google_bigtable import BigtableEngine

engine = await BigtableEngine.async_initialize(project_id=PROJECT_ID)

collection_name = "my_docs"

from langchain_google_bigtable import ColumnConfig, VectorMetadataMapping, Encoding

# Define mappings for metadata fields you want to filter on.
metadata_mappings = [
    VectorMetadataMapping(metadata_key="author", encoding=Encoding.UTF8),
    VectorMetadataMapping(metadata_key="year", encoding=Encoding.INT_BIG_ENDIAN),
    VectorMetadataMapping(metadata_key="category", encoding=Encoding.UTF8),
    VectorMetadataMapping(metadata_key="rating", encoding=Encoding.FLOAT),
]

# Define the optional column for storing all metadata as a single JSON string.
metadata_as_json_column = ColumnConfig(
    column_family=DATA_COLUMN_FAMILY, column_qualifier="metadata_json"
)

# Configure the columns for your store.
content_column = ColumnConfig(
    column_family=DATA_COLUMN_FAMILY, column_qualifier="content"
)
embedding_column = ColumnConfig(
    column_family=DATA_COLUMN_FAMILY, column_qualifier="embedding"
)

from langchain_google_bigtable import BigtableVectorStore

vector_store = await BigtableVectorStore.create(
    project_id=PROJECT_ID,
    instance_id=INSTANCE_ID,
    table_id=TABLE_ID,
    engine=engine,
    embedding_service=embeddings,
    collection=collection_name,
    metadata_mappings=metadata_mappings,
    metadata_as_json_column=metadata_as_json_column,
    content_column=content_column,
    embedding_column=embedding_column,
)

from langchain_core.documents import Document

docs_to_add = [
    Document(
        page_content="A young farm boy, Luke Skywalker, is thrust into a galactic conflict.",
        id="doc_1",
        metadata={
            "author": "George Lucas",
            "year": 1977,
            "category": "sci-fi",
            "rating": 4.8,
        },
    ),
    Document(
        page_content="A hobbit named Frodo Baggins must destroy a powerful ring.",
        id="doc_2",
        metadata={
            "author": "J.R.R. Tolkien",
            "year": 1954,
            "category": "fantasy",
            "rating": 4.9,
        },
    ),
    # Document without a pre-defined ID, one will be generated.
    Document(
        page_content="A group of children confront an evil entity emerging from the sewers.",
        metadata={"author": "Stephen King", "year": 1986, "category": "horror"},
    ),
    Document(
        page_content="In a distant future, the noble House Atreides rules the desert planet Arrakis.",
        id="doc_3",
        metadata={
            "author": "Frank Herbert",
            "year": 1965,
            "category": "sci-fi",
            "rating": 4.9,
        },
    ),
]

added_ids = await vector_store.aadd_documents(docs_to_add)
print(f"Added documents with IDs: {added_ids}")

doc_to_update = [
    Document(
        page_content="An old hobbit, Frodo Baggins, must take a powerful ring to be destroyed.",  # Updated content
        id="doc_2",  # Same ID
        metadata={
            "author": "J.R.R. Tolkien",
            "year": 1954,
            "category": "epic-fantasy",
            "rating": 4.9,
        },  # Updated metadata
    )
]

await vector_store.aadd_documents(doc_to_update)
print("Document 'doc_2' has been updated.")

is_deleted = await vector_store.adelete(ids=["doc_2"])

results = await vector_store.asimilarity_search("a story about a powerful ring", k=1)
print(results[0].page_content)

from langchain_google_bigtable.vector_store import QueryParameters

complex_filter = {
    "ColumnValueFilter": {
        "ColumnValueUnionFilter": {  # OR
            "ColumnValueChainFilter": {  # First AND condition
                "category": {"==": "sci-fi"},
                "year": {">": 1970, "<": 2000},
            },
            "author": {"==": "J.R.R. Tolkien"},
        }
    }
}

query_params_complex = QueryParameters(filters=complex_filter)

complex_results = await vector_store.asimilarity_search(
    "a story about a hero's journey", k=5, query_parameters=query_params_complex
)

print(f"Found {len(complex_results)} documents matching the complex filter:")
for doc in complex_results:
    print(f"- ID: {doc.id}, Metadata: {doc.metadata}")

results_with_scores = await vector_store.asimilarity_search_with_score(
    query="an evil entity", k=1
)
for doc, score in results_with_scores:
    print(f"* [SCORE={score:.4f}] {doc.page_content} [{doc.metadata}]")

# Define a filter to use with the retriever
retriever_filter = {"ColumnValueFilter": {"category": {"==": "horror"}}}
retriever_query_params = QueryParameters(filters=retriever_filter)

retriever = vector_store.as_retriever(
    search_type="mmr",  # Specify MMR for retrieval
    search_kwargs={
        "k": 1,
        "lambda_mult": 0.8,
        "query_parameters": retriever_query_params,  # Pass filter parameters
    },
)
retrieved_docs = await retriever.ainvoke("a story about a hobbit")
print(retrieved_docs[0].page_content)