PolarDB:Build Long-Term and Short-Term Memory for AI Agents with PolarSearch Memory Container

Save the following as trip_agent.py. Set environment variables as shown in the comments before running.

<details> <summary>Click to expand the complete code</summary>

#!/usr/bin/env python3
"""
Trip Memory Agent - A simplified memory agent that tracks user information
and manages memories using OpenSearch as the backend.
"""
import os
import json
import requests
from requests.auth import HTTPBasicAuth
import urllib3
from typing import Optional, List, Dict, Any

# Disable HTTPS warnings for local development
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

# Import OpenAI library for LLM operations
import openai

# --- Configuration from Environment Variables ---
OPENSEARCH_URL = os.getenv('OPENSEARCH_URL')
OPENSEARCH_USERNAME = os.getenv('OPENSEARCH_USERNAME', 'admin')
OPENSEARCH_PASSWORD = os.getenv('OPENSEARCH_PASSWORD')
CONTAINER_ID = os.getenv('MEM_CONTAINER_ID')
QWEN_API_KEY = os.getenv('QWEN_API_KEY')
QWEN_BASE_URL = os.getenv('QWEN_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1')
QWEN_MODEL_NAME = os.getenv('QWEN_MODEL_NAME', 'qwen-max')

# Global variable to store current user ID
current_user_id = None

# Global OpenAI client instance
llm_client = None


# --- OpenSearch Helper ---
def opensearch_request(method: str, endpoint: str, json_data: Optional[Dict] = None) -> requests.Response:
    """Make an HTTP request to the OpenSearch API."""
    if not all([OPENSEARCH_URL, OPENSEARCH_USERNAME, OPENSEARCH_PASSWORD, CONTAINER_ID]):
        raise ValueError("OpenSearch environment variables are not fully configured.")

    url = f"{OPENSEARCH_URL}{endpoint}"
    auth = HTTPBasicAuth(OPENSEARCH_USERNAME, OPENSEARCH_PASSWORD)
    try:
        response = getattr(requests, method.lower())(url, json=json_data, auth=auth, verify=False, timeout=20)
        response.raise_for_status()
        return response
    except requests.exceptions.RequestException as e:
        print(f"Error communicating with OpenSearch: {e}")
        raise e


# --- User Management ---
def get_user_id():
    """Get user ID from global variable or prompt the user."""
    global current_user_id
    if not current_user_id:
        print("=" * 50)
        print("Welcome to the Trip Memory Agent!")
        print("=" * 50)
        print("Hello! I don't know you yet. Please tell me your name:")
        user_input = input("> ")
        current_user_id = user_input.strip()
        print(f"Nice to meet you, {current_user_id}!")
    return current_user_id


# --- LLM-Powered Helper Functions ---
def is_trip_plan_request(query: str) -> bool:
    """Determine if a user's query is a request for a trip plan using an LLM."""
    if not query.strip():
        return False
    try:
        prompt = f"""Please determine if the following user input is a request to plan a trip.
Answer only "YES" or "NO". Do not explain.
User input: {query}"""
        response = llm_client.chat.completions.create(
            model=QWEN_MODEL_NAME,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=10,
            temperature=0.0
        )
        result = response.choices[0].message.content.strip().upper()
        return result == "YES"
    except Exception as e:
        print(f"Error detecting trip plan request with LLM: {e}. Falling back to keywords.")
        keywords = ['plan a trip', 'trip plan', 'plan my travel', 'travel plan']
        return any(keyword in query.lower() for keyword in keywords)


def extract_keywords(query: str) -> List[str]:
    """Extract keywords from a query using an LLM for better search results."""
    try:
        prompt = f"""Extract the most important keywords from the following query for search purposes.
Focus on locations, activities, and specific preferences.
Return only a JSON array of strings (keywords).
Query: {query}"""
        response = llm_client.chat.completions.create(
            model=QWEN_MODEL_NAME,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=100,
            temperature=0.0
        )
        extracted_content = response.choices[0].message.content.strip()
        return json.loads(extracted_content)
    except (Exception, json.JSONDecodeError) as e:
        print(f"Error extracting keywords with LLM: {e}. Falling back to splitting query.")
        return query.split()


