long_term_memory.py

import chromadb
from chromadb.utils import embedding_functions

embedding_function = embedding_functions.DefaultEmbeddingFunction()
client = chromadb.PersistentClient(path="_memories/long_term_memory")
long_term_memory = client.get_or_create_collection(name="long_term_memory", embedding_function=embedding_function, metadata={"hnsw:space": "cosine"})


def get_ltm():  # Only get documents and ids from long term memory
    res = long_term_memory.get(
        include=["documents", "metadatas"]
    )
    return res


# def add_content_to_ltm(content_id, content, author, iteration, virality_score, sentiment_score):
#     ids: list = []
#     metadatas: list = []
#     documents: list = []

#     ids.append(content_id)
#     metadatas.append({"Author": author, "Iteration": iteration, "Virality Score": virality_score, "Sentiment Score": sentiment_score})
#     documents.append(content)

#     try:
#         long_term_memory.add(
#             ids=ids,
#             metadatas=metadatas,
#             documents=documents
#         )
#     except Exception as e:
#         print("Add data to db failed: ", e)


def add_content_to_ltm( content_id, content, author, iteration, virality_score, sentiment_score, original_content_id=None, direct_interaction_id=None):
    ids = [content_id]
    metadatas = [{
        "Author": author,
        "Iteration": iteration,
        "Virality Score": virality_score,
        "Sentiment Score": sentiment_score,
        "Original_Content_ID": original_content_id if original_content_id else content_id,
        "Direct_Interaction_ID": direct_interaction_id if direct_interaction_id else content_id
    }]
    documents = [content]

    try:
        long_term_memory.add(
            ids=ids,
            metadatas=metadatas,
            documents=documents
        )
    except Exception as e:
        print("Add data to db failed: ", e)


def is_content_in_ltm(content_id):
    res = long_term_memory.get(
        ids=content_id,
        include=["documents"],
    )
    return bool(res["documents"])


def modify_ltm_virality_score(content_id, new_virality_score):
    res = long_term_memory.get(
        ids=content_id,
        include=["metadatas"],
    )

    virality_score = res["metadatas"][0]["Virality Score"]

    long_term_memory.update(
        ids=content_id,
        metadatas=[{"Virality Score": virality_score + new_virality_score}]
    )


def modify_ltm_sentiment_score(content_id, new_sentiment_score):
    res = long_term_memory.get(
        ids=content_id,
        include=["metadatas"],
    )

    current_sentiment_score = res["metadatas"][0]["Sentiment Score"]

    long_term_memory.update(
        ids=content_id,
        metadatas=[{"Sentiment Score": current_sentiment_score + new_sentiment_score}]
    )


def get_source_agent_from_ltm(content_id):
    res = long_term_memory.get(
        ids=content_id,
        include=["documents", "metadatas"],
    )

    return res["metadatas"][0]["Author"]


def get_feedbacks_from_ltm(agent):
    res = long_term_memory.get(
        where={"Author": str(agent.name.lower())},
        include=["documents", "metadatas"],
    )

    num_documents = len(res['documents'])

    output_strings = []

    for i in range(num_documents):
        document = res['documents'][i]
        virality_score = res['metadatas'][i]['Virality Score']
        content_score = res['metadatas'][i]['Sentiment Score']
        output_string = f"{document} - Virality Score: {virality_score} - Content Score: {content_score}"
        output_strings.append(output_string)
    final_output = "\n".join(output_strings)
    return final_output


def clear_ltm():
    long_term_memory.delete(
        where={"Iteration": {"$gte": 0}}
    )