How to Implement Working Memory in AI Agents and Agentic Systems for Real-time AI Applications

1
from datasets import load_dataset
2
3
import pandas as pd
4
5
# Make sure you have an HF_TOKEN in your environment variables to access the dataset on hugging face 
6
7
product_dataset = load_dataset("philschmid/amazon-product-descriptions-vlm")
8
9
# Convert product_dataset to pandas dataframe
10
11
product_dataframe = pd.DataFrame(product_dataset['train'])

1
try:
2
3
  documents = product_dataframe.to_dict('records')
4
5
  product_collection.insert_many(documents)
6
7
  print("Data ingestion into MongoDB completed")
8
9
except Exception as e:
10
11
  print(f"Error during data ingestion into MongoDB: {e}")

1
import cohere
2
3
co = cohere.ClientV2()
4
5
def get_embedding(texts, model="embed-english-v3.0", input_type="search_document"):
6
7
  """Gets embeddings for a list of texts using the Cohere API.
8
9
  Args:
10
11
    texts: A list of texts to embed.
12
13
    model: The Cohere embedding model to use.
14
15
    input_type: The input type for the embedding model.
16
17
  Returns:
18
19
    A list of embeddings, where each embedding is a list of floats.
20
21
  """
22
23
  try:
24
25
    response = co.embed(
26
27
        texts=[texts],
28
29
        model=model,
30
31
        input_type=input_type,
32
33
        embedding_types=["float"],
34
35
    )
36
37
    # Extract and return the embeddings
38
39
    return response.embeddings.float[0]
40
41
  except Exception as e:
42
43
    print(f"Error generating embeddings: {e}")
44
45
    print("Couldn't generate emebedding for text: ")
46
47
    print(texts)
48
49
    return None

1
# Generate an embedding attribute for each data point in the dataset
2
3
# Embedding is generated from the new product semantics attribute
4
5
try:
6
7
  product_dataframe['embedding'] = product_dataframe['product_semantics'].apply(get_embedding)
8
9
  print("Embeddings generated successfully")
10
11
except Exception as e:
12
13
  print(f"Error generating embeddings: {e}")

1
# The field containing the text embeddings on each document
2
3
embedding_field_name = "embedding"
4
5
# MongoDB Atlas Vector Search index name
6
7
vector_search_index_name = "vector_index"

1
def setup_vector_search_index(collection, index_definition, index_name="vector_index"):
2
3
  """
4
5
  Setup a vector search index for a MongoDB collection and wait for 30 seconds.
6
7
  Args:
8
9
  collection: MongoDB collection object
10
11
  index_definition: Dictionary containing the index definition
12
13
  index_name: Name of the index (default: "vector_index")
14
15
  """
16
17
  new_vector_search_index_model = SearchIndexModel(
18
19
      definition=index_definition,
20
21
      name=index_name,
22
23
      type="vectorSearch"
24
25
  )
26
27
  # Create the new index
28
29
  try:
30
31
    result = collection.create_search_index(model=new_vector_search_index_model)
32
33
    print(f"Creating index '{index_name}'...")
34
35
    # Sleep for 30 seconds
36
37
    print(f"Waiting for 30 seconds to allow index '{index_name}' to be created...")
38
39
    time.sleep(30)
40
41
    print(f"30-second wait completed for index '{index_name}'.")
42
43
    return result
44
45
  except Exception as e:
46
47
    print(f"Error creating new vector search index '{index_name}': {str(e)}")
48
49
    return None

1
from pymongo import MongoClient
2
3
from tavily import TavilyHybridClient
4
5
hybrid_rag = TavilyHybridClient(
6
7
    api_key=os.environ.get("TAVILY_API_KEY"),
8
9
    db_provider="mongodb",
10
11
    collection=product_collection,
12
13
    index=vector_search_index_name,
14
15
    embeddings_field="embedding",
16
17
    content_field="product_semantics"
18
19
)

1
results = hybrid_rag.search("Get me a black laptop to use in a office", max_local=5, max_foreign=2)

1
results = hybrid_rag.search("Get me a black laptop to use in a office", max_local=5, max_foreign=2, save_foreign=True)