def generate_response_with_memories(query: str, search_results: List[Dict]) -> str:
    """Generate a response using an LLM with memory context."""
    if not search_results:
        return "I couldn't find any relevant memories for your query."

    context = "Here is relevant information I found in your memories:\n"
    for result in search_results:
        context += f"- {result['content']}\n"

    prompt = f"""You are a helpful trip planning assistant. Based on the user's past memories and their current request, generate a helpful and personalized response.

{context}
User's current request: {query}

Your response:"""
    try:
        response = llm_client.chat.completions.create(
            model=QWEN_MODEL_NAME,
            messages=[{"role": "system", "content": prompt}],
            max_tokens=1024,
            temperature=0.7
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        print(f"Error generating response with LLM: {e}")
        return "I found some memories, but I encountered an error while trying to generate a response."


# --- Core Memory Management Functions (Our "Tools") ---

def add_memory(message: str) -> Dict[str, Any]:
    """Add a new piece of information (a memory) for the current user."""
    global current_user_id
    payload = {
        "messages": [{"role": "user", "content": [{"text": message, "type": "text"}]}],
        "infer": True,
        "memory_type": "conversation",
        "namespace": {"user_id": current_user_id}
    }
    response = opensearch_request("POST", f"/_plugins/_ml/memory_containers/{CONTAINER_ID}/memories", payload)
    return response.json()


def search_memories(query: str) -> Dict[str, Any]:
    """Search the user's memories based on a query."""
    global current_user_id
    keywords = extract_keywords(query)

    search_query = {
        "query": {
            "bool": {
                "must": [{"term": {"namespace.user_id": current_user_id}}],
                "should": [
                    {"multi_match": {"query": " ".join(keywords), "fields": ["memory", "messages.content.text"]}}],
                "minimum_should_match": 1 if keywords else 0
            }
        },
        "sort": [{"created_time": {"order": "desc"}}],
        "size": 10
    }

    all_hits = []
    for index_suffix in ["long-term", "working"]:
        try:
            response = opensearch_request("POST",
                                          f"/_plugins/_ml/memory_containers/{CONTAINER_ID}/memories/{index_suffix}/_search",
                                          search_query)
            result = response.json()
            if "hits" in result and "hits" in result["hits"]:
                all_hits.extend(result["hits"]["hits"])
        except Exception as e:
            print(f"Error searching {index_suffix} memory: {e}")

    all_hits.sort(key=lambda x: x.get('_source', {}).get('created_time', ''), reverse=True)

    formatted_results = []
    for hit in all_hits:
        source = hit.get('_source', {})
        content = source.get('memory', '')
        if not content and 'messages' in source:
            content = " ".join(
                item['text'] for msg in source['messages'] if 'content' in msg
                for item in msg['content'] if isinstance(item, dict) and 'text' in item
            )

        index_type = "long-term" if 'long-term' in hit.get('_index', '') else "working"
        formatted_results.append({
            "memory_id": hit.get('_id'),
            "index_type": index_type,
            "timestamp": source.get('created_time'),
            "content": content.strip()
        })
    return {"status": "success", "results": formatted_results, "total_found": len(formatted_results)}


def find_and_update_memory(query: str, new_text: str) -> Dict[str, Any]:
    """Find a memory based on a query and update its content."""
    search_result = search_memories(query)
    if not search_result["results"]:
        return {"status": "error", "message": "No matching memory found to update."}

    memory_to_update = search_result["results"][0]
    memory_id = memory_to_update["memory_id"]
    index_type = memory_to_update["index_type"]

    payload = {"memory": new_text}
    response = opensearch_request("PUT",
                                  f"/_plugins/_ml/memory_containers/{CONTAINER_ID}/memories/{index_type}/{memory_id}",
                                  payload)
    return {"status": "success", "updated_memory_id": memory_id, "details": response.json()}


def find_and_delete_memory(query: str) -> Dict[str, Any]:
    """Find a memory based on a query and delete it."""
    search_result = search_memories(query)
    if not search_result["results"]:
        return {"status": "error", "message": "No matching memory found to delete."}

    memory_to_delete = search_result["results"][0]
    memory_id = memory_to_delete["memory_id"]
    index_type = memory_to_delete["index_type"]

    response = opensearch_request("DELETE",
                                  f"/_plugins/_ml/memory_containers/{CONTAINER_ID}/memories/{index_type}/{memory_id}")
    return {"status": "success", "deleted_memory_id": memory_id, "details": response.json()}


# --- Agent Logic ---

def main():
    """Main function to run the trip memory agent."""
    global llm_client

    if not all([QWEN_API_KEY, QWEN_BASE_URL, QWEN_MODEL_NAME]):
        print("Error: Qwen LLM environment variables are not configured.")
        return

    llm_client = openai.OpenAI(api_key=QWEN_API_KEY, base_url=QWEN_BASE_URL)
    get_user_id()

    tools = [
        {
            "type": "function",
            "function": {
                "name": "add_memory",
                "description": "Adds a new piece of information or a memory about the user. Use this to remember user preferences, facts, or past events.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "message": {
                            "type": "string",
                            "description": "The information or content of the memory to be saved. e.g., 'I love to visit historical museums.'"
                        }
                    },
                    "required": ["message"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "search_memories",
                "description": "Searches for and retrieves existing memories about the user. Use this when the user asks what you know about them, asks for past information, or wants a trip plan based on their preferences.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query to find relevant memories. e.g., 'my preferences', 'what do you know about me', 'trip to Paris'"
                        }
                    },
                    "required": ["query"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "find_and_update_memory",
                "description": "Finds a specific memory using a search query and updates its content with new text.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {"type": "string", "description": "A query to identify the memory to update. e.g., 'my favorite food'"},
                        "new_text": {"type": "string", "description": "The new content for the memory. e.g., 'My favorite food is now ramen.'"}
                    },
                    "required": ["query", "new_text"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "find_and_delete_memory",
                "description": "Finds a specific memory using a search query and deletes it.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {"type": "string", "description": "A query to identify the memory to delete. e.g., 'my old address'"}
                    },
                    "required": ["query"]
                }
            }
        }
    ]

    available_tools = {
        "add_memory": add_memory,
        "search_memories": search_memories,
        "find_and_update_memory": find_and_update_memory,
        "find_and_delete_memory": find_and_delete_memory,
    }

    system_prompt = f"""You are a trip assistant for a user named {current_user_id}. Your job is to help them by managing their memories and planning trips.
You have access to the following tools: `add_memory`, `search_memories`, `update_memory`, and `delete_memory`.
- When a user provides new information about themselves, use `add_memory`.
- When a user asks what you know about them, asks for their preferences, or wants a trip plan, use `search_memories` to get context first.
- When a user wants to change a piece of information, use `find_and_update_memory`.
- When a user wants to forget something, use `find_and_delete_memory`.
- Always respond to the user in a friendly and conversational manner. After a tool is used, explain what you did in simple terms."""

    messages = [{"role": "system", "content": system_prompt}]

    print("\nTrip Memory Agent is ready!")
    print(f"Hello {current_user_id}! I can remember things for you, search your memories, and help plan trips.")

    while True:
        try:
            user_input = input(f"\n{current_user_id}> ")
            if user_input.lower() in ['quit', 'exit', 'bye']:
                print("Goodbye!")
                break

            messages.append({"role": "user", "content": user_input})

            if is_trip_plan_request(user_input):
                print("This looks like a trip plan request. Searching memories for context...")
                search_result = search_memories(user_input)
                if search_result.get("total_found", 0) > 0:
                    print("Found relevant memories! Generating a personalized plan...")
                    response = generate_response_with_memories(user_input, search_result["results"])
                    print(f"\nAgent: {response}")
                else:
                    print("I couldn't find any relevant memories to help with that plan. Let's start from scratch!")
                messages.pop()
                continue

            print("Thinking...")
            response = llm_client.chat.completions.create(
                model=QWEN_MODEL_NAME,
                messages=messages,
                tools=tools,
                tool_choice="auto"
            )
            response_message = response.choices[0].message

            if response_message.tool_calls:
                messages.append(response_message)

                for tool_call in response_message.tool_calls:
                    function_name = tool_call.function.name
                    function_to_call = available_tools.get(function_name)

                    if not function_to_call:
                        print(f"Error: LLM tried to call an unknown function '{function_name}'")
                        continue

                    try:
                        function_args = json.loads(tool_call.function.arguments)
                        print(f"Calling tool: `{function_name}` with args: {function_args}")
                        function_response = function_to_call(**function_args)
                        messages.append({
                            "tool_call_id": tool_call.id,
                            "role": "tool",
                            "name": function_name,
                            "content": json.dumps(function_response)
                        })
                    except Exception as e:
                        print(f"Error executing tool '{function_name}': {e}")
                        messages.append({
                            "tool_call_id": tool_call.id,
                            "role": "tool",
                            "name": function_name,
                            "content": f'{{"status": "error", "message": "{str(e)}"}}'
                        })

                print("Summarizing tool results...")
                second_response = llm_client.chat.completions.create(
                    model=QWEN_MODEL_NAME,
                    messages=messages
                )
                final_response = second_response.choices[0].message.content
                print(f"\nAgent: {final_response}")
                messages.append({"role": "assistant", "content": final_response})

            else:
                final_response = response_message.content
                print(f"\nAgent: {final_response}")
                messages.append({"role": "assistant", "content": final_response})

        except Exception as e:
            print(f"\nAn unexpected error occurred: {e}")
            messages = [{"role": "system", "content": system_prompt}]


