Langchain Tools

pip install langchain_core
pip install mem0ai

from langchain_core.tools import StructuredTool
from mem0 import MemoryClient
from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional
import os

os.environ["MEM0_API_KEY"] = "your-api-key"

client = MemoryClient()

class Message(BaseModel):
    role: str = Field(description="Role of the message sender (user or assistant)")
    content: str = Field(description="Content of the message")

class AddMemoryInput(BaseModel):
    messages: List[Message] = Field(description="List of messages to add to memory")
    user_id: str = Field(description="ID of the user associated with these messages")
    metadata: Optional[Dict[str, Any]] = Field(description="Additional metadata for the messages", default=None)

    class Config:
        json_schema_extra = {
            "examples": [{
                "messages": [
                    {"role": "user", "content": "Hi, I'm Alex. I'm a vegetarian and I'm allergic to nuts."},
                    {"role": "assistant", "content": "Hello Alex! I've noted that you're a vegetarian and have a nut allergy."}
                ],
                "user_id": "alex",
                "metadata": {"food": "vegan"}
            }]
        }

def add_memory(messages: List[Message], user_id: str, metadata: Optional[Dict[str, Any]] = None) -> Any:
    """Add messages to memory with associated user ID and metadata."""
    message_dicts = [msg.dict() for msg in messages]
    return client.add(message_dicts, user_id=user_id, metadata=metadata)

add_tool = StructuredTool(
    name="add_memory",
    description="Add new messages to memory with associated metadata",
    func=add_memory,
    args_schema=AddMemoryInput
)

{
  "results": [
    {
      "memory": "Name is Alex",
      "event": "ADD"
    },
    {
      "memory": "Is a vegetarian", 
      "event": "ADD"
    },
    {
      "memory": "Is allergic to nuts",
      "event": "ADD"
    }
  ]
}

class SearchMemoryInput(BaseModel):
    query: str = Field(description="The search query string")
    filters: Dict[str, Any] = Field(description="Filters to apply to the search")

    class Config:
        json_schema_extra = {
            "examples": [{
                "query": "tell me about my allergies?",
                "filters": {
                    "AND": [
                        {"user_id": "alex"},
                        {"created_at": {"gte": "2024-01-01", "lte": "2024-12-31"}}
                    ]
                }
            }]
        }

def search_memory(query: str, filters: Dict[str, Any]) -> Any:
    """Search memory with the given query and filters."""
    return client.search(query=query, filters=filters)

search_tool = StructuredTool(
    name="search_memory",
    description="Search through memories with a query and filters",
    func=search_memory,
    args_schema=SearchMemoryInput
)

[
  {
    "id": "1a75e827-7eca-45ea-8c5c-cfd43299f061",
    "memory": "Name is Alex",
    "user_id": "alex", 
    "hash": "d0fccc8fa47f7a149ee95750c37bb0ca",
    "metadata": {
      "food": "vegan"
    },
    "categories": [
      "personal_details"
    ],
    "created_at": "2024-11-27T16:53:43.276872-08:00",
    "updated_at": "2024-11-27T16:53:43.276885-08:00",
    "score": 0.3810526501504994
  }
]

class GetAllMemoryInput(BaseModel):
    filters: Dict[str, Any] = Field(description="Filters to apply to the retrieval")
    page: Optional[int] = Field(description="Page number for pagination", default=1)
    page_size: Optional[int] = Field(description="Number of items per page", default=50)

    class Config:
        json_schema_extra = {
            "examples": [{
                "filters": {
                    "AND": [
                        {"user_id": "alex"},
                        {"created_at": {"gte": "2024-07-01", "lte": "2024-07-31"}},
                        {"categories": {"contains": "food_preferences"}}
                    ]
                },
                "page": 1,
                "page_size": 50
            }]
        }

def get_all_memory(filters: Dict[str, Any], page: int = 1, page_size: int = 50) -> Any:
    """Retrieve all memories matching the specified criteria."""
    return client.get_all(filters=filters, page=page, page_size=page_size)

get_all_tool = StructuredTool(
    name="get_all_memory",
    description="Retrieve all memories matching specified filters",
    func=get_all_memory,
    args_schema=GetAllMemoryInput
)

{
  "count": 3,
  "next": null,
  "previous": null,
  "results": [
    {
      "id": "1a75e827-7eca-45ea-8c5c-cfd43299f061",
      "memory": "Name is Alex",
      "user_id": "alex", 
      "hash": "d0fccc8fa47f7a149ee95750c37bb0ca",
      "metadata": {
        "food": "vegan"
      },
      "categories": [
        "personal_details"
      ],
      "created_at": "2024-11-27T16:53:43.276872-08:00",
      "updated_at": "2024-11-27T16:53:43.276885-08:00"
    },
    {
      "id": "91509588-0b39-408a-8df3-84b3bce8c521",
      "memory": "Is a vegetarian",
      "user_id": "alex",
      "hash": "ce6b1c84586772ab9995a9477032df99", 
      "metadata": {
        "food": "vegan"
      },
      "categories": [
        "user_preferences",
        "food"
      ],
      "created_at": "2024-11-27T16:53:43.308027-08:00",
      "updated_at": "2024-11-27T16:53:43.308037-08:00"
    },
    {
      "id": "8d74f7a0-6107-4589-bd6f-210f6bf4fbbb",
      "memory": "Is allergic to nuts",
      "user_id": "alex",
      "hash": "7873cd0e5a29c513253d9fad038e758b",
      "metadata": {
        "food": "vegan"
      },
      "categories": [
        "health"
      ],
      "created_at": "2024-11-27T16:53:43.337253-08:00",
      "updated_at": "2024-11-27T16:53:43.337262-08:00"
    }
  ]
}