Shared Memory Between Agents
This example shows how to use one Agent
to load content into TaskMemory
and get that content from another Agent
using TaskMemoryClient
.
The first Agent
uses a remote vector store (MongoDbAtlasVectorStoreDriver
in this example) to handle memory operations. The second Agent
uses the same instance of TaskMemory
and the TaskMemoryClient
with the same MongoDbAtlasVectorStoreDriver
to get the data.
The MongoDbAtlasVectorStoreDriver
assumes that you have a vector index configured where the path to the content is called vector
, and the number of dimensions set on the index is 1536
(this is a commonly used number of dimensions for embedding models).
asker
uses the same instance of TaskMemory
as loader
so that asker
has access to the namespace_storages
that loader
has set.
import os
from griptape.tools import WebScraper, VectorStoreClient, TaskMemoryClient
from griptape.structures import Agent
from griptape.drivers import AzureOpenAiChatPromptDriver, AzureOpenAiEmbeddingDriver, AzureMongoDbVectorStoreDriver
from griptape.engines import VectorQueryEngine, PromptSummaryEngine, CsvExtractionEngine, JsonExtractionEngine
from griptape.memory import TaskMemory
from griptape.artifacts import TextArtifact
from griptape.memory.task.storage import TextArtifactStorage
from griptape.config import StructureConfig, StructureGlobalDriversConfig
AZURE_OPENAI_ENDPOINT_1 = os.environ["AZURE_OPENAI_ENDPOINT_1"]
AZURE_OPENAI_API_KEY_1 = os.environ["AZURE_OPENAI_API_KEY_1"]
MONGODB_HOST = os.environ["MONGODB_HOST"]
MONGODB_USERNAME = os.environ["MONGODB_USERNAME"]
MONGODB_PASSWORD = os.environ["MONGODB_PASSWORD"]
MONGODB_DATABASE_NAME = os.environ["MONGODB_DATABASE_NAME"]
MONGODB_COLLECTION_NAME = os.environ["MONGODB_COLLECTION_NAME"]
MONGODB_INDEX_NAME = os.environ["MONGODB_INDEX_NAME"]
MONGODB_VECTOR_PATH = os.environ["MONGODB_VECTOR_PATH"]
MONGODB_CONNECTION_STRING = f"mongodb+srv://{MONGODB_USERNAME}:{MONGODB_PASSWORD}@{MONGODB_HOST}/{MONGODB_DATABASE_NAME}?tls=true&authMechanism=SCRAM-SHA-256&retrywrites=false&maxIdleTimeMS=120000"
azure_embedding_driver = AzureOpenAiEmbeddingDriver(
model='text-embedding-ada-002',
azure_endpoint=AZURE_OPENAI_ENDPOINT_1,
api_key=AZURE_OPENAI_API_KEY_1,
azure_deployment='text-embedding-ada-002'
)
azure_prompt_driver = AzureOpenAiChatPromptDriver(
model='gpt-4',
azure_endpoint=AZURE_OPENAI_ENDPOINT_1,
api_key=AZURE_OPENAI_API_KEY_1,
azure_deployment='gpt-4'
)
mongo_driver = AzureMongoDbVectorStoreDriver(
connection_string=MONGODB_CONNECTION_STRING,
database_name=MONGODB_DATABASE_NAME,
collection_name=MONGODB_COLLECTION_NAME,
embedding_driver=azure_embedding_driver,
index_name=MONGODB_INDEX_NAME,
vector_path=MONGODB_VECTOR_PATH
)
loader = Agent(
tools=[
WebScraper()
],
config=StructureConfig(
global_drivers=StructureGlobalDriversConfig(
prompt_driver=azure_prompt_driver,
vector_store_driver=mongo_driver,
embedding_driver=azure_embedding_driver
)
),
)
asker = Agent(
tools=[
TaskMemoryClient(off_prompt=False),
],
meta_memory=loader.meta_memory,
task_memory=loader.task_memory,
)
if __name__ == "__main__":
loader.run("Load https://medium.com/enterprise-rag/a-first-intro-to-complex-rag-retrieval-augmented-generation-a8624d70090f")
asker.run("why is retrieval augmented generation useful?")