if __name__ == "__main__":
    # Set environment variables before running:
    # export OPENSEARCH_URL="https://..."
    # export OPENSEARCH_USERNAME="admin"
    # export OPENSEARCH_PASSWORD="your_password"
    # export MEM_CONTAINER_ID="your_container_id"
    # export QWEN_API_KEY="your_api_key"
    main()

</details>

1. Set the required environment variables:

   export OPENSEARCH_URL="https://pc-xxx.polardbsearch.rds.aliyuncs.com:3001"
   export OPENSEARCH_USERNAME="polarsearch_user"
   export OPENSEARCH_PASSWORD="your_password"
   export MEM_CONTAINER_ID="your_memory_container_id"
   export QWEN_API_KEY="your_api_key"

2. Run the script:

   python trip_agent.py

3. Example interaction:

<details> <summary>Click to expand an example interaction</summary>

==================================================
Welcome to the Trip Memory Agent!
==================================================
Hello! I don't know you yet. Please tell me your name:
> ABC
Nice to meet you, ABC!

Trip Memory Agent is ready!
Hello ABC! I can remember things for you, search your memories, and help plan trips.

ABC> i like hiking and travelling
Thinking...
Calling tool: `add_memory` with args: {'message': 'i like hiking and travelling'}
Summarizing tool results...

