Zep Cloud

import getpass
import time
from uuid import uuid4

from langchain_community.memory.zep_cloud_memory import ZepCloudMemory
from langchain_community.retrievers import ZepCloudRetriever
from langchain_core.messages import AIMessage, HumanMessage

# Provide your Zep API key.
zep_api_key = getpass.getpass()

session_id = str(uuid4())  # This is a unique identifier for the user/session

# Initialize the Zep Memory Class
zep_memory = ZepCloudMemory(session_id=session_id, api_key=zep_api_key)

# Preload some messages into the memory. The default message window is 4 messages. We want to push beyond this to demonstrate auto-summarization.
test_history = [
    {"role": "human", "role_type": "user", "content": "Who was Octavia Butler?"},
    {
        "role": "ai",
        "role_type": "assistant",
        "content": (
            "Octavia Estelle Butler (June 22, 1947 – February 24, 2006) was an American"
            " science fiction author."
        ),
    },
    {
        "role": "human",
        "role_type": "user",
        "content": "Which books of hers were made into movies?",
    },
    {
        "role": "ai",
        "role_type": "assistant",
        "content": (
            "The most well-known adaptation of Octavia Butler's work is the FX series"
            " Kindred, based on her novel of the same name."
        ),
    },
    {"role": "human", "role_type": "user", "content": "Who were her contemporaries?"},
    {
        "role": "ai",
        "role_type": "assistant",
        "content": (
            "Octavia Butler's contemporaries included Ursula K. Le Guin, Samuel R."
            " Delany, and Joanna Russ."
        ),
    },
    {"role": "human", "role_type": "user", "content": "What awards did she win?"},
    {
        "role": "ai",
        "role_type": "assistant",
        "content": (
            "Octavia Butler won the Hugo Award, the Nebula Award, and the MacArthur"
            " Fellowship."
        ),
    },
    {
        "role": "human",
        "role_type": "user",
        "content": "Which other women sci-fi writers might I want to read?",
    },
    {
        "role": "ai",
        "role_type": "assistant",
        "content": "You might want to read Ursula K. Le Guin or Joanna Russ.",
    },
    {
        "role": "human",
        "role_type": "user",
        "content": (
            "Write a short synopsis of Butler's book, Parable of the Sower. What is it"
            " about?"
        ),
    },
    {
        "role": "ai",
        "role_type": "assistant",
        "content": (
            "Parable of the Sower is a science fiction novel by Octavia Butler,"
            " published in 1993. It follows the story of Lauren Olamina, a young woman"
            " living in a dystopian future where society has collapsed due to"
            " environmental disasters, poverty, and violence."
        ),
    },
    {
        "role": "human",
        "role_type": "user",
        "content": "What is the setting of the book?",
    },
    {
        "role": "ai",
        "role_type": "assistant",
        "content": (
            "The book is set in a dystopian future in the 2020s, where society has"
            " collapsed due to climate change and economic crises."
        ),
    },
    {"role": "human", "role_type": "user", "content": "Who is the protagonist?"},
    {
        "role": "ai",
        "role_type": "assistant",
        "content": (
            "The protagonist of the book is Lauren Olamina, a young woman who possesses"
            " 'hyperempathy', the ability to feel pain and other sensations she"
            " witnesses."
        ),
    },
    {
        "role": "human",
        "role_type": "user",
        "content": "What is the main theme of the book?",
    },
    {
        "role": "ai",
        "role_type": "assistant",
        "content": (
            "The main theme of the book is survival in the face of drastic societal"
            " change and collapse. It also explores themes of adaptability, community,"
            " and the human capacity for change."
        ),
    },
    {
        "role": "human",
        "role_type": "user",
        "content": "What is the 'Parable of the Sower'?",
    },
    {
        "role": "ai",
        "role_type": "assistant",
        "content": (
            "The 'Parable of the Sower' is a biblical parable that Butler uses as a"
            " metaphor in the book. In the parable, a sower scatters seeds, some of"
            " which fall on fertile ground and grow, while others fall on rocky ground"
            " or among thorns and fail to grow. The parable is used to illustrate the"
            " importance of receptivity and preparedness in the face of change."
        ),
    },
    {
        "role": "human",
        "role_type": "user",
        "content": "What is Butler's writing style like?",
    },
    {
        "role": "ai",
        "role_type": "assistant",
        "content": (
            "Butler's writing style is known for its clarity, directness, and"
            " psychological insight. Her narratives often involve complex, diverse"
            " characters and explore themes of race, gender, and power."
        ),
    },
    {
        "role": "human",
        "role_type": "user",
        "content": "What other books has she written?",
    },
    {
        "role": "ai",
        "content": (
            "In addition to 'Parable of the Sower', Butler has written several other"
            " notable works, including 'Kindred', 'Dawn', and 'Parable of the Talents'."
        ),
    },
]

for msg in test_history:
    zep_memory.chat_memory.add_message(
        HumanMessage(content=msg["content"])
        if msg["role"] == "human"
        else AIMessage(content=msg["content"])
    )

time.sleep(
    10
)  # Wait for the messages to be embedded and summarized, this happens asynchronously.

zep_retriever = ZepCloudRetriever(
    api_key=zep_api_key,
    session_id=session_id,  # Ensure that you provide the session_id when instantiating the Retriever
    top_k=5,
)

await zep_retriever.ainvoke("Who wrote Parable of the Sower?")

zep_retriever.invoke("Who wrote Parable of the Sower?")

zep_retriever = ZepCloudRetriever(
    api_key=zep_api_key,
    session_id=session_id,  # Ensure that you provide the session_id when instantiating the Retriever
    top_k=5,
    search_type="mmr",
    mmr_lambda=0.5,
)

await zep_retriever.ainvoke("Who wrote Parable of the Sower?")

filter = {"where": {"jsonpath": '$[*] ? (@.baz == "qux")'}}

await zep_retriever.ainvoke(
    "Who wrote Parable of the Sower?", config={"metadata": filter}
)

zep_retriever = ZepCloudRetriever(
    api_key=zep_api_key,
    session_id=session_id,  # Ensure that you provide the session_id when instantiating the Retriever
    top_k=3,
    search_scope="summary",
    search_type="mmr",
    mmr_lambda=0.5,
)

await zep_retriever.ainvoke("Who wrote Parable of the Sower?")