> For the complete documentation index, see [llms.txt](https://docs.canso.ai/llms.txt). Markdown versions of documentation pages are available by appending `.md` to page URLs; this page is available as [Markdown](https://docs.canso.ai/ai-agents/use-cases/agent-with-memory.md).

# Agent with Memory

This guide will walk you through setting up and deploying an AI agent with memory capabilities using the Canso AI Agentic System. We'll create a `memory-agent` that can store and retrieve information from a vector database, enhancing its ability to maintain context across interactions.

## Prerequisites

Before proceeding, ensure you have:

1. A [Canso compatible Kubernetes cluster](/getting-started/provision-k8s-cluster.md)set up
2. [Canso Helm charts](/getting-started/canso-helm-charts.md) installed on your cluster
3. Gru CLI installed (follow the instructions [here](/getting-started/canso-py-client.md))

## Setting Up Components

Our `memory-agent` requires several components:

* A **Vector Database** for memory storage
* A **Broker** for task orchestration
* A **Checkpoint DB** for saving agent state
* A **Task Server** for executing long-running tasks

First, create a file named `config.yaml` with the following configuration:

```yaml
vector_db:
  type: milvus
  name: my-vector-db
  size: 4Gi
  image_pull_secret: docker-secret-cred-agents

broker:
  type: redis
  name: my-redis

checkpoint_db:
  type: postgres
  name: my-postgres
  size: 4Gi

task_server:
  type: celery
  name: my-task-server
  replicas: 1
  concurrency_per_replica: 1
  broker_resource_name: my-redis
```

Next, run the command to set up all components:

```bash
gru component setup --cluster-name <name_of_your_cluster> --config-file config.yaml
```

## Creating the Project Bootstrap

Generate the scaffold for our project:

```bash
gru agent create_bootstrap
```

When prompted, provide the following inputs:

```
agent_name (Agent Name): memory-agent
agent_framework (Langgraph): Langgraph        
task_server_name: my-task-server
checkpoint_db_name: my-postgres
replicas (1): 1
vector_db_name: my-vector-db
```

This creates a folder named `memory-agent` with the following structure:

```
memory-agent/
├── .dockerignore
├── .env
├── Dockerfile
├── README.md
├── config.yaml
├── requirements.txt
└── src/
    └── main.py
```

## Updating Dependencies

Update `requirements.txt` with all necessary dependencies:

```
gru==0.0.1rc22.dev16
python-dotenv==1.0.1
langchain-openai==0.2.14
openai==1.14.1
boto3==1.35.95
```

## Developing the Memory-Enabled Agent

Replace the content of `src/main.py` with the following code:

```python
import os
from typing import Annotated, TypedDict
from dotenv import load_dotenv
from gru.agents.tools.langgraph.python_arguments import PythonArgumentsTool
from gru.agents.tools.langgraph import GitRepoContentRetriever, PythonCodeRunner, PythonRunStatusChecker, DataQualityTool, MemoryRetrievalTool
from langgraph.graph.message import AnyMessage, add_messages
from langchain_core.runnables import Runnable, RunnableConfig
from langgraph.prebuilt import tools_condition
from langgraph.graph import END, StateGraph, START
from langchain_openai import ChatOpenAI
from openai import OpenAI
from langchain_core.prompts import ChatPromptTemplate
from langgraph.prebuilt import ToolNode
from gru.agents import CansoLanggraphAgent
from gru.agents import CansoMemory
from gru.agents.tools.core.vector_db.vectordb_factory import VectorDBType
from gru.agents.tools.core.embeddings.embedding_factory import EmbeddingType

# This loads the environment variables from .env file.
# It is recommended to keep this as the first statement in main.py and not to be removed.
load_dotenv()

GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
OPEN_AI_KEY = os.getenv("OPENAI_API_KEY", "")

# Initialize OpenAI client for embeddings
embedding_client = OpenAI()

# Initialize CansoMemory with embedding client
memory = CansoMemory(
    client=embedding_client, 
    embedding_type=EmbeddingType.OPENAI, 
    vector_db_type=VectorDBType.MILVUS
)

# Initialize the LLM
model = ChatOpenAI(model="gpt-4o", temperature=0)

# Initialize tools
git_content_retriever = GitRepoContentRetriever(GITHUB_TOKEN)
python_argument_tool = PythonArgumentsTool(GITHUB_TOKEN)
python_code_runner = PythonCodeRunner(GITHUB_TOKEN)
python_run_status_tool = PythonRunStatusChecker()
check_data_quality = DataQualityTool()
memory_retrieval_tool = MemoryRetrievalTool(memory=memory)

# Group tools by sensitivity
safe_tools = [python_argument_tool, python_run_status_tool, check_data_quality, memory_retrieval_tool]
sensitive_tools = [git_content_retriever, python_code_runner]

sensitive_tool_names = {tool.name for tool in sensitive_tools}

class State(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]

class Assistant:
    def __init__(self, runnable: Runnable):
        self.runnable = runnable

    def __call__(self, state: State, config: RunnableConfig):
        while True:
            result = self.runnable.invoke(state)
            # If the LLM happens to return an empty response, we will re-prompt it
            # for an actual response.
            if not result.tool_calls and (
                not result.content
                or isinstance(result.content, list)
                and not result.content[0].get("text")
            ):
                messages = state["messages"] + [("user", "Respond with a real output.")]
                state = {**state, "messages": messages}
            else:
                break
        return {"messages": result}

def route_tools(state: State):
    """Route to different tool nodes based on the tool being called."""
    next_node = tools_condition(state)
    # If no tools are invoked, return END
    if next_node == END:
        return END

    ai_message = state["messages"][-1]
    # Handle the first tool call (assuming single tool calls)
    first_tool_call = ai_message.tool_calls[0]

    # Route to sensitive tools if the tool name is in sensitive_tool_names
    if first_tool_call["name"] in sensitive_tool_names:
        return "sensitive_tools"
    return "safe_tools"

# Define the system prompt
assistant_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a helpful assistant with memory capabilities. You can store and retrieve information "
            "from your memory to maintain context across conversations. "
            "Use the query_memory tool to search for relevant information stored in your memory. "
            "You can analyze data quality, run code from repositories, and perform other advanced tasks "
            "using the provided tools. Always ask for confirmation before running sensitive operations."
        ),
        ("placeholder", "{messages}"),
    ]
)

# Combine all tools and create the runnable
all_tools = safe_tools + sensitive_tools
assistant_runnable = assistant_prompt | model.bind_tools(all_tools)

# Create the graph
graph = StateGraph(State)

graph.add_node("assistant", Assistant(assistant_runnable))
graph.add_node("safe_tools", ToolNode(safe_tools))
graph.add_node("sensitive_tools", ToolNode(sensitive_tools))

graph.add_edge(START, "assistant")
graph.add_conditional_edges(
    "assistant",
    route_tools,
    ["safe_tools", "sensitive_tools", END]
)
graph.add_edge("safe_tools", "assistant")
graph.add_edge("sensitive_tools", "assistant")

# Initialize the Canso agent with memory
canso_agent = CansoLanggraphAgent(
    stateGraph=graph, 
    interrupt_before=["sensitive_tools"], 
    memory=memory
)
canso_agent.run()
```

## Setting Up Environment Variables

Create or update the `.env` file with the following content:

```
OPENAI_API_KEY=<your_openai_api_key>
GITHUB_TOKEN=<your_github_token>
```

## Building and Pushing the Docker Image

Build the Docker image and push it to your container registry:

```bash
docker build -t <your-account>/memory-agent:0.0.1 .
docker push <your-account>/memory-agent:0.0.1
```

## Registering and Deploying the Agent

Register and deploy your agent with:

```bash
# Register agent
gru agent register . --cluster-name <name_of_your_cluster> --image <your-account>/memory-agent:0.0.1 --image_pull_secret <image_pull_secret>

# Deploy agent
gru agent deploy memory-agent
```

## Storing Initial Memory Data

Before interacting with your agent, let's store some initial data in its memory.

Create a JSON file named `customer_data.json`:

```json
{
  "collection_name": "canso_domain_knowledge",
  "text": "Customer support policy: All premium customers receive priority support with a response time of maximum 4 hours. Standard customers receive support within 24 hours. For technical issues, we offer screen sharing sessions for premium customers. Refunds are processed within 7 business days.",
  "data": {
    "domain": "customer_support",
    "tags": ["support", "policy", "premium", "standard"],
    "metadata": {
      "version": "1.2",
      "last_updated": "2025-02-10",
      "created_by": "policy-team",
      "criticality": "high"
    }
  }
}
```

Insert this data into agent's memory

```bash
gru agent memory insert --agent-name memory-agent --file customer_data.json
```

This command will return the name of the memory where the data was inserted, which you can use for update and delete operations.

```bash
Memory registration request queued successfully with memory name = 8a0a3e62-de36-403f-a094-606c15f5a54f
```

## Updating existing Data in Memory

With this operation, you can also update certain attributes of the data you previously added.

Create a JSON file named customer\_update\_data.json:

```json
{
  "collection_name": "canso_domain_knowledge",
  "memory_name": "8a0a3e62-de36-403f-a094-606c15f5a54f",
  "text": "Customer support policy: All premium customers receive priority support with a response time of maximum 3 hours. Standard customers receive support within 48 hours. For technical issues, we offer screen sharing sessions and dedicated agents for premium customers. Refunds are processed within 5 business days.",
  "data": {
    "domain": "customer_success",
    "tags": ["support", "policy", "premium", "standard", "refund"],
    "metadata": {
      "version": "1.3",
      "last_updated": "2025-03-15",
      "created_by": "success-team",
      "criticality": "critical"
    }
  }
}
```

Update this data in the agent's memory.

```bash
gru agent memory update --agent-name memory-agent --file customer_update_data.json
```

## Delete the existing Data from Memory

You can also delete the previously stored elements from memory by.

```bash
gru agent memory delete --agent-name memory-agent --file customer_delete_data.json
```

## Interacting with the Memory-Enabled Agent

Start a conversation with your agent:

```bash
gru agent converse memory-agent
```

Example interaction:

```
Conversation ID: abcd1234
User: What do you know about our customer support policies?

Agent: I'll check my memory for information about customer support policies.

I've found some relevant information in my knowledge base about our customer support policies:

We have different service levels based on customer type:
- Premium customers receive priority support with a maximum response time of 4 hours
- Standard customers receive support within 24 hours
- Premium customers are eligible for screen sharing sessions for technical issues
- All refunds are processed within 7 business days

This information was last updated on March 15, 2025 (version 1.2) and is considered high-criticality business information.

Is there anything specific about these policies you'd like me to elaborate on?

User: Can you also run a data quality check on our customer database?

Agent: I'll be calling the tool check_data_quality with following arguments:
bucket: customer-data-bucket
file_path: customers.csv
Do you approve of this action? Type 'y' to continue; otherwise, explain your requested changed.

User: y

Agent: [Results of data quality check would appear here]
```

## Additional Resources

* [Memory API Reference](/ai-agents/api-summary/memory-api.md)
* [Memory Examples](/ai-agents/use-cases/examples.md)
* [SQL Runner Tool](/ai-agents/toolkit/sql-runner.md)