Agent: Great to know that you enjoy hiking and traveling, ABC! I've added this information to your profile. Whenever you're looking for trip ideas, I'll make sure to suggest some amazing hiking spots and travel destinations for you.

ABC> I like listen to music
Thinking...
Calling tool: `add_memory` with args: {'message': 'I like listen to music'}
Summarizing tool results...

Agent: Awesome, ABC! I've noted that you like listening to music. Whether it's for your travels or just everyday life, I can help you find some great playlists and music recommendations.

ABC> what do you know about me
Thinking...
Calling tool: `search_memories` with args: {'query': 'what do you know about me'}
Summarizing tool results...

Agent: I know a few things about you, ABC! Here's what I have:

1. You like hiking and traveling.
2. You enjoy listening to music.

If there's anything else you'd like to add, just let me know!

ABC> I live in Beijing
Thinking...
Calling tool: `add_memory` with args: {'message': 'I live in Beijing'}
Summarizing tool results...

Agent: Thanks for letting me know, ABC! I've added that you live in Beijing to your profile.

ABC> help me plan a weekend trip
This looks like a trip plan request. Searching memories for context...
Found relevant memories! Generating a personalized plan...

Agent: Given your love for hiking, traveling, and music, here's a weekend plan from Beijing:

**Yunmeng Mountain, Huairou District**

Day 1:
- Drive to Yunmeng Mountain (about 2 hours from central Beijing)
- Hike the scenic trails and enjoy a picnic lunch with your favorite music
- Stay overnight at a local guesthouse

Day 2:
- Short morning hike or walk around the area
- Brunch at a local restaurant (try Huairou roast duck)
- Return to Beijing in the afternoon

ABC> I live in Shanghai now
Thinking...
Calling tool: `find_and_update_memory` with args: {'query': 'I live in Beijing', 'new_text': 'I live in Shanghai now'}
Summarizing tool results...

Agent: Got it, ABC! I've updated your location from Beijing to Shanghai.

ABC> help me plan a weekend trip
This looks like a trip plan request. Searching memories for context...
Found relevant memories! Generating a personalized plan...

Agent: Since you now live in Shanghai and enjoy hiking and music, here's a weekend plan:

**Moganshan, Zhejiang Province** (about 3 hours from Shanghai)

Day 1:
- Drive to Moganshan in the morning
- Afternoon hike on scenic trails (Jianchi or Daken Scenic Area)
- Stay at a boutique hotel or guesthouse

Day 2:
- Morning walk near your accommodation
- Visit a local tea plantation for a tea-picking experience
- Return to Shanghai in the afternoon

Pack appropriate hiking gear and book accommodation in advance.

</details>