Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e77a10a8da | ||
|
|
e4307ae420 | ||
|
|
7c89d00079 | ||
|
|
563eaae5ee | ||
|
|
6733f78f81 | ||
|
|
c11637bd2f | ||
|
|
ff30cb8ddd | ||
|
|
2e853c3d22 | ||
|
|
3cc7013fde | ||
|
|
8e6a08aa83 | ||
|
|
8b9a8e5825 |
@@ -84,6 +84,7 @@ Here's a comprehensive list of all parameters that can be used across different
|
||||
| `memory_add_embedding_type` | The type of embedding to use for the add memory action | VertexAI |
|
||||
| `memory_update_embedding_type` | The type of embedding to use for the update memory action | VertexAI |
|
||||
| `memory_search_embedding_type` | The type of embedding to use for the search memory action | VertexAI |
|
||||
| `lmstudio_base_url` | Base URL for LM Studio API | LM Studio |
|
||||
</Tab>
|
||||
<Tab title="TypeScript">
|
||||
| Parameter | Description | Provider |
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
You can use embedding models from LM Studio to run Mem0 locally.
|
||||
|
||||
### Usage
|
||||
|
||||
```python
|
||||
import os
|
||||
from mem0 import Memory
|
||||
|
||||
os.environ["OPENAI_API_KEY"] = "your_api_key" # For LLM
|
||||
|
||||
config = {
|
||||
"embedder": {
|
||||
"provider": "lmstudio",
|
||||
"config": {
|
||||
"model": "nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
m = Memory.from_config(config)
|
||||
messages = [
|
||||
{"role": "user", "content": "I'm planning to watch a movie tonight. Any recommendations?"},
|
||||
{"role": "assistant", "content": "How about a thriller movies? They can be quite engaging."},
|
||||
{"role": "user", "content": "I’m not a big fan of thriller movies but I love sci-fi movies."},
|
||||
{"role": "assistant", "content": "Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future."}
|
||||
]
|
||||
m.add(messages, user_id="john")
|
||||
```
|
||||
|
||||
### Config
|
||||
|
||||
Here are the parameters available for configuring Ollama embedder:
|
||||
|
||||
| Parameter | Description | Default Value |
|
||||
| --- | --- | --- |
|
||||
| `model` | The name of the OpenAI model to use | `nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf` |
|
||||
| `embedding_dims` | Dimensions of the embedding model | `1536` |
|
||||
| `lmstudio_base_url` | Base URL for LM Studio connection | `http://localhost:1234/v1` |
|
||||
@@ -22,6 +22,7 @@ See the list of supported embedders below.
|
||||
<Card title="Gemini" href="/components/embedders/models/gemini"></Card>
|
||||
<Card title="Vertex AI" href="/components/embedders/models/vertexai"></Card>
|
||||
<Card title="Together" href="/components/embedders/models/together"></Card>
|
||||
<Card title="LM Studio" href="/components/embedders/models/lmstudio"></Card>
|
||||
</CardGroup>
|
||||
|
||||
## Usage
|
||||
|
||||
@@ -108,6 +108,7 @@ Here's a comprehensive list of all parameters that can be used across different
|
||||
| `azure_kwargs` | Azure LLM args for initialization | AzureOpenAI |
|
||||
| `deepseek_base_url` | Base URL for DeepSeek API | DeepSeek |
|
||||
| `xai_base_url` | Base URL for XAI API | XAI |
|
||||
| `lmstudio_base_url` | Base URL for LM Studio API | LM Studio |
|
||||
</Tab>
|
||||
<Tab title="TypeScript">
|
||||
| Parameter | Description | Provider |
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
---
|
||||
title: LM Studio
|
||||
---
|
||||
|
||||
To use LM Studio with Mem0, you'll need to have LM Studio running locally with its server enabled. LM Studio provides a way to run local LLMs with an OpenAI-compatible API.
|
||||
|
||||
## Usage
|
||||
|
||||
<CodeGroup>
|
||||
```python Python
|
||||
import os
|
||||
from mem0 import Memory
|
||||
|
||||
os.environ["OPENAI_API_KEY"] = "your-api-key" # used for embedding model
|
||||
|
||||
config = {
|
||||
"llm": {
|
||||
"provider": "lmstudio",
|
||||
"config": {
|
||||
"model": "lmstudio-community/Meta-Llama-3.1-70B-Instruct-GGUF/Meta-Llama-3.1-70B-Instruct-IQ2_M.gguf",
|
||||
"temperature": 0.2,
|
||||
"max_tokens": 2000,
|
||||
"lmstudio_base_url": "http://localhost:1234/v1", # default LM Studio API URL
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
m = Memory.from_config(config)
|
||||
messages = [
|
||||
{"role": "user", "content": "I'm planning to watch a movie tonight. Any recommendations?"},
|
||||
{"role": "assistant", "content": "How about a thriller movies? They can be quite engaging."},
|
||||
{"role": "user", "content": "I'm not a big fan of thriller movies but I love sci-fi movies."},
|
||||
{"role": "assistant", "content": "Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future."}
|
||||
]
|
||||
m.add(messages, user_id="alice", metadata={"category": "movies"})
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
### Running Completely Locally
|
||||
|
||||
You can also use LM Studio for both LLM and embedding to run Mem0 entirely locally:
|
||||
|
||||
```python
|
||||
from mem0 import Memory
|
||||
|
||||
# No external API keys needed!
|
||||
config = {
|
||||
"llm": {
|
||||
"provider": "lmstudio"
|
||||
},
|
||||
"embedder": {
|
||||
"provider": "lmstudio"
|
||||
}
|
||||
}
|
||||
|
||||
m = Memory.from_config(config)
|
||||
messages = [
|
||||
{"role": "user", "content": "I'm planning to watch a movie tonight. Any recommendations?"},
|
||||
{"role": "assistant", "content": "How about a thriller movies? They can be quite engaging."},
|
||||
{"role": "user", "content": "I'm not a big fan of thriller movies but I love sci-fi movies."},
|
||||
{"role": "assistant", "content": "Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future."}
|
||||
]
|
||||
m.add(messages, user_id="alice123", metadata={"category": "movies"})
|
||||
```
|
||||
|
||||
<Note>
|
||||
When using LM Studio for both LLM and embedding, make sure you have:
|
||||
1. An LLM model loaded for generating responses
|
||||
2. An embedding model loaded for vector embeddings
|
||||
3. The server enabled with the correct endpoints accessible
|
||||
</Note>
|
||||
|
||||
<Note>
|
||||
To use LM Studio, you need to:
|
||||
1. Download and install [LM Studio](https://lmstudio.ai/)
|
||||
2. Start a local server from the "Server" tab
|
||||
3. Set the appropriate `lmstudio_base_url` in your configuration (default is usually http://localhost:1234/v1)
|
||||
</Note>
|
||||
|
||||
## Config
|
||||
|
||||
All available parameters for the `lmstudio` config are present in [Master List of All Params in Config](../config).
|
||||
@@ -32,6 +32,7 @@ To view all supported llms, visit the [Supported LLMs](./models).
|
||||
<Card title="Gemini" href="/components/llms/models/gemini" />
|
||||
<Card title="DeepSeek" href="/components/llms/models/deepseek" />
|
||||
<Card title="xAI" href="/components/llms/models/xAI" />
|
||||
<Card title="LM Studio" href="/components/llms/models/lmstudio" />
|
||||
</CardGroup>
|
||||
|
||||
## Structured vs Unstructured Outputs
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
# Azure AI Search
|
||||
|
||||
[Azure AI Search](https://learn.microsoft.com/azure/search/search-what-is-azure-search/) (formerly known as "Azure Cognitive Search") provides secure information retrieval at scale over user-owned content in traditional and generative AI search applications.
|
||||
|
||||
## Usage
|
||||
@@ -50,6 +48,24 @@ config = {
|
||||
}
|
||||
```
|
||||
|
||||
## Using hybrid search
|
||||
|
||||
```python
|
||||
config = {
|
||||
"vector_store": {
|
||||
"provider": "azure_ai_search",
|
||||
"config": {
|
||||
"service_name": "ai-search-test",
|
||||
"api_key": "*****",
|
||||
"collection_name": "mem0",
|
||||
"embedding_model_dims": 1536,
|
||||
"hybrid_search": True,
|
||||
"vector_filter_mode": "postFilter"
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Configuration Parameters
|
||||
|
||||
| Parameter | Description | Default Value | Options |
|
||||
@@ -60,6 +76,8 @@ config = {
|
||||
| `embedding_model_dims` | Dimensions of the embedding model | `1536` | Any integer value |
|
||||
| `compression_type` | Type of vector compression to use | `none` | `none`, `scalar`, `binary` |
|
||||
| `use_float16` | Store vectors in half precision (Edm.Half) | `False` | `True`, `False` |
|
||||
| `vector_filter_mode` | Vector filter mode to use | `preFilter` | `postFilter`, `preFilter` |
|
||||
| `hybrid_search` | Use hybrid search | `False` | `True`, `False` |
|
||||
|
||||
## Notes on Configuration Options
|
||||
|
||||
@@ -68,6 +86,10 @@ config = {
|
||||
- `scalar`: Scalar quantization with reasonable balance of speed and accuracy
|
||||
- `binary`: Binary quantization for maximum compression with some accuracy trade-off
|
||||
|
||||
- **vector_filter_mode**:
|
||||
- `preFilter`: Applies filters before vector search (faster)
|
||||
- `postFilter`: Applies filters after vector search (may provide better relevance)
|
||||
|
||||
- **use_float16**: Using half precision (float16) reduces storage requirements but may slightly impact accuracy. Useful for very large vector collections.
|
||||
|
||||
- **Filterable Fields**: The implementation automatically extracts `user_id`, `run_id`, and `agent_id` fields from payloads for filtering.
|
||||
@@ -1,7 +1,7 @@
|
||||
# Pinecone
|
||||
|
||||
[Pinecone](https://www.pinecone.io/) is a fully managed vector database designed for machine learning applications, offering high performance vector search with low latency at scale. It's particularly well-suited for semantic search, recommendation systems, and other AI-powered applications.
|
||||
|
||||
> **Note**: Before configuring Pinecone, you need to select an embedding model (e.g., OpenAI, Cohere, or custom models) and ensure the `embedding_model_dims` in your config matches your chosen model's dimensions. For example, OpenAI's text-embedding-ada-002 uses 1536 dimensions.
|
||||
|
||||
### Usage
|
||||
|
||||
```python
|
||||
@@ -11,13 +11,17 @@ from mem0 import Memory
|
||||
os.environ["OPENAI_API_KEY"] = "sk-xx"
|
||||
os.environ["PINECONE_API_KEY"] = "your-api-key"
|
||||
|
||||
# Example using serverless configuration
|
||||
config = {
|
||||
"vector_store": {
|
||||
"provider": "pinecone",
|
||||
"config": {
|
||||
"collection_name": "memory_index",
|
||||
"embedding_model_dims": 1536,
|
||||
"environment": "us-west1-gcp",
|
||||
"collection_name": "testing",
|
||||
"embedding_model_dims": 1536, # Matches OpenAI's text-embedding-3-small
|
||||
"serverless_config": {
|
||||
"cloud": "aws", # Choose between 'aws' or 'gcp' or 'azure'
|
||||
"region": "us-east-1"
|
||||
},
|
||||
"metric": "cosine"
|
||||
}
|
||||
}
|
||||
@@ -40,28 +44,29 @@ Here are the parameters available for configuring Pinecone:
|
||||
| Parameter | Description | Default Value |
|
||||
| --- | --- | --- |
|
||||
| `collection_name` | Name of the index/collection | Required |
|
||||
| `embedding_model_dims` | Dimensions of the embedding model | Required |
|
||||
| `embedding_model_dims` | Dimensions of the embedding model (must match your chosen embedding model) | Required |
|
||||
| `client` | Existing Pinecone client instance | `None` |
|
||||
| `api_key` | API key for Pinecone | Environment variable: `PINECONE_API_KEY` |
|
||||
| `environment` | Pinecone environment | `None` |
|
||||
| `serverless_config` | Configuration for serverless deployment | `None` |
|
||||
| `serverless_config` | Configuration for serverless deployment (AWS or GCP or Azure) | `None` |
|
||||
| `pod_config` | Configuration for pod-based deployment | `None` |
|
||||
| `hybrid_search` | Whether to enable hybrid search | `False` |
|
||||
| `metric` | Distance metric for vector similarity | `"cosine"` |
|
||||
| `batch_size` | Batch size for operations | `100` |
|
||||
|
||||
#### Serverless Config Example
|
||||
> **Important**: You must choose either `serverless_config` or `pod_config` for your deployment, but not both.
|
||||
|
||||
#### Serverless Config Example
|
||||
```python
|
||||
config = {
|
||||
"vector_store": {
|
||||
"provider": "pinecone",
|
||||
"config": {
|
||||
"collection_name": "memory_index",
|
||||
"embedding_model_dims": 1536,
|
||||
"embedding_model_dims": 1536, # For OpenAI's text-embedding-3-small
|
||||
"serverless_config": {
|
||||
"cloud": "aws",
|
||||
"region": "us-west-2"
|
||||
"cloud": "aws", # or "gcp" or "azure"
|
||||
"region": "us-east-1" # Choose appropriate region
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -69,14 +74,13 @@ config = {
|
||||
```
|
||||
|
||||
#### Pod Config Example
|
||||
|
||||
```python
|
||||
config = {
|
||||
"vector_store": {
|
||||
"provider": "pinecone",
|
||||
"config": {
|
||||
"collection_name": "memory_index",
|
||||
"embedding_model_dims": 1536,
|
||||
"embedding_model_dims": 1536, # For OpenAI's text-embedding-ada-002
|
||||
"pod_config": {
|
||||
"environment": "gcp-starter",
|
||||
"replicas": 1,
|
||||
|
||||
+2
-1
@@ -193,7 +193,8 @@
|
||||
"examples/multimodal-demo",
|
||||
"examples/personalized-deep-research",
|
||||
"examples/mem0-agentic-tool",
|
||||
"examples/openai-inbuilt-tools"
|
||||
"examples/openai-inbuilt-tools",
|
||||
"examples/mem0-openai-voice-demo"
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
@@ -0,0 +1,538 @@
|
||||
---
|
||||
title: 'Mem0 with OpenAI Agents SDK for Voice'
|
||||
description: 'Integrate memory capabilities into your voice agents using Mem0 and OpenAI Agents SDK'
|
||||
---
|
||||
|
||||
# Building Voice Agents with Memory using Mem0 and OpenAI Agents SDK
|
||||
|
||||
This guide demonstrates how to combine OpenAI's Agents SDK for voice applications with Mem0's memory capabilities to create a voice assistant that remembers user preferences and past interactions.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Before you begin, make sure you have:
|
||||
|
||||
1. Installed OpenAI Agents SDK with voice dependencies:
|
||||
```bash
|
||||
pip install 'openai-agents[voice]'
|
||||
```
|
||||
|
||||
2. Installed Mem0 SDK:
|
||||
```bash
|
||||
pip install mem0ai
|
||||
```
|
||||
|
||||
3. Installed other required dependencies:
|
||||
```bash
|
||||
pip install numpy sounddevice pydantic
|
||||
```
|
||||
|
||||
4. Set up your API keys:
|
||||
- OpenAI API key for the Agents SDK
|
||||
- Mem0 API key from the Mem0 Platform
|
||||
|
||||
## Code Breakdown
|
||||
|
||||
Let's break down the key components of this implementation:
|
||||
|
||||
### 1. Setting Up Dependencies and Environment
|
||||
|
||||
```python
|
||||
# OpenAI Agents SDK imports
|
||||
from agents import (
|
||||
Agent,
|
||||
function_tool
|
||||
)
|
||||
from agents.voice import (
|
||||
AudioInput,
|
||||
SingleAgentVoiceWorkflow,
|
||||
VoicePipeline
|
||||
)
|
||||
from agents.extensions.handoff_prompt import prompt_with_handoff_instructions
|
||||
|
||||
# Mem0 imports
|
||||
from mem0 import AsyncMemoryClient
|
||||
|
||||
# Set up API keys (replace with your actual keys)
|
||||
os.environ["OPENAI_API_KEY"] = "your-openai-api-key"
|
||||
os.environ["MEM0_API_KEY"] = "your-mem0-api-key"
|
||||
|
||||
# Define a global user ID for simplicity
|
||||
USER_ID = "voice_user"
|
||||
|
||||
# Initialize Mem0 client
|
||||
mem0_client = AsyncMemoryClient()
|
||||
```
|
||||
|
||||
This section handles:
|
||||
- Importing required modules from OpenAI Agents SDK and Mem0
|
||||
- Setting up environment variables for API keys
|
||||
- Defining a simple user identification system (using a global variable)
|
||||
- Initializing the Mem0 client that will handle memory operations
|
||||
|
||||
### 2. Memory Tools with Function Decorators
|
||||
|
||||
The `@function_tool` decorator transforms Python functions into callable tools for the OpenAI agent. Here are the key memory tools:
|
||||
|
||||
#### Storing User Memories
|
||||
|
||||
```python
|
||||
import logging
|
||||
|
||||
# Set up logging at the top of your file
|
||||
logging.basicConfig(
|
||||
level=logging.DEBUG,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
||||
force=True
|
||||
)
|
||||
logger = logging.getLogger("memory_voice_agent")
|
||||
|
||||
# Then use logger in your function tools
|
||||
@function_tool
|
||||
async def save_memories(
|
||||
memory: str
|
||||
) -> str:
|
||||
"""Store a user memory in memory."""
|
||||
# This will be visible in your console
|
||||
logger.debug(f"Saving memory: {memory} for user {USER_ID}")
|
||||
|
||||
# Store the preference in Mem0
|
||||
memory_content = f"User memory - {memory}"
|
||||
await mem0_client.add(
|
||||
memory_content,
|
||||
user_id=USER_ID,
|
||||
)
|
||||
|
||||
return f"I've saved your memory: {memory}"
|
||||
```
|
||||
|
||||
This function:
|
||||
- Takes a memory string
|
||||
- Creates a formatted memory string
|
||||
- Stores it in Mem0 using the `add()` method
|
||||
- Includes metadata to categorize the memory for easier retrieval
|
||||
- Returns a confirmation message that the agent will speak
|
||||
|
||||
#### Finding Relevant Memories
|
||||
|
||||
```python
|
||||
@function_tool
|
||||
async def search_memories(
|
||||
query: str
|
||||
) -> str:
|
||||
"""
|
||||
Find memories relevant to the current conversation.
|
||||
Args:
|
||||
query: The search query to find relevant memories
|
||||
"""
|
||||
print(f"Finding memories related to: {query}")
|
||||
results = await mem0_client.search(
|
||||
query,
|
||||
user_id=USER_ID,
|
||||
limit=5,
|
||||
threshold=0.7, # Higher threshold for more relevant results
|
||||
output_format="v1.1"
|
||||
)
|
||||
|
||||
# Format and return the results
|
||||
if not results.get('results', []):
|
||||
return "I don't have any relevant memories about this topic."
|
||||
|
||||
memories = [f"• {result['memory']}" for result in results.get('results', [])]
|
||||
return "Here's what I remember that might be relevant:\n" + "\n".join(memories)
|
||||
```
|
||||
|
||||
This tool:
|
||||
- Takes a search query string
|
||||
- Passes it to Mem0's semantic search to find related memories
|
||||
- Sets a threshold for relevance to ensure quality results
|
||||
- Returns a formatted list of relevant memories or a default message
|
||||
|
||||
### 3. Creating the Voice Agent
|
||||
|
||||
```python
|
||||
def create_memory_voice_agent():
|
||||
# Create the agent with memory-enabled tools
|
||||
agent = Agent(
|
||||
name="Memory Assistant",
|
||||
instructions=prompt_with_handoff_instructions(
|
||||
"""You're speaking to a human, so be polite and concise.
|
||||
Always respond in clear, natural English.
|
||||
You have the ability to remember information about the user.
|
||||
Use the save_memories tool when the user shares an important information worth remembering.
|
||||
Use the search_memories tool when you need context from past conversations or user asks you to recall something.
|
||||
""",
|
||||
),
|
||||
model="gpt-4o",
|
||||
tools=[save_memories, search_memories],
|
||||
)
|
||||
|
||||
return agent
|
||||
```
|
||||
|
||||
This function:
|
||||
- Creates an OpenAI Agent with specific instructions
|
||||
- Configures it to use gpt-4o (you can use other models)
|
||||
- Registers the memory-related tools with the agent
|
||||
- Uses `prompt_with_handoff_instructions` to include standard voice agent behaviors
|
||||
|
||||
### 4. Microphone Recording Functionality
|
||||
|
||||
```python
|
||||
async def record_from_microphone(duration=5, samplerate=24000):
|
||||
"""Record audio from the microphone for a specified duration."""
|
||||
print(f"Recording for {duration} seconds...")
|
||||
|
||||
# Create a buffer to store the recorded audio
|
||||
frames = []
|
||||
|
||||
# Callback function to store audio data
|
||||
def callback(indata, frames_count, time_info, status):
|
||||
frames.append(indata.copy())
|
||||
|
||||
# Start recording
|
||||
with sd.InputStream(samplerate=samplerate, channels=1, callback=callback, dtype=np.int16):
|
||||
await asyncio.sleep(duration)
|
||||
|
||||
# Combine all frames into a single numpy array
|
||||
audio_data = np.concatenate(frames)
|
||||
return audio_data
|
||||
```
|
||||
|
||||
This function:
|
||||
- Creates a simple asynchronous microphone recording function
|
||||
- Uses the sounddevice library to capture audio input
|
||||
- Stores frames in a buffer during recording
|
||||
- Combines frames into a single numpy array when complete
|
||||
- Returns the audio data for processing
|
||||
|
||||
### 5. Main Loop and Voice Processing
|
||||
|
||||
```python
|
||||
async def main():
|
||||
# Create the agent
|
||||
agent = create_memory_voice_agent()
|
||||
|
||||
# Set up the voice pipeline
|
||||
pipeline = VoicePipeline(
|
||||
workflow=SingleAgentVoiceWorkflow(agent)
|
||||
)
|
||||
|
||||
# Configure TTS settings
|
||||
pipeline.config.tts_settings.voice = "alloy"
|
||||
pipeline.config.tts_settings.speed = 1.0
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Get user input
|
||||
print("\nPress Enter to start recording (or 'q' to quit)...")
|
||||
user_input = input()
|
||||
if user_input.lower() == 'q':
|
||||
break
|
||||
|
||||
# Record and process audio
|
||||
audio_data = await record_from_microphone(duration=5)
|
||||
audio_input = AudioInput(buffer=audio_data)
|
||||
result = await pipeline.run(audio_input)
|
||||
|
||||
# Play response and handle events
|
||||
player = sd.OutputStream(samplerate=24000, channels=1, dtype=np.int16)
|
||||
player.start()
|
||||
|
||||
agent_response = ""
|
||||
print("\nAgent response:")
|
||||
|
||||
async for event in result.stream():
|
||||
if event.type == "voice_stream_event_audio":
|
||||
player.write(event.data)
|
||||
elif event.type == "voice_stream_event_content":
|
||||
content = event.data
|
||||
agent_response += content
|
||||
print(content, end="", flush=True)
|
||||
|
||||
# Save the agent's response to memory
|
||||
if agent_response:
|
||||
try:
|
||||
await mem0_client.add(
|
||||
f"Agent response: {agent_response}",
|
||||
user_id=USER_ID,
|
||||
metadata={"type": "agent_response"}
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Failed to store memory: {e}")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\nExiting...")
|
||||
```
|
||||
|
||||
This main function orchestrates the entire process:
|
||||
1. Creates the memory-enabled voice agent
|
||||
2. Sets up the voice pipeline with TTS settings
|
||||
3. Implements an interactive loop for recording and processing voice input
|
||||
4. Handles streaming of response events (both audio and text)
|
||||
5. Automatically saves the agent's responses to memory
|
||||
6. Includes proper error handling and exit mechanisms
|
||||
|
||||
## Create a Memory-Enabled Voice Agent
|
||||
|
||||
Now that we've explained each component, here's the complete implementation that combines OpenAI Agents SDK for voice with Mem0's memory capabilities:
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
import os
|
||||
import logging
|
||||
from typing import Optional, List, Dict, Any
|
||||
import numpy as np
|
||||
import sounddevice as sd
|
||||
from pydantic import BaseModel
|
||||
|
||||
# OpenAI Agents SDK imports
|
||||
from agents import (
|
||||
Agent,
|
||||
function_tool
|
||||
)
|
||||
from agents.voice import (
|
||||
AudioInput,
|
||||
SingleAgentVoiceWorkflow,
|
||||
VoicePipeline
|
||||
)
|
||||
from agents.extensions.handoff_prompt import prompt_with_handoff_instructions
|
||||
|
||||
# Mem0 imports
|
||||
from mem0 import AsyncMemoryClient
|
||||
|
||||
# Set up API keys (replace with your actual keys)
|
||||
os.environ["OPENAI_API_KEY"] = "your-openai-api-key"
|
||||
os.environ["MEM0_API_KEY"] = "your-mem0-api-key"
|
||||
|
||||
# Define a global user ID for simplicity
|
||||
USER_ID = "voice_user"
|
||||
|
||||
# Initialize Mem0 client
|
||||
mem0_client = AsyncMemoryClient()
|
||||
|
||||
# Create tools that utilize Mem0's memory
|
||||
@function_tool
|
||||
async def save_memories(
|
||||
memory: str
|
||||
) -> str:
|
||||
"""
|
||||
Store a user memory in memory.
|
||||
Args:
|
||||
memory: The memory to save
|
||||
"""
|
||||
print(f"Saving memory: {memory} for user {USER_ID}")
|
||||
|
||||
# Store the preference in Mem0
|
||||
memory_content = f"User memory - {memory}"
|
||||
await mem0_client.add(
|
||||
memory_content,
|
||||
user_id=USER_ID,
|
||||
)
|
||||
|
||||
return f"I've saved your memory: {memory}"
|
||||
|
||||
@function_tool
|
||||
async def search_memories(
|
||||
query: str
|
||||
) -> str:
|
||||
"""
|
||||
Find memories relevant to the current conversation.
|
||||
Args:
|
||||
query: The search query to find relevant memories
|
||||
"""
|
||||
print(f"Finding memories related to: {query}")
|
||||
results = await mem0_client.search(
|
||||
query,
|
||||
user_id=USER_ID,
|
||||
limit=5,
|
||||
threshold=0.7, # Higher threshold for more relevant results
|
||||
output_format="v1.1"
|
||||
)
|
||||
|
||||
# Format and return the results
|
||||
if not results.get('results', []):
|
||||
return "I don't have any relevant memories about this topic."
|
||||
|
||||
memories = [f"• {result['memory']}" for result in results.get('results', [])]
|
||||
return "Here's what I remember that might be relevant:\n" + "\n".join(memories)
|
||||
|
||||
# Create the agent with memory-enabled tools
|
||||
def create_memory_voice_agent():
|
||||
# Create the agent with memory-enabled tools
|
||||
agent = Agent(
|
||||
name="Memory Assistant",
|
||||
instructions=prompt_with_handoff_instructions(
|
||||
"""You're speaking to a human, so be polite and concise.
|
||||
Always respond in clear, natural English.
|
||||
You have the ability to remember information about the user.
|
||||
Use the save_memories tool when the user shares an important information worth remembering.
|
||||
Use the search_memories tool when you need context from past conversations or user asks you to recall something.
|
||||
""",
|
||||
),
|
||||
model="gpt-4o",
|
||||
tools=[save_memories, search_memories],
|
||||
)
|
||||
|
||||
return agent
|
||||
|
||||
async def record_from_microphone(duration=5, samplerate=24000):
|
||||
"""Record audio from the microphone for a specified duration."""
|
||||
print(f"Recording for {duration} seconds...")
|
||||
|
||||
# Create a buffer to store the recorded audio
|
||||
frames = []
|
||||
|
||||
# Callback function to store audio data
|
||||
def callback(indata, frames_count, time_info, status):
|
||||
frames.append(indata.copy())
|
||||
|
||||
# Start recording
|
||||
with sd.InputStream(samplerate=samplerate, channels=1, callback=callback, dtype=np.int16):
|
||||
await asyncio.sleep(duration)
|
||||
|
||||
# Combine all frames into a single numpy array
|
||||
audio_data = np.concatenate(frames)
|
||||
return audio_data
|
||||
|
||||
async def main():
|
||||
print("Starting Memory Voice Agent")
|
||||
|
||||
# Create the agent and context
|
||||
agent = create_memory_voice_agent()
|
||||
|
||||
# Set up the voice pipeline
|
||||
pipeline = VoicePipeline(
|
||||
workflow=SingleAgentVoiceWorkflow(agent)
|
||||
)
|
||||
|
||||
# Configure TTS settings
|
||||
pipeline.config.tts_settings.voice = "alloy"
|
||||
pipeline.config.tts_settings.speed = 1.0
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Get user input
|
||||
print("\nPress Enter to start recording (or 'q' to quit)...")
|
||||
user_input = input()
|
||||
if user_input.lower() == 'q':
|
||||
break
|
||||
|
||||
# Record and process audio
|
||||
audio_data = await record_from_microphone(duration=5)
|
||||
audio_input = AudioInput(buffer=audio_data)
|
||||
|
||||
print("Processing your request...")
|
||||
|
||||
# Process the audio input
|
||||
result = await pipeline.run(audio_input)
|
||||
|
||||
# Create an audio player
|
||||
player = sd.OutputStream(samplerate=24000, channels=1, dtype=np.int16)
|
||||
player.start()
|
||||
|
||||
# Store the agent's response for adding to memory
|
||||
agent_response = ""
|
||||
|
||||
print("\nAgent response:")
|
||||
# Play the audio stream as it comes in
|
||||
async for event in result.stream():
|
||||
if event.type == "voice_stream_event_audio":
|
||||
player.write(event.data)
|
||||
elif event.type == "voice_stream_event_content":
|
||||
# Accumulate and print the text response
|
||||
content = event.data
|
||||
agent_response += content
|
||||
print(content, end="", flush=True)
|
||||
|
||||
print("\n")
|
||||
|
||||
# Example of saving the conversation to Mem0 after completion
|
||||
if agent_response:
|
||||
try:
|
||||
await mem0_client.add(
|
||||
f"Agent response: {agent_response}",
|
||||
user_id=USER_ID,
|
||||
metadata={"type": "agent_response"}
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Failed to store memory: {e}")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\nExiting...")
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
## Key Features of This Implementation
|
||||
|
||||
This implementation offers several key features:
|
||||
|
||||
1. **Simplified User Management**: Uses a global `USER_ID` variable for simplicity, but can be extended to manage multiple users.
|
||||
|
||||
2. **Real Microphone Input**: Includes a `record_from_microphone()` function that captures actual voice input from your microphone.
|
||||
|
||||
3. **Interactive Voice Loop**: Implements a continuous interaction loop, allowing for multiple back-and-forth exchanges.
|
||||
|
||||
4. **Memory Management Tools**:
|
||||
- `save_memories`: Stores user memories in Mem0
|
||||
- `search_memories`: Searches for relevant past information
|
||||
|
||||
5. **Voice Configuration**: Demonstrates how to configure TTS settings for the voice response.
|
||||
|
||||
## Running the Example
|
||||
|
||||
To run this example:
|
||||
|
||||
1. Replace the placeholder API keys with your actual keys
|
||||
2. Make sure your microphone is properly connected
|
||||
3. Run the script with Python 3.8 or newer
|
||||
4. Press Enter to start recording, then speak your request
|
||||
5. Press 'q' to quit the application
|
||||
|
||||
The agent will listen to your request, process it through the OpenAI model, utilize Mem0 for memory operations as needed, and respond both through text output and voice speech.
|
||||
|
||||
## Best Practices for Voice Agents with Memory
|
||||
|
||||
1. **Optimizing Memory for Voice**: Keep memories concise and relevant for voice responses.
|
||||
|
||||
2. **Forgetting Mechanism**: Implement a way to delete or expire memories that are no longer relevant.
|
||||
|
||||
3. **Context Preservation**: Store enough context with each memory to make retrieval effective.
|
||||
|
||||
4. **Error Handling**: Implement robust error handling for memory operations, as voice interactions should continue smoothly even if memory operations fail.
|
||||
|
||||
## Conclusion
|
||||
|
||||
By combining OpenAI's Agents SDK with Mem0's memory capabilities, you can create voice agents that maintain persistent memory of user preferences and past interactions. This significantly enhances the user experience by making conversations more natural and personalized.
|
||||
|
||||
As you build your voice application, experiment with different memory strategies and filtering approaches to find the optimal balance between comprehensive memory and efficient retrieval for your specific use case.
|
||||
|
||||
## Debugging Function Tools
|
||||
|
||||
When working with the OpenAI Agents SDK, you might notice that regular `print()` statements inside `@function_tool` decorated functions don't appear in your console output. This is because the Agents SDK captures and redirects standard output when executing these functions.
|
||||
|
||||
To effectively debug your function tools, use Python's `logging` module instead:
|
||||
|
||||
```python
|
||||
import logging
|
||||
|
||||
# Set up logging at the top of your file
|
||||
logging.basicConfig(
|
||||
level=logging.DEBUG,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
||||
force=True
|
||||
)
|
||||
logger = logging.getLogger("memory_voice_agent")
|
||||
|
||||
# Then use logger in your function tools
|
||||
@function_tool
|
||||
async def save_memories(
|
||||
memory: str
|
||||
) -> str:
|
||||
"""Store a user memory in memory."""
|
||||
# This will be visible in your console
|
||||
logger.debug(f"Saving memory: {memory} for user {USER_ID}")
|
||||
|
||||
# Rest of your function...
|
||||
```
|
||||
@@ -66,4 +66,8 @@ Explore how **Mem0** can power real-world applications and bring personalized, i
|
||||
<Card title="OpenAI Inbuilt Tools" icon="robot" href="/examples/openai-inbuilt-tools">
|
||||
Use Mem0's memory capabilities with OpenAI's Inbuilt Tools to create AI agents with persistent memory.
|
||||
</Card>
|
||||
|
||||
<Card title="Mem0 OpenAI Voice Demo" icon="robot" href="/examples/mem0-openai-voice-demo">
|
||||
Use Mem0's memory capabilities with OpenAI's Inbuilt Tools to create AI agents with persistent memory.
|
||||
</Card>
|
||||
</CardGroup>
|
||||
|
||||
+35
-1
@@ -783,6 +783,13 @@
|
||||
"title": "Immutable",
|
||||
"default": false
|
||||
},
|
||||
"expiration_date": {
|
||||
"type": "string",
|
||||
"format": "date-time",
|
||||
"description": "The date and time when the memory will expire. The default expiration date is 30 days from the date of creation. Format: YYYY-MM-DD",
|
||||
"title": "Expiration date",
|
||||
"nullable": true
|
||||
},
|
||||
"organization": {
|
||||
"type": "string"
|
||||
},
|
||||
@@ -1183,6 +1190,13 @@
|
||||
"title": "Immutable",
|
||||
"default": false
|
||||
},
|
||||
"expiration_date": {
|
||||
"type": "string",
|
||||
"format": "date-time",
|
||||
"description": "The date and time when the memory will expire. The default expiration date is 30 days from the date of creation. Format: YYYY-MM-DD",
|
||||
"title": "Expiration date",
|
||||
"nullable": true
|
||||
},
|
||||
"organization": {
|
||||
"type": "string"
|
||||
},
|
||||
@@ -1324,6 +1338,13 @@
|
||||
"title": "Immutable",
|
||||
"default": false
|
||||
},
|
||||
"expiration_date": {
|
||||
"type": "string",
|
||||
"format": "date-time",
|
||||
"description": "The date and time when the memory will expire. The default expiration date is 30 days from the date of creation. Format: YYYY-MM-DD",
|
||||
"title": "Expiration date",
|
||||
"nullable": true
|
||||
},
|
||||
"created_at": {
|
||||
"type": "string",
|
||||
"format": "date-time",
|
||||
@@ -1451,6 +1472,13 @@
|
||||
"title": "Immutable",
|
||||
"default": false
|
||||
},
|
||||
"expiration_date": {
|
||||
"type": "string",
|
||||
"format": "date-time",
|
||||
"description": "The date and time when the memory will expire. The default expiration date is 30 days from the date of creation. Format: YYYY-MM-DD",
|
||||
"title": "Expiration date",
|
||||
"nullable": true
|
||||
},
|
||||
"created_at": {
|
||||
"type": "string",
|
||||
"format": "date-time",
|
||||
@@ -1761,7 +1789,7 @@
|
||||
},
|
||||
{
|
||||
"lang": "JavaScript",
|
||||
"source": "const options = {\n method: 'PUT',\n headers: {\n 'Authorization': 'Token <api-key>',\n 'Content-Type': 'application/json'\n },\n body: JSON.stringify({\n text: 'Your updated memory text here'\n })\n};\n\nfetch('https://api.mem0.ai/v1/memories/{memory_id}/', options)\n .then(response => response.json())\n .then(response => console.log(response))\n .catch(err => console.error(err));"
|
||||
"source": "// To use the JavaScript SDK, install the package:\n// npm i mem0ai\n\nimport MemoryClient from 'mem0ai';\nconst client = new MemoryClient({ apiKey: \"your-api-key\" });\n\n// Update a specific memory\nconst memory_id=<memory_id>\nconst message=\"Your updated memory message here\"\nclient.update(memory_id, message)\n .then(result => console.log(result))\n .catch(error => console.error(error));"
|
||||
},
|
||||
{
|
||||
"lang": "cURL",
|
||||
@@ -4798,6 +4826,12 @@
|
||||
"type": "boolean",
|
||||
"default": false
|
||||
},
|
||||
"expiration_date": {
|
||||
"description": "The date and time when the memory will expire. The default expiration date is 30 days from the date of creation. Format: YYYY-MM-DD",
|
||||
"title": "Expiration date",
|
||||
"type": "string",
|
||||
"nullable": true
|
||||
},
|
||||
"org_id": {
|
||||
"description": "The unique identifier of the organization associated with this memory.",
|
||||
"title": "Organization id",
|
||||
|
||||
+122
-35
@@ -416,6 +416,89 @@ curl -X POST "https://api.mem0.ai/v1/memories/" \
|
||||
|
||||
</CodeGroup>
|
||||
|
||||
#### Setting Memory Expiration Date
|
||||
|
||||
You can set an expiration date for memories, after which they will no longer be retrieved in searches. This is useful for creating temporary memories or memories that are only relevant for a specific time period.
|
||||
|
||||
<CodeGroup>
|
||||
|
||||
```python Python
|
||||
import datetime
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "I'll be in San Francisco until August 31st."
|
||||
}
|
||||
]
|
||||
|
||||
# Set an expiration date for this memory
|
||||
client.add(messages=messages, user_id="alex", expiration_date=str(datetime.datetime.now().date() + datetime.timedelta(days=30)))
|
||||
|
||||
# You can also use an explicit date string
|
||||
client.add(messages=messages, user_id="alex", expiration_date="2023-08-31")
|
||||
```
|
||||
|
||||
```javascript JavaScript
|
||||
const messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "I'll be in San Francisco until August 31st."
|
||||
}
|
||||
];
|
||||
|
||||
// Set an expiration date 30 days from now
|
||||
const expirationDate = new Date();
|
||||
expirationDate.setDate(expirationDate.getDate() + 30);
|
||||
client.add(messages, {
|
||||
user_id: "alex",
|
||||
expiration_date: expirationDate.toISOString().split('T')[0]
|
||||
})
|
||||
.then(response => console.log(response))
|
||||
.catch(error => console.error(error));
|
||||
|
||||
// You can also use an explicit date string
|
||||
client.add(messages, {
|
||||
user_id: "alex",
|
||||
expiration_date: "2023-08-31"
|
||||
})
|
||||
.then(response => console.log(response))
|
||||
.catch(error => console.error(error));
|
||||
```
|
||||
|
||||
```bash cURL
|
||||
curl -X POST "https://api.mem0.ai/v1/memories/" \
|
||||
-H "Authorization: Token your-api-key" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "I'll be in San Francisco until August 31st."
|
||||
}
|
||||
],
|
||||
"user_id": "alex",
|
||||
"expiration_date": "2023-08-31"
|
||||
}'
|
||||
```
|
||||
|
||||
```json Output
|
||||
[
|
||||
{
|
||||
"id": "a1b2c3d4-e5f6-4g7h-8i9j-k0l1m2n3o4p5",
|
||||
"data": {
|
||||
"memory": "In San Francisco until August 31st"
|
||||
},
|
||||
"event": "ADD"
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
</CodeGroup>
|
||||
|
||||
<Note>
|
||||
Once a memory reaches its expiration date, it won't be included in search or get results.
|
||||
</Note>
|
||||
|
||||
#### Monitor Memories
|
||||
|
||||
@@ -470,6 +553,7 @@ curl -X POST "https://api.mem0.ai/v1/memories/search/" \
|
||||
"metadata": {"food": "vegan"},
|
||||
"categories": ["food_preferences"],
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at": "2024-07-20T01:30:36.275141-07:00",
|
||||
"updated_at": "2024-07-20T01:30:36.275172-07:00"
|
||||
}
|
||||
@@ -486,6 +570,7 @@ curl -X POST "https://api.mem0.ai/v1/memories/search/" \
|
||||
"metadata": {"food": "vegan"},
|
||||
"categories": ["food_preferences"],
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at": "2024-07-20T01:30:36.275141-07:00",
|
||||
"updated_at": "2024-07-20T01:30:36.275172-07:00"
|
||||
}
|
||||
@@ -531,6 +616,7 @@ curl -X POST "https://api.mem0.ai/v1/memories/search/?version=v2" \
|
||||
"metadata": {"food": "vegan"},
|
||||
"categories": ["food_preferences"],
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at": "2024-07-20T01:30:36.275141-07:00",
|
||||
"updated_at": "2024-07-20T01:30:36.275172-07:00"
|
||||
}
|
||||
@@ -621,6 +707,7 @@ curl -X POST "https://api.mem0.ai/v1/memories/search/?version=v2" \
|
||||
"metadata": null,
|
||||
"categories": ["food_preferences"],
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at": "2024-07-20T01:30:36.275141-07:00",
|
||||
"updated_at": "2024-07-20T01:30:36.275172-07:00"
|
||||
}
|
||||
@@ -686,6 +773,7 @@ curl -X POST "https://api.mem0.ai/v1/memories/search/?version=v2" \
|
||||
"metadata": null,
|
||||
"categories": ["food_preferences"],
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at": "2024-07-20T01:30:36.275141-07:00",
|
||||
"updated_at": "2024-07-20T01:30:36.275172-07:00"
|
||||
}
|
||||
@@ -760,6 +848,7 @@ curl -X POST "https://api.mem0.ai/v1/memories/search/?version=v2" \
|
||||
"metadata": {"food": "vegan"},
|
||||
"categories": ["food_preferences"],
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at": "2024-07-20T01:30:36.275141-07:00",
|
||||
"updated_at": "2024-07-20T01:30:36.275172-07:00"
|
||||
}
|
||||
@@ -860,6 +949,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?user_id=alex&page=1&page_size=50"
|
||||
"hash":"62bc074f56d1f909f1b4c2b639f56f6a",
|
||||
"metadata":None,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at":"2024-07-25T23:57:00.108347-07:00",
|
||||
"updated_at":"2024-07-25T23:57:00.108367-07:00",
|
||||
"categories":None
|
||||
@@ -871,6 +961,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?user_id=alex&page=1&page_size=50"
|
||||
"hash":"35a305373d639b0bffc6c2a3e2eb4244",
|
||||
"metadata":"None",
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at":"2024-07-26T00:31:03.543759-07:00",
|
||||
"updated_at":"2024-07-26T00:31:03.543778-07:00",
|
||||
"categories":None
|
||||
@@ -894,6 +985,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?user_id=alex&page=1&page_size=50"
|
||||
"hash":"62bc074f56d1f909f1b4c2b639f56f6a",
|
||||
"metadata":None,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at":"2024-07-25T23:57:00.108347-07:00",
|
||||
"updated_at":"2024-07-25T23:57:00.108367-07:00",
|
||||
"categories":None
|
||||
@@ -905,6 +997,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?user_id=alex&page=1&page_size=50"
|
||||
"hash":"35a305373d639b0bffc6c2a3e2eb4244",
|
||||
"metadata":"None",
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at":"2024-07-26T00:31:03.543759-07:00",
|
||||
"updated_at":"2024-07-26T00:31:03.543778-07:00",
|
||||
"categories":None
|
||||
@@ -949,6 +1042,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?agent_id=ai-tutor&page=1&page_size
|
||||
"hash":"62bc074f56d1f909f1b4c2b639f56f6a",
|
||||
"metadata":None,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at":"2024-07-25T23:57:00.108347-07:00",
|
||||
"updated_at":"2024-07-25T23:57:00.108367-07:00",
|
||||
"categories":None
|
||||
@@ -960,6 +1054,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?agent_id=ai-tutor&page=1&page_size
|
||||
"hash":"35a305373d639b0bffc6c2a3e2eb4244",
|
||||
"metadata":None,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at":"2024-07-26T00:31:03.543759-07:00",
|
||||
"updated_at":"2024-07-26T00:31:03.543778-07:00",
|
||||
"categories":None
|
||||
@@ -983,6 +1078,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?agent_id=ai-tutor&page=1&page_size
|
||||
"hash": "62bc074f56d1f909f1b4c2b639f56f6a",
|
||||
"metadata":None,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at": "2024-07-25T23:57:00.108347-07:00",
|
||||
"updated_at": "2024-07-25T23:57:00.108367-07:00"
|
||||
},
|
||||
@@ -993,6 +1089,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?agent_id=ai-tutor&page=1&page_size
|
||||
"hash": "35a305373d639b0bffc6c2a3e2eb4244",
|
||||
"metadata":None,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at": "2024-07-26T00:31:03.543759-07:00",
|
||||
"updated_at": "2024-07-26T00:31:03.543778-07:00"
|
||||
}
|
||||
@@ -1035,6 +1132,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?user_id=alex&run_id=trip-planning-
|
||||
"hash":"d2088c936e259f2f5d2d75543d31401c",
|
||||
"metadata":None,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at":"2024-07-26T00:25:16.566471-07:00",
|
||||
"updated_at":"2024-07-26T00:25:16.566492-07:00",
|
||||
"categories":None
|
||||
@@ -1046,6 +1144,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?user_id=alex&run_id=trip-planning-
|
||||
"hash":"d2088c936e259f2f5d2d75543d31401c",
|
||||
"metadata":None,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at":"2024-07-26T00:33:20.350542-07:00",
|
||||
"updated_at":"2024-07-26T00:33:20.350560-07:00",
|
||||
"categories":None
|
||||
@@ -1057,6 +1156,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?user_id=alex&run_id=trip-planning-
|
||||
"hash":"d2088c936e259f2f5d2d75543d31401c",
|
||||
"metadata":None,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at":"2024-07-26T00:51:09.642275-07:00",
|
||||
"updated_at":"2024-07-26T00:51:09.642295-07:00",
|
||||
"categories":None
|
||||
@@ -1080,6 +1180,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?user_id=alex&run_id=trip-planning-
|
||||
"hash": "d2088c936e259f2f5d2d75543d31401c",
|
||||
"metadata":None,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at": "2024-07-26T00:25:16.566471-07:00",
|
||||
"updated_at": "2024-07-26T00:25:16.566492-07:00",
|
||||
"categories": ["food_preferences"]
|
||||
@@ -1091,6 +1192,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?user_id=alex&run_id=trip-planning-
|
||||
"hash": "d2088c936e259f2f5d2d75543d31401c",
|
||||
"metadata":None,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at": "2024-07-26T00:33:20.350542-07:00",
|
||||
"updated_at": "2024-07-26T00:33:20.350560-07:00",
|
||||
"categories": ["food_preferences"]
|
||||
@@ -1102,6 +1204,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?user_id=alex&run_id=trip-planning-
|
||||
"hash": "d2088c936e259f2f5d2d75543d31401c",
|
||||
"metadata":None,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at": "2024-07-26T00:51:09.642275-07:00",
|
||||
"updated_at": "2024-07-26T00:51:09.642295-07:00",
|
||||
"categories": None
|
||||
@@ -1141,6 +1244,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/582bbe6d-506b-48c6-a4c6-5df3b1e6342
|
||||
"hash":"d2088c936e259f2f5d2d75543d31401c",
|
||||
"metadata":"None",
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at":"2024-07-26T00:25:16.566471-07:00",
|
||||
"updated_at":"2024-07-26T00:25:16.566492-07:00",
|
||||
"categories": ["travel"]
|
||||
@@ -1221,6 +1325,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?user_id=alex&keywords=to play&page
|
||||
"hash": "d2088c936e259f2f5d2d75543d31401c",
|
||||
"metadata": null,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at": "2024-07-26T00:25:16.566471-07:00",
|
||||
"updated_at": "2024-07-26T00:25:16.566492-07:00",
|
||||
"categories": ["likes", "food_preferences"]
|
||||
@@ -1232,6 +1337,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?user_id=alex&keywords=to play&page
|
||||
"hash": "d2088c936e259f2f5d2d75543d31401c",
|
||||
"metadata": null,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at": "2024-07-26T00:33:20.350542-07:00",
|
||||
"updated_at": "2024-07-26T00:33:20.350560-07:00",
|
||||
"categories": ["likes"]
|
||||
@@ -1359,6 +1465,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?version=v2&page=1&page_size=50" \
|
||||
"hash":"62bc074f56d1f909f1b4c2b639f56f6a",
|
||||
"metadata":null,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at":"2024-07-25T23:57:00.108347-07:00",
|
||||
"updated_at":"2024-07-25T23:57:00.108367-07:00",
|
||||
"categories": ["food_preferences"]
|
||||
@@ -1379,6 +1486,7 @@ curl -X GET "https://api.mem0.ai/v1/memories/?version=v2&page=1&page_size=50" \
|
||||
"hash":"62bc074f56d1f909f1b4c2b639f56f6a",
|
||||
"metadata":null,
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at":"2024-07-25T23:57:00.108347-07:00",
|
||||
"updated_at":"2024-07-25T23:57:00.108367-07:00",
|
||||
"categories": ["food_preferences"]
|
||||
@@ -1468,45 +1576,25 @@ curl -X GET "https://api.mem0.ai/v1/memories/?version=v2&page=1&page_size=50" \
|
||||
}'
|
||||
```
|
||||
|
||||
```json Output (Default)
|
||||
```json Output
|
||||
[
|
||||
{
|
||||
"id":"f38b689d-6b24-45b7-bced-17fbb4d8bac7",
|
||||
"memory":"Name: Alex. Vegetarian. Allergic to nuts.",
|
||||
"user_id":"alex",
|
||||
"hash":"62bc074f56d1f909f1b4c2b639f56f6a",
|
||||
"metadata":{"food":"vegan"},
|
||||
"immutable": false,
|
||||
"created_at":"2024-07-25T23:57:00.108347-07:00",
|
||||
"updated_at":"2024-07-25T23:57:00.108367-07:00",
|
||||
"categories": ["food_preferences"]
|
||||
}
|
||||
{
|
||||
"id": "654fee-b411-4afe-b7e5-35789b72c4a5",
|
||||
"memory": "Name: Alex. Vegetarian. Allergic to nuts.",
|
||||
"user_id": "alex",
|
||||
"metadata": {"food": "vegan"},
|
||||
"categories": ["food_preferences"],
|
||||
"immutable": false,
|
||||
"expiration_date": null,
|
||||
"created_at": "2024-07-20T01:30:36.275141-07:00",
|
||||
"updated_at": "2024-07-20T01:30:36.275172-07:00"
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
```json Output (Paginated)
|
||||
{
|
||||
"count": 1,
|
||||
"next": null,
|
||||
"previous": null,
|
||||
"results": [
|
||||
{
|
||||
"id":"f38b689d-6b24-45b7-bced-17fbb4d8bac7",
|
||||
"memory":"Name: Alex. Vegetarian. Allergic to nuts.",
|
||||
"user_id":"alex",
|
||||
"hash":"62bc074f56d1f909f1b4c2b639f56f6a",
|
||||
"metadata":{"food":"vegan"},
|
||||
"immutable": false,
|
||||
"created_at":"2024-07-25T23:57:00.108347-07:00",
|
||||
"updated_at":"2024-07-25T23:57:00.108367-07:00",
|
||||
"categories": ["food_preferences"]
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
</CodeGroup>
|
||||
|
||||
|
||||
### 4.5 Memory History
|
||||
|
||||
Get history of how a memory has changed over time.
|
||||
@@ -1819,8 +1907,7 @@ response = client.batch_update(update_memories)
|
||||
print(response)
|
||||
```
|
||||
```javascript JavaScript
|
||||
const updateMemories = [
|
||||
{"memory_id": "285ed74b-6e05-4043-b16b-3abd5b533496",
|
||||
const updateMemories = [{"memory_id": "285ed74b-6e05-4043-b16b-3abd5b533496",
|
||||
text: "Watches football"
|
||||
},
|
||||
{"memory_id": "2c9bd859-d1b7-4d33-a6b8-94e0147c4f07",
|
||||
|
||||
@@ -56,7 +56,13 @@
|
||||
"sourcemap": true,
|
||||
"clean": true,
|
||||
"treeshake": true,
|
||||
"minify": false
|
||||
"minify": false,
|
||||
"external": [
|
||||
"@mem0/community"
|
||||
],
|
||||
"noExternal": [
|
||||
"!src/community/**"
|
||||
]
|
||||
},
|
||||
"keywords": [
|
||||
"mem0",
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
# Dependencies
|
||||
node_modules
|
||||
.pnp
|
||||
.pnp.js
|
||||
|
||||
# Build outputs
|
||||
dist
|
||||
build
|
||||
|
||||
# Lock files
|
||||
package-lock.json
|
||||
yarn.lock
|
||||
pnpm-lock.yaml
|
||||
|
||||
# Coverage
|
||||
coverage
|
||||
|
||||
# Misc
|
||||
.DS_Store
|
||||
.env.local
|
||||
.env.development.local
|
||||
.env.test.local
|
||||
.env.production.local
|
||||
|
||||
# Logs
|
||||
npm-debug.log*
|
||||
yarn-debug.log*
|
||||
yarn-error.log*
|
||||
@@ -0,0 +1,91 @@
|
||||
{
|
||||
"name": "@mem0/community",
|
||||
"version": "0.0.1",
|
||||
"description": "Community features for Mem0",
|
||||
"main": "./dist/index.js",
|
||||
"module": "./dist/index.mjs",
|
||||
"types": "./dist/index.d.ts",
|
||||
"exports": {
|
||||
".": {
|
||||
"types": "./dist/index.d.ts",
|
||||
"require": "./dist/index.js",
|
||||
"import": "./dist/index.mjs"
|
||||
},
|
||||
"./langchain": {
|
||||
"types": "./dist/integrations/langchain/index.d.ts",
|
||||
"require": "./dist/integrations/langchain/index.js",
|
||||
"import": "./dist/integrations/langchain/index.mjs"
|
||||
}
|
||||
},
|
||||
"files": [
|
||||
"dist"
|
||||
],
|
||||
"scripts": {
|
||||
"clean": "rimraf dist",
|
||||
"build": "npm run clean && npx prettier --check . && npx tsup",
|
||||
"dev": "npx nodemon",
|
||||
"test": "jest",
|
||||
"test:ts": "jest --config jest.config.js",
|
||||
"test:watch": "jest --config jest.config.js --watch",
|
||||
"format": "npm run clean && prettier --write .",
|
||||
"format:check": "npm run clean && prettier --check .",
|
||||
"prepublishOnly": "npm run build"
|
||||
},
|
||||
"tsup": {
|
||||
"entry": {
|
||||
"index": "src/index.ts",
|
||||
"integrations/langchain/index": "src/integrations/langchain/index.ts"
|
||||
},
|
||||
"format": [
|
||||
"cjs",
|
||||
"esm"
|
||||
],
|
||||
"dts": {
|
||||
"resolve": true,
|
||||
"compilerOptions": {
|
||||
"rootDir": "src"
|
||||
}
|
||||
},
|
||||
"splitting": false,
|
||||
"sourcemap": true,
|
||||
"clean": true,
|
||||
"treeshake": true,
|
||||
"minify": false,
|
||||
"outDir": "dist",
|
||||
"tsconfig": "./tsconfig.json"
|
||||
},
|
||||
"keywords": [
|
||||
"mem0",
|
||||
"community",
|
||||
"ai",
|
||||
"memory"
|
||||
],
|
||||
"author": "Deshraj Yadav",
|
||||
"license": "Apache-2.0",
|
||||
"devDependencies": {
|
||||
"@types/node": "^22.7.6",
|
||||
"@types/uuid": "^9.0.8",
|
||||
"dotenv": "^16.4.5",
|
||||
"jest": "^29.7.0",
|
||||
"nodemon": "^3.0.1",
|
||||
"prettier": "^3.5.2",
|
||||
"rimraf": "^5.0.5",
|
||||
"ts-jest": "^29.2.6",
|
||||
"tsup": "^8.3.0",
|
||||
"typescript": "5.5.4"
|
||||
},
|
||||
"dependencies": {
|
||||
"@langchain/community": "^0.3.36",
|
||||
"@langchain/core": "^0.3.42",
|
||||
"axios": "1.7.7",
|
||||
"mem0ai": "^2.1.8",
|
||||
"uuid": "9.0.1",
|
||||
"zod": "3.22.4"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=18"
|
||||
},
|
||||
"publishConfig": {
|
||||
"access": "public"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
export * from "./integrations/langchain";
|
||||
@@ -0,0 +1 @@
|
||||
export * from "./mem0";
|
||||
@@ -0,0 +1,314 @@
|
||||
import { MemoryClient } from "mem0ai";
|
||||
import type { Memory, MemoryOptions, SearchOptions } from "mem0ai";
|
||||
|
||||
import {
|
||||
InputValues,
|
||||
OutputValues,
|
||||
MemoryVariables,
|
||||
getInputValue,
|
||||
getOutputValue,
|
||||
} from "@langchain/core/memory";
|
||||
import {
|
||||
AIMessage,
|
||||
BaseMessage,
|
||||
ChatMessage,
|
||||
getBufferString,
|
||||
HumanMessage,
|
||||
SystemMessage,
|
||||
} from "@langchain/core/messages";
|
||||
import {
|
||||
BaseChatMemory,
|
||||
BaseChatMemoryInput,
|
||||
} from "@langchain/community/memory/chat_memory";
|
||||
|
||||
/**
|
||||
* Extracts and formats memory content into a system prompt
|
||||
* @param memory Array of Memory objects from mem0ai
|
||||
* @returns Formatted system prompt string
|
||||
*/
|
||||
export const mem0MemoryContextToSystemPrompt = (memory: Memory[]): string => {
|
||||
if (!memory || !Array.isArray(memory)) {
|
||||
return "";
|
||||
}
|
||||
|
||||
return memory
|
||||
.filter((m) => m?.memory)
|
||||
.map((m) => m.memory)
|
||||
.join("\n");
|
||||
};
|
||||
|
||||
/**
|
||||
* Condenses memory content into a single HumanMessage with context
|
||||
* @param memory Array of Memory objects from mem0ai
|
||||
* @returns HumanMessage containing formatted memory context
|
||||
*/
|
||||
export const condenseMem0MemoryIntoHumanMessage = (
|
||||
memory: Memory[],
|
||||
): HumanMessage => {
|
||||
const basePrompt =
|
||||
"These are the memories I have stored. Give more weightage to the question by users and try to answer that first. You have to modify your answer based on the memories I have provided. If the memories are irrelevant you can ignore them. Also don't reply to this section of the prompt, or the memories, they are only for your reference. The MEMORIES of the USER are: \n\n";
|
||||
const systemPrompt = mem0MemoryContextToSystemPrompt(memory);
|
||||
|
||||
return new HumanMessage(`${basePrompt}\n${systemPrompt}`);
|
||||
};
|
||||
|
||||
/**
|
||||
* Converts Mem0 memories to a list of BaseMessages
|
||||
* @param memories Array of Memory objects from mem0ai
|
||||
* @returns Array of BaseMessage objects
|
||||
*/
|
||||
export const mem0MemoryToMessages = (memories: Memory[]): BaseMessage[] => {
|
||||
if (!memories || !Array.isArray(memories)) {
|
||||
return [];
|
||||
}
|
||||
|
||||
const messages: BaseMessage[] = [];
|
||||
|
||||
// Add memories as system message if present
|
||||
const memoryContent = memories
|
||||
.filter((m) => m?.memory)
|
||||
.map((m) => m.memory)
|
||||
.join("\n");
|
||||
|
||||
if (memoryContent) {
|
||||
messages.push(new SystemMessage(memoryContent));
|
||||
}
|
||||
|
||||
// Add conversation messages
|
||||
memories.forEach((memory) => {
|
||||
if (memory.messages) {
|
||||
memory.messages.forEach((msg) => {
|
||||
const content =
|
||||
typeof msg.content === "string"
|
||||
? msg.content
|
||||
: JSON.stringify(msg.content);
|
||||
if (msg.role === "user") {
|
||||
messages.push(new HumanMessage(content));
|
||||
} else if (msg.role === "assistant") {
|
||||
messages.push(new AIMessage(content));
|
||||
} else if (content) {
|
||||
messages.push(new ChatMessage(content, msg.role));
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
return messages;
|
||||
};
|
||||
|
||||
/**
|
||||
* Interface defining the structure of the input data for the Mem0Client
|
||||
*/
|
||||
export interface ClientOptions {
|
||||
apiKey: string;
|
||||
host?: string;
|
||||
organizationName?: string;
|
||||
projectName?: string;
|
||||
organizationId?: string;
|
||||
projectId?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Interface defining the structure of the input data for the Mem0Memory
|
||||
* class. It includes properties like memoryKey, sessionId, and apiKey.
|
||||
*/
|
||||
export interface Mem0MemoryInput extends BaseChatMemoryInput {
|
||||
sessionId: string;
|
||||
apiKey: string;
|
||||
humanPrefix?: string;
|
||||
aiPrefix?: string;
|
||||
memoryOptions?: MemoryOptions | SearchOptions;
|
||||
mem0Options?: ClientOptions;
|
||||
separateMessages?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Class used to manage the memory of a chat session using the Mem0 service.
|
||||
* It handles loading and saving chat history, and provides methods to format
|
||||
* the memory content for use in chat models.
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* const memory = new Mem0Memory({
|
||||
* sessionId: "user123" // or use user_id inside of memoryOptions (recommended),
|
||||
* apiKey: "your-api-key",
|
||||
* memoryOptions: {
|
||||
* user_id: "user123",
|
||||
* run_id: "run123"
|
||||
* },
|
||||
* });
|
||||
*
|
||||
* // Use with a chat model
|
||||
* const model = new ChatOpenAI({
|
||||
* modelName: "gpt-3.5-turbo",
|
||||
* temperature: 0,
|
||||
* });
|
||||
*
|
||||
* const chain = new ConversationChain({ llm: model, memory });
|
||||
* ```
|
||||
*/
|
||||
export class Mem0Memory extends BaseChatMemory implements Mem0MemoryInput {
|
||||
memoryKey = "history";
|
||||
|
||||
apiKey: string;
|
||||
|
||||
sessionId: string;
|
||||
|
||||
humanPrefix = "Human";
|
||||
|
||||
aiPrefix = "AI";
|
||||
|
||||
mem0Client: InstanceType<typeof MemoryClient>;
|
||||
|
||||
memoryOptions: MemoryOptions | SearchOptions;
|
||||
|
||||
mem0Options: ClientOptions;
|
||||
|
||||
// Whether to return separate messages for chat history with a SystemMessage containing (facts and summary) or return a single HumanMessage with the entire memory context.
|
||||
// Defaults to false (return a single HumanMessage) in order to allow more flexibility with different models.
|
||||
separateMessages?: boolean;
|
||||
|
||||
constructor(fields: Mem0MemoryInput) {
|
||||
if (!fields.apiKey) {
|
||||
throw new Error("apiKey is required for Mem0Memory");
|
||||
}
|
||||
if (!fields.sessionId) {
|
||||
throw new Error("sessionId is required for Mem0Memory");
|
||||
}
|
||||
|
||||
super({
|
||||
returnMessages: fields?.returnMessages ?? false,
|
||||
inputKey: fields?.inputKey,
|
||||
outputKey: fields?.outputKey,
|
||||
});
|
||||
|
||||
this.apiKey = fields.apiKey;
|
||||
this.sessionId = fields.sessionId;
|
||||
this.humanPrefix = fields.humanPrefix ?? this.humanPrefix;
|
||||
this.aiPrefix = fields.aiPrefix ?? this.aiPrefix;
|
||||
this.memoryOptions = fields.memoryOptions ?? {};
|
||||
this.mem0Options = fields.mem0Options ?? {
|
||||
apiKey: this.apiKey,
|
||||
};
|
||||
this.separateMessages = fields.separateMessages ?? false;
|
||||
try {
|
||||
this.mem0Client = new MemoryClient({
|
||||
...this.mem0Options,
|
||||
apiKey: this.apiKey,
|
||||
});
|
||||
} catch (error) {
|
||||
console.error("Failed to initialize Mem0Client:", error);
|
||||
throw new Error(
|
||||
"Failed to initialize Mem0Client. Please check your configuration.",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
get memoryKeys(): string[] {
|
||||
return [this.memoryKey];
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieves memories from the Mem0 service and formats them for use
|
||||
* @param values Input values containing optional search query
|
||||
* @returns Promise resolving to formatted memory variables
|
||||
*/
|
||||
async loadMemoryVariables(values: InputValues): Promise<MemoryVariables> {
|
||||
const searchType = values.input ? "search" : "get_all";
|
||||
let memories: Memory[] = [];
|
||||
|
||||
try {
|
||||
if (searchType === "get_all") {
|
||||
memories = await this.mem0Client.getAll({
|
||||
user_id: this.sessionId,
|
||||
...this.memoryOptions,
|
||||
});
|
||||
} else {
|
||||
memories = await this.mem0Client.search(values.input, {
|
||||
user_id: this.sessionId,
|
||||
...this.memoryOptions,
|
||||
});
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("Error loading memories:", error);
|
||||
return this.returnMessages
|
||||
? { [this.memoryKey]: [] }
|
||||
: { [this.memoryKey]: "" };
|
||||
}
|
||||
|
||||
if (this.returnMessages) {
|
||||
return {
|
||||
[this.memoryKey]: this.separateMessages
|
||||
? mem0MemoryToMessages(memories)
|
||||
: [condenseMem0MemoryIntoHumanMessage(memories)],
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
[this.memoryKey]: this.separateMessages
|
||||
? getBufferString(
|
||||
mem0MemoryToMessages(memories),
|
||||
this.humanPrefix,
|
||||
this.aiPrefix,
|
||||
)
|
||||
: (condenseMem0MemoryIntoHumanMessage(memories).content ?? ""),
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Saves the current conversation context to the Mem0 service
|
||||
* @param inputValues Input messages to be saved
|
||||
* @param outputValues Output messages to be saved
|
||||
* @returns Promise resolving when the context has been saved
|
||||
*/
|
||||
async saveContext(
|
||||
inputValues: InputValues,
|
||||
outputValues: OutputValues,
|
||||
): Promise<void> {
|
||||
const input = getInputValue(inputValues, this.inputKey);
|
||||
const output = getOutputValue(outputValues, this.outputKey);
|
||||
|
||||
if (!input || !output) {
|
||||
console.warn("Missing input or output values, skipping memory save");
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const messages = [
|
||||
{
|
||||
role: "user",
|
||||
content: `${input}`,
|
||||
},
|
||||
{
|
||||
role: "assistant",
|
||||
content: `${output}`,
|
||||
},
|
||||
];
|
||||
|
||||
await this.mem0Client.add(messages, {
|
||||
user_id: this.sessionId,
|
||||
...this.memoryOptions,
|
||||
});
|
||||
} catch (error) {
|
||||
console.error("Error saving memory context:", error);
|
||||
// Continue execution even if memory save fails
|
||||
}
|
||||
|
||||
await super.saveContext(inputValues, outputValues);
|
||||
}
|
||||
|
||||
/**
|
||||
* Clears all memories for the current session
|
||||
* @returns Promise resolving when memories have been cleared
|
||||
*/
|
||||
async clear(): Promise<void> {
|
||||
try {
|
||||
// Note: Implement clear functionality if Mem0Client provides it
|
||||
// await this.mem0Client.clear(this.sessionId);
|
||||
} catch (error) {
|
||||
console.error("Error clearing memories:", error);
|
||||
}
|
||||
|
||||
await super.clear();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"target": "ES2020",
|
||||
"module": "ESNext",
|
||||
"lib": ["ES2020"],
|
||||
"declaration": true,
|
||||
"declarationMap": true,
|
||||
"sourceMap": true,
|
||||
"outDir": "./dist",
|
||||
"rootDir": "./src",
|
||||
"strict": true,
|
||||
"moduleResolution": "node",
|
||||
"esModuleInterop": true,
|
||||
"skipLibCheck": true,
|
||||
"forceConsistentCasingInFileNames": true,
|
||||
"types": ["node"],
|
||||
"typeRoots": ["./node_modules/@types"]
|
||||
},
|
||||
"include": ["src/**/*.ts"],
|
||||
"exclude": ["node_modules", "dist", "**/*.test.ts"]
|
||||
}
|
||||
@@ -3,11 +3,13 @@ export * from "./memory/memory.types";
|
||||
export * from "./types";
|
||||
export * from "./embeddings/base";
|
||||
export * from "./embeddings/openai";
|
||||
export * from "./embeddings/ollama";
|
||||
export * from "./llms/base";
|
||||
export * from "./llms/openai";
|
||||
export * from "./llms/openai_structured";
|
||||
export * from "./llms/anthropic";
|
||||
export * from "./llms/groq";
|
||||
export * from "./llms/ollama";
|
||||
export * from "./vector_stores/base";
|
||||
export * from "./vector_stores/memory";
|
||||
export * from "./vector_stores/qdrant";
|
||||
|
||||
@@ -30,6 +30,8 @@ class BaseEmbedderConfig(ABC):
|
||||
memory_add_embedding_type: Optional[str] = None,
|
||||
memory_update_embedding_type: Optional[str] = None,
|
||||
memory_search_embedding_type: Optional[str] = None,
|
||||
# LM Studio specific
|
||||
lmstudio_base_url: Optional[str] = "http://localhost:1234/v1",
|
||||
):
|
||||
"""
|
||||
Initializes a configuration class instance for the Embeddings.
|
||||
@@ -58,6 +60,8 @@ class BaseEmbedderConfig(ABC):
|
||||
:type memory_update_embedding_type: Optional[str], optional
|
||||
:param memory_search_embedding_type: The type of embedding to use for the search memory action, defaults to None
|
||||
:type memory_search_embedding_type: Optional[str], optional
|
||||
:param lmstudio_base_url: LM Studio base URL to be use, defaults to "http://localhost:1234/v1"
|
||||
:type lmstudio_base_url: Optional[str], optional
|
||||
"""
|
||||
|
||||
self.model = model
|
||||
@@ -82,3 +86,6 @@ class BaseEmbedderConfig(ABC):
|
||||
self.memory_add_embedding_type = memory_add_embedding_type
|
||||
self.memory_update_embedding_type = memory_update_embedding_type
|
||||
self.memory_search_embedding_type = memory_search_embedding_type
|
||||
|
||||
# LM Studio specific
|
||||
self.lmstudio_base_url = lmstudio_base_url
|
||||
|
||||
@@ -39,6 +39,8 @@ class BaseLlmConfig(ABC):
|
||||
deepseek_base_url: Optional[str] = None,
|
||||
# XAI specific
|
||||
xai_base_url: Optional[str] = None,
|
||||
# LM Studio specific
|
||||
lmstudio_base_url: Optional[str] = "http://localhost:1234/v1",
|
||||
):
|
||||
"""
|
||||
Initializes a configuration class instance for the LLM.
|
||||
@@ -83,6 +85,8 @@ class BaseLlmConfig(ABC):
|
||||
:type deepseek_base_url: Optional[str], optional
|
||||
:param xai_base_url: XAI base URL to be use, defaults to None
|
||||
:type xai_base_url: Optional[str], optional
|
||||
:param lmstudio_base_url: LM Studio base URL to be use, defaults to "http://localhost:1234/v1"
|
||||
:type lmstudio_base_url: Optional[str], optional
|
||||
"""
|
||||
|
||||
self.model = model
|
||||
@@ -116,3 +120,6 @@ class BaseLlmConfig(ABC):
|
||||
|
||||
# XAI specific
|
||||
self.xai_base_url = xai_base_url
|
||||
|
||||
# LM Studio specific
|
||||
self.lmstudio_base_url = lmstudio_base_url
|
||||
|
||||
@@ -8,21 +8,26 @@ class AzureAISearchConfig(BaseModel):
|
||||
api_key: str = Field(None, description="API key for the Azure AI Search service")
|
||||
embedding_model_dims: int = Field(None, description="Dimension of the embedding vector")
|
||||
compression_type: Optional[str] = Field(
|
||||
None,
|
||||
description="Type of vector compression to use. Options: 'scalar', 'binary', or None"
|
||||
None, description="Type of vector compression to use. Options: 'scalar', 'binary', or None"
|
||||
)
|
||||
use_float16: bool = Field(
|
||||
False,
|
||||
description="Whether to store vectors in half precision (Edm.Half) instead of full precision (Edm.Single)"
|
||||
False,
|
||||
description="Whether to store vectors in half precision (Edm.Half) instead of full precision (Edm.Single)",
|
||||
)
|
||||
|
||||
hybrid_search: bool = Field(
|
||||
False, description="Whether to use hybrid search. If True, vector_filter_mode must be 'preFilter'"
|
||||
)
|
||||
vector_filter_mode: Optional[str] = Field(
|
||||
"preFilter", description="Mode for vector filtering. Options: 'preFilter', 'postFilter'"
|
||||
)
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def validate_extra_fields(cls, values: Dict[str, Any]) -> Dict[str, Any]:
|
||||
allowed_fields = set(cls.model_fields.keys())
|
||||
input_fields = set(values.keys())
|
||||
extra_fields = input_fields - allowed_fields
|
||||
|
||||
|
||||
# Check for use_compression to provide a helpful error
|
||||
if "use_compression" in extra_fields:
|
||||
raise ValueError(
|
||||
@@ -30,13 +35,13 @@ class AzureAISearchConfig(BaseModel):
|
||||
"Please use 'compression_type=\"scalar\"' instead of 'use_compression=True' "
|
||||
"or 'compression_type=None' instead of 'use_compression=False'."
|
||||
)
|
||||
|
||||
|
||||
if extra_fields:
|
||||
raise ValueError(
|
||||
f"Extra fields not allowed: {', '.join(extra_fields)}. "
|
||||
f"Please input only the following fields: {', '.join(allowed_fields)}"
|
||||
)
|
||||
|
||||
|
||||
# Validate compression_type values
|
||||
if "compression_type" in values and values["compression_type"] is not None:
|
||||
valid_types = ["scalar", "binary"]
|
||||
@@ -45,9 +50,9 @@ class AzureAISearchConfig(BaseModel):
|
||||
f"Invalid compression_type: {values['compression_type']}. "
|
||||
f"Must be one of: {', '.join(valid_types)}, or None"
|
||||
)
|
||||
|
||||
|
||||
return values
|
||||
|
||||
|
||||
model_config = {
|
||||
"arbitrary_types_allowed": True,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,8 +17,7 @@ class ElasticsearchConfig(BaseModel):
|
||||
use_ssl: bool = Field(True, description="Use SSL for connection")
|
||||
auto_create_index: bool = Field(True, description="Automatically create index during initialization")
|
||||
custom_search_query: Optional[Callable[[List[float], int, Optional[Dict]], Dict]] = Field(
|
||||
None,
|
||||
description="Custom search query function. Parameters: (query, limit, filters) -> Dict"
|
||||
None, description="Custom search query function. Parameters: (query, limit, filters) -> Dict"
|
||||
)
|
||||
|
||||
@model_validator(mode="before")
|
||||
|
||||
@@ -14,9 +14,7 @@ class GoogleMatchingEngineConfig(BaseModel):
|
||||
credentials_path: Optional[str] = Field(None, description="Path to service account credentials file")
|
||||
vector_search_api_endpoint: Optional[str] = Field(None, description="Vector search API endpoint")
|
||||
|
||||
model_config = {
|
||||
"extra": "forbid"
|
||||
}
|
||||
model_config = {"extra": "forbid"}
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
@@ -26,4 +24,4 @@ class GoogleMatchingEngineConfig(BaseModel):
|
||||
def model_post_init(self, _context) -> None:
|
||||
"""Set collection_name to index_id if not provided"""
|
||||
if self.collection_name is None:
|
||||
self.collection_name = self.index_id
|
||||
self.collection_name = self.index_id
|
||||
|
||||
@@ -13,7 +13,7 @@ class EmbedderConfig(BaseModel):
|
||||
@field_validator("config")
|
||||
def validate_config(cls, v, values):
|
||||
provider = values.data.get("provider")
|
||||
if provider in ["openai", "ollama", "huggingface", "azure_openai", "gemini", "vertexai", "together"]:
|
||||
if provider in ["openai", "ollama", "huggingface", "azure_openai", "gemini", "vertexai", "together", "lmstudio"]:
|
||||
return v
|
||||
else:
|
||||
raise ValueError(f"Unsupported embedding provider: {provider}")
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
from typing import Literal, Optional
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
from mem0.configs.embeddings.base import BaseEmbedderConfig
|
||||
from mem0.embeddings.base import EmbeddingBase
|
||||
|
||||
|
||||
class LMStudioEmbedding(EmbeddingBase):
|
||||
def __init__(self, config: Optional[BaseEmbedderConfig] = None):
|
||||
super().__init__(config)
|
||||
|
||||
self.config.model = self.config.model or "nomic-ai/nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf"
|
||||
self.config.embedding_dims = self.config.embedding_dims or 1536
|
||||
self.config.api_key = self.config.api_key or "lm-studio"
|
||||
|
||||
self.client = OpenAI(base_url=self.config.lmstudio_base_url, api_key=self.config.api_key)
|
||||
|
||||
def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
|
||||
"""
|
||||
Get the embedding for the given text using LM Studio.
|
||||
Args:
|
||||
text (str): The text to embed.
|
||||
memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
|
||||
Returns:
|
||||
list: The embedding vector.
|
||||
"""
|
||||
text = text.replace("\n", " ")
|
||||
return (
|
||||
self.client.embeddings.create(input=[text], model=self.config.model)
|
||||
.data[0]
|
||||
.embedding
|
||||
)
|
||||
@@ -24,6 +24,7 @@ class LlmConfig(BaseModel):
|
||||
"gemini",
|
||||
"deepseek",
|
||||
"xai",
|
||||
"lmstudio",
|
||||
):
|
||||
return v
|
||||
else:
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
from mem0.configs.llms.base import BaseLlmConfig
|
||||
from mem0.llms.base import LLMBase
|
||||
|
||||
|
||||
class LMStudioLLM(LLMBase):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
super().__init__(config)
|
||||
|
||||
self.config.model = self.config.model or "lmstudio-community/Meta-Llama-3.1-70B-Instruct-GGUF/Meta-Llama-3.1-70B-Instruct-IQ2_M.gguf"
|
||||
self.config.api_key = self.config.api_key or "lm-studio"
|
||||
|
||||
self.client = OpenAI(base_url=self.config.lmstudio_base_url, api_key=self.config.api_key)
|
||||
|
||||
def generate_response(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
response_format: dict = {"type": "json_object"},
|
||||
tools: Optional[List[Dict]] = None,
|
||||
tool_choice: str = "auto"
|
||||
):
|
||||
"""
|
||||
Generate a response based on the given messages using LM Studio.
|
||||
|
||||
Args:
|
||||
messages (list): List of message dicts containing 'role' and 'content'.
|
||||
response_format (str or object, optional): Format of the response. Defaults to "text".
|
||||
tools (list, optional): List of tools that the model can call. Defaults to None.
|
||||
tool_choice (str, optional): Tool choice method. Defaults to "auto".
|
||||
|
||||
Returns:
|
||||
str: The generated response.
|
||||
"""
|
||||
params = {
|
||||
"model": self.config.model,
|
||||
"messages": messages,
|
||||
"temperature": self.config.temperature,
|
||||
"max_tokens": self.config.max_tokens,
|
||||
"top_p": self.config.top_p
|
||||
}
|
||||
if response_format:
|
||||
params["response_format"] = response_format
|
||||
|
||||
response = self.client.chat.completions.create(**params)
|
||||
return response.choices[0].message.content
|
||||
+9
-5
@@ -71,13 +71,14 @@ class Memory(MemoryBase):
|
||||
if "vector_store" not in config_dict and "embedder" in config_dict:
|
||||
config_dict["vector_store"] = {}
|
||||
config_dict["vector_store"]["config"] = {}
|
||||
config_dict["vector_store"]["config"]["embedding_model_dims"] = config_dict["embedder"]["config"]["embedding_dims"]
|
||||
config_dict["vector_store"]["config"]["embedding_model_dims"] = config_dict["embedder"]["config"][
|
||||
"embedding_dims"
|
||||
]
|
||||
try:
|
||||
return config_dict
|
||||
except ValidationError as e:
|
||||
logger.error(f"Configuration validation error: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def add(
|
||||
self,
|
||||
@@ -204,7 +205,8 @@ class Memory(MemoryBase):
|
||||
messages_embeddings = self.embedding_model.embed(new_mem, "add")
|
||||
new_message_embeddings[new_mem] = messages_embeddings
|
||||
existing_memories = self.vector_store.search(
|
||||
query=messages_embeddings,
|
||||
query=new_mem,
|
||||
vectors=messages_embeddings,
|
||||
limit=5,
|
||||
filters=filters,
|
||||
)
|
||||
@@ -222,7 +224,9 @@ class Memory(MemoryBase):
|
||||
temp_uuid_mapping[str(idx)] = item["id"]
|
||||
retrieved_old_memory[idx]["id"] = str(idx)
|
||||
|
||||
function_calling_prompt = get_update_memory_messages(retrieved_old_memory, new_retrieved_facts, self.custom_update_memory_prompt)
|
||||
function_calling_prompt = get_update_memory_messages(
|
||||
retrieved_old_memory, new_retrieved_facts, self.custom_update_memory_prompt
|
||||
)
|
||||
|
||||
try:
|
||||
new_memories_with_actions = self.llm.generate_response(
|
||||
@@ -479,7 +483,7 @@ class Memory(MemoryBase):
|
||||
|
||||
def _search_vector_store(self, query, filters, limit):
|
||||
embeddings = self.embedding_model.embed(query, "search")
|
||||
memories = self.vector_store.search(query=embeddings, limit=limit, filters=filters)
|
||||
memories = self.vector_store.search(query=query, vectors=embeddings, limit=limit, filters=filters)
|
||||
|
||||
excluded_keys = {
|
||||
"user_id",
|
||||
|
||||
@@ -25,6 +25,7 @@ class LlmFactory:
|
||||
"gemini": "mem0.llms.gemini.GeminiLLM",
|
||||
"deepseek": "mem0.llms.deepseek.DeepSeekLLM",
|
||||
"xai": "mem0.llms.xai.XAILLM",
|
||||
"lmstudio": "mem0.llms.lmstudio.LMStudioLLM",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
@@ -47,6 +48,7 @@ class EmbedderFactory:
|
||||
"gemini": "mem0.embeddings.gemini.GoogleGenAIEmbedding",
|
||||
"vertexai": "mem0.embeddings.vertexai.VertexAIEmbedding",
|
||||
"together": "mem0.embeddings.together.TogetherEmbedding",
|
||||
"lmstudio": "mem0.embeddings.lmstudio.LMStudioEmbedding",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
|
||||
@@ -45,8 +45,10 @@ class AzureAISearch(VectorStoreBase):
|
||||
collection_name,
|
||||
api_key,
|
||||
embedding_model_dims,
|
||||
compression_type: Optional[str] = None,
|
||||
compression_type: Optional[str] = None,
|
||||
use_float16: bool = False,
|
||||
hybrid_search: bool = False,
|
||||
vector_filter_mode: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the Azure AI Search vector store.
|
||||
@@ -60,13 +62,17 @@ class AzureAISearch(VectorStoreBase):
|
||||
Allowed values are None (no quantization), "scalar", or "binary".
|
||||
use_float16 (bool): Whether to store vectors in half precision (Edm.Half) or full precision (Edm.Single).
|
||||
(Note: This flag is preserved from the initial implementation per feedback.)
|
||||
hybrid_search (bool): Whether to use hybrid search. Default is False.
|
||||
vector_filter_mode (Optional[str]): Mode for vector filtering. Default is "preFilter".
|
||||
"""
|
||||
self.index_name = collection_name
|
||||
self.collection_name = collection_name
|
||||
self.embedding_model_dims = embedding_model_dims
|
||||
# If compression_type is None, treat it as "none".
|
||||
self.compression_type = (compression_type or "none").lower()
|
||||
self.compression_type = (compression_type or "none").lower()
|
||||
self.use_float16 = use_float16
|
||||
self.hybrid_search = hybrid_search
|
||||
self.vector_filter_mode = vector_filter_mode
|
||||
|
||||
self.search_client = SearchClient(
|
||||
endpoint=f"https://{service_name}.search.windows.net",
|
||||
@@ -113,8 +119,6 @@ class AzureAISearch(VectorStoreBase):
|
||||
)
|
||||
]
|
||||
# If no compression is desired, compression_configurations remains empty.
|
||||
|
||||
|
||||
fields = [
|
||||
SimpleField(name="id", type=SearchFieldDataType.String, key=True),
|
||||
SimpleField(name="user_id", type=SearchFieldDataType.String, filterable=True),
|
||||
@@ -123,11 +127,11 @@ class AzureAISearch(VectorStoreBase):
|
||||
SearchField(
|
||||
name="vector",
|
||||
type=vector_type,
|
||||
searchable=True,
|
||||
searchable=True,
|
||||
vector_search_dimensions=self.embedding_model_dims,
|
||||
vector_search_profile_name="my-vector-config",
|
||||
),
|
||||
SimpleField(name="payload", type=SearchFieldDataType.String, searchable=True),
|
||||
SearchField(name="payload", type=SearchFieldDataType.String, searchable=True),
|
||||
]
|
||||
|
||||
vector_search = VectorSearch(
|
||||
@@ -135,7 +139,7 @@ class AzureAISearch(VectorStoreBase):
|
||||
VectorSearchProfile(
|
||||
name="my-vector-config",
|
||||
algorithm_configuration_name="my-algorithms-config",
|
||||
compression_name=compression_name if self.compression_type != "none" else None
|
||||
compression_name=compression_name if self.compression_type != "none" else None,
|
||||
)
|
||||
],
|
||||
algorithms=[HnswAlgorithmConfiguration(name="my-algorithms-config")],
|
||||
@@ -164,8 +168,7 @@ class AzureAISearch(VectorStoreBase):
|
||||
"""
|
||||
logger.info(f"Inserting {len(vectors)} vectors into index {self.index_name}")
|
||||
documents = [
|
||||
self._generate_document(vector, payload, id)
|
||||
for id, vector, payload in zip(ids, vectors, payloads)
|
||||
self._generate_document(vector, payload, id) for id, vector, payload in zip(ids, vectors, payloads)
|
||||
]
|
||||
response = self.search_client.upload_documents(documents)
|
||||
for doc in response:
|
||||
@@ -189,12 +192,13 @@ class AzureAISearch(VectorStoreBase):
|
||||
filter_expression = " and ".join(filter_conditions)
|
||||
return filter_expression
|
||||
|
||||
def search(self, query, limit=5, filters=None):
|
||||
def search(self, query, vectors, limit=5, filters=None):
|
||||
"""
|
||||
Search for similar vectors.
|
||||
|
||||
Args:
|
||||
query (List[float]): Query vector.
|
||||
query (str): Query.
|
||||
vectors (List[float]): Query vector.
|
||||
limit (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (Dict, optional): Filters to apply to the search. Defaults to None.
|
||||
|
||||
@@ -205,23 +209,28 @@ class AzureAISearch(VectorStoreBase):
|
||||
if filters:
|
||||
filter_expression = self._build_filter_expression(filters)
|
||||
|
||||
vector_query = VectorizedQuery(
|
||||
vector=query, k_nearest_neighbors=limit, fields="vector"
|
||||
)
|
||||
search_results = self.search_client.search(
|
||||
vector_queries=[vector_query],
|
||||
filter=filter_expression,
|
||||
top=limit
|
||||
)
|
||||
vector_query = VectorizedQuery(vector=vectors, k_nearest_neighbors=limit, fields="vector")
|
||||
if self.hybrid_search:
|
||||
search_results = self.search_client.search(
|
||||
search_text=query,
|
||||
vector_queries=[vector_query],
|
||||
filter=filter_expression,
|
||||
top=limit,
|
||||
vector_filter_mode=self.vector_filter_mode,
|
||||
search_fields=["payload"],
|
||||
)
|
||||
else:
|
||||
search_results = self.search_client.search(
|
||||
vector_queries=[vector_query],
|
||||
filter=filter_expression,
|
||||
top=limit,
|
||||
vector_filter_mode=self.vector_filter_mode,
|
||||
)
|
||||
|
||||
results = []
|
||||
for result in search_results:
|
||||
payload = json.loads(result["payload"])
|
||||
results.append(
|
||||
OutputData(
|
||||
id=result["id"], score=result["@search.score"], payload=payload
|
||||
)
|
||||
)
|
||||
results.append(OutputData(id=result["id"], score=result["@search.score"], payload=payload))
|
||||
return results
|
||||
|
||||
def delete(self, vector_id):
|
||||
@@ -275,9 +284,7 @@ class AzureAISearch(VectorStoreBase):
|
||||
result = self.search_client.get_document(key=vector_id)
|
||||
except ResourceNotFoundError:
|
||||
return None
|
||||
return OutputData(
|
||||
id=result["id"], score=None, payload=json.loads(result["payload"])
|
||||
)
|
||||
return OutputData(id=result["id"], score=None, payload=json.loads(result["payload"]))
|
||||
|
||||
def list_cols(self) -> List[str]:
|
||||
"""
|
||||
@@ -321,17 +328,11 @@ class AzureAISearch(VectorStoreBase):
|
||||
if filters:
|
||||
filter_expression = self._build_filter_expression(filters)
|
||||
|
||||
search_results = self.search_client.search(
|
||||
search_text="*", filter=filter_expression, top=limit
|
||||
)
|
||||
search_results = self.search_client.search(search_text="*", filter=filter_expression, top=limit)
|
||||
results = []
|
||||
for result in search_results:
|
||||
payload = json.loads(result["payload"])
|
||||
results.append(
|
||||
OutputData(
|
||||
id=result["id"], score=result["@search.score"], payload=payload
|
||||
)
|
||||
)
|
||||
results.append(OutputData(id=result["id"], score=result["@search.score"], payload=payload))
|
||||
return [results]
|
||||
|
||||
def __del__(self):
|
||||
|
||||
@@ -13,7 +13,7 @@ class VectorStoreBase(ABC):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def search(self, query, limit=5, filters=None):
|
||||
def search(self, query, vectors, limit=5, filters=None):
|
||||
"""Search for similar vectors."""
|
||||
pass
|
||||
|
||||
|
||||
@@ -127,19 +127,22 @@ class ChromaDB(VectorStoreBase):
|
||||
logger.info(f"Inserting {len(vectors)} vectors into collection {self.collection_name}")
|
||||
self.collection.add(ids=ids, embeddings=vectors, metadatas=payloads)
|
||||
|
||||
def search(self, query: List[list], limit: int = 5, filters: Optional[Dict] = None) -> List[OutputData]:
|
||||
def search(
|
||||
self, query: str, vectors: List[list], limit: int = 5, filters: Optional[Dict] = None
|
||||
) -> List[OutputData]:
|
||||
"""
|
||||
Search for similar vectors.
|
||||
|
||||
Args:
|
||||
query (List[list]): Query vector.
|
||||
query (str): Query.
|
||||
vectors (List[list]): List of vectors to search.
|
||||
limit (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (Optional[Dict], optional): Filters to apply to the search. Defaults to None.
|
||||
|
||||
Returns:
|
||||
List[OutputData]: Search results.
|
||||
"""
|
||||
results = self.collection.query(query_embeddings=query, where=filters, n_results=limit)
|
||||
results = self.collection.query(query_embeddings=vectors, where=filters, n_results=limit)
|
||||
final_results = self._parse_output(results)
|
||||
return final_results
|
||||
|
||||
|
||||
@@ -45,7 +45,7 @@ class ElasticsearchDB(VectorStoreBase):
|
||||
# Create index only if auto_create_index is True
|
||||
if config.auto_create_index:
|
||||
self.create_index()
|
||||
|
||||
|
||||
if config.custom_search_query:
|
||||
self.custom_search_query = config.custom_search_query
|
||||
else:
|
||||
@@ -121,21 +121,25 @@ class ElasticsearchDB(VectorStoreBase):
|
||||
)
|
||||
return results
|
||||
|
||||
def search(self, query: List[float], limit: int = 5, filters: Optional[Dict] = None) -> List[OutputData]:
|
||||
def search(
|
||||
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
|
||||
) -> List[OutputData]:
|
||||
"""
|
||||
Search with two options:
|
||||
1. Use custom search query if provided
|
||||
2. Use KNN search on vectors with pre-filtering if no custom search query is provided
|
||||
"""
|
||||
if self.custom_search_query:
|
||||
search_query = self.custom_search_query(query, limit, filters)
|
||||
search_query = self.custom_search_query(vectors, limit, filters)
|
||||
else:
|
||||
search_query = {"knn": {"field": "vector", "query_vector": query, "k": limit, "num_candidates": limit * 2}}
|
||||
search_query = {
|
||||
"knn": {"field": "vector", "query_vector": vectors, "k": limit, "num_candidates": limit * 2}
|
||||
}
|
||||
if filters:
|
||||
filter_conditions = []
|
||||
for key, value in filters.items():
|
||||
filter_conditions.append({"term": {f"metadata.{key}": value}})
|
||||
search_query["filter"] = {"bool": {"must": filter_conditions}}
|
||||
search_query["knn"]["filter"] = {"bool": {"must": filter_conditions}}
|
||||
|
||||
response = self.client.search(index=self.collection_name, body=search_query)
|
||||
|
||||
|
||||
@@ -134,12 +134,13 @@ class MilvusDB(VectorStoreBase):
|
||||
|
||||
return memory
|
||||
|
||||
def search(self, query: list, limit: int = 5, filters: dict = None) -> list:
|
||||
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> list:
|
||||
"""
|
||||
Search for similar vectors.
|
||||
|
||||
Args:
|
||||
query (List[float]): Query vector.
|
||||
query (str): Query.
|
||||
vectors (List[float]): Query vector.
|
||||
limit (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (Dict, optional): Filters to apply to the search. Defaults to None.
|
||||
|
||||
@@ -149,7 +150,7 @@ class MilvusDB(VectorStoreBase):
|
||||
query_filter = self._create_filter(filters) if filters else None
|
||||
hits = self.client.search(
|
||||
collection_name=self.collection_name,
|
||||
data=[query],
|
||||
data=[vectors],
|
||||
limit=limit,
|
||||
filter=query_filter,
|
||||
output_fields=["*"],
|
||||
|
||||
@@ -28,10 +28,12 @@ class OpenSearchDB(VectorStoreBase):
|
||||
# Initialize OpenSearch client
|
||||
self.client = OpenSearch(
|
||||
hosts=[{"host": config.host, "port": config.port or 9200}],
|
||||
http_auth=config.http_auth if config.http_auth else ((config.user, config.password) if (config.user and config.password) else None),
|
||||
http_auth=config.http_auth
|
||||
if config.http_auth
|
||||
else ((config.user, config.password) if (config.user and config.password) else None),
|
||||
use_ssl=config.use_ssl,
|
||||
verify_certs=config.verify_certs,
|
||||
connection_class=RequestsHttpConnection
|
||||
connection_class=RequestsHttpConnection,
|
||||
)
|
||||
|
||||
self.collection_name = config.collection_name
|
||||
@@ -115,14 +117,16 @@ class OpenSearchDB(VectorStoreBase):
|
||||
results.append(OutputData(id=id_, score=1.0, payload=payloads[i]))
|
||||
return results
|
||||
|
||||
def search(self, query: List[float], limit: int = 5, filters: Optional[Dict] = None) -> List[OutputData]:
|
||||
def search(
|
||||
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
|
||||
) -> List[OutputData]:
|
||||
"""Search for similar vectors using OpenSearch k-NN search with pre-filtering."""
|
||||
search_query = {
|
||||
"size": limit,
|
||||
"query": {
|
||||
"knn": {
|
||||
"vector": {
|
||||
"vector": query,
|
||||
"vector": vectors,
|
||||
"k": limit,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -120,12 +120,13 @@ class PGVector(VectorStoreBase):
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
def search(self, query, limit=5, filters=None):
|
||||
def search(self, query, vectors, limit=5, filters=None):
|
||||
"""
|
||||
Search for similar vectors.
|
||||
|
||||
Args:
|
||||
query (List[float]): Query vector.
|
||||
query (str): Query.
|
||||
vectors (List[float]): Query vector.
|
||||
limit (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (Dict, optional): Filters to apply to the search. Defaults to None.
|
||||
|
||||
@@ -150,7 +151,7 @@ class PGVector(VectorStoreBase):
|
||||
ORDER BY distance
|
||||
LIMIT %s
|
||||
""",
|
||||
(query, *filter_params, limit),
|
||||
(vectors, *filter_params, limit),
|
||||
)
|
||||
|
||||
results = self.cur.fetchall()
|
||||
|
||||
@@ -199,12 +199,13 @@ class PineconeDB(VectorStoreBase):
|
||||
|
||||
return pinecone_filter
|
||||
|
||||
def search(self, query: List[float], limit: int = 5, filters: Optional[Dict] = None) -> List[OutputData]:
|
||||
def search(self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None) -> List[OutputData]:
|
||||
"""
|
||||
Search for similar vectors.
|
||||
|
||||
Args:
|
||||
query (list): Query vector.
|
||||
query (str): Query.
|
||||
vectors (list): List of vectors to search.
|
||||
limit (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (dict, optional): Filters to apply to the search. Defaults to None.
|
||||
|
||||
@@ -214,7 +215,7 @@ class PineconeDB(VectorStoreBase):
|
||||
filter_dict = self._create_filter(filters) if filters else None
|
||||
|
||||
query_params = {
|
||||
"vector": query,
|
||||
"vector": vectors,
|
||||
"top_k": limit,
|
||||
"include_metadata": True,
|
||||
"include_values": False,
|
||||
|
||||
@@ -127,12 +127,13 @@ class Qdrant(VectorStoreBase):
|
||||
conditions.append(FieldCondition(key=key, match=MatchValue(value=value)))
|
||||
return Filter(must=conditions) if conditions else None
|
||||
|
||||
def search(self, query: list, limit: int = 5, filters: dict = None) -> list:
|
||||
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> list:
|
||||
"""
|
||||
Search for similar vectors.
|
||||
|
||||
Args:
|
||||
query (list): Query vector.
|
||||
query (str): Query.
|
||||
vectors (list): Query vector.
|
||||
limit (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (dict, optional): Filters to apply to the search. Defaults to None.
|
||||
|
||||
@@ -142,7 +143,7 @@ class Qdrant(VectorStoreBase):
|
||||
query_filter = self._create_filter(filters) if filters else None
|
||||
hits = self.client.query_points(
|
||||
collection_name=self.collection_name,
|
||||
query=query,
|
||||
query=vectors,
|
||||
query_filter=query_filter,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
@@ -101,12 +101,12 @@ class RedisDB(VectorStoreBase):
|
||||
data.append(entry)
|
||||
self.index.load(data, id_field="memory_id")
|
||||
|
||||
def search(self, query: list, limit: int = 5, filters: dict = None):
|
||||
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None):
|
||||
conditions = [Tag(key) == value for key, value in filters.items() if value is not None]
|
||||
filter = reduce(lambda x, y: x & y, conditions)
|
||||
|
||||
v = VectorQuery(
|
||||
vector=np.array(query, dtype=np.float32).tobytes(),
|
||||
vector=np.array(vectors, dtype=np.float32).tobytes(),
|
||||
vector_field_name="embedding",
|
||||
return_fields=["memory_id", "hash", "agent_id", "run_id", "user_id", "memory", "metadata", "created_at"],
|
||||
filter_expression=filter,
|
||||
|
||||
@@ -112,16 +112,18 @@ class Supabase(VectorStoreBase):
|
||||
payloads = [{} for _ in vectors]
|
||||
|
||||
records = [(id, vector, payload) for id, vector, payload in zip(ids, vectors, payloads)]
|
||||
print(records)
|
||||
|
||||
self.collection.upsert(records)
|
||||
|
||||
def search(self, query: List[float], limit: int = 5, filters: Optional[dict] = None) -> List[OutputData]:
|
||||
def search(
|
||||
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[dict] = None
|
||||
) -> List[OutputData]:
|
||||
"""
|
||||
Search for similar vectors.
|
||||
|
||||
Args:
|
||||
query (List[float]): Query vector
|
||||
query (str): Query.
|
||||
vectors (List[float]): Query vector.
|
||||
limit (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (Dict, optional): Filters to apply to the search. Defaults to None.
|
||||
|
||||
@@ -129,11 +131,9 @@ class Supabase(VectorStoreBase):
|
||||
List[OutputData]: Search results
|
||||
"""
|
||||
filters = self._preprocess_filters(filters)
|
||||
print(filters)
|
||||
results = self.collection.query(
|
||||
data=query, limit=limit, filters=filters, include_metadata=True, include_value=True
|
||||
data=vectors, limit=limit, filters=filters, include_metadata=True, include_value=True
|
||||
)
|
||||
print(results)
|
||||
|
||||
return [OutputData(id=str(result[0]), score=float(result[1]), payload=result[2]) for result in results]
|
||||
|
||||
|
||||
@@ -32,19 +32,19 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
def __init__(self, **kwargs):
|
||||
"""Initialize Google Matching Engine client."""
|
||||
logger.debug("Initializing Google Matching Engine with kwargs: %s", kwargs)
|
||||
|
||||
|
||||
# If collection_name is passed, use it as deployment_index_id if deployment_index_id is not provided
|
||||
if 'collection_name' in kwargs and 'deployment_index_id' not in kwargs:
|
||||
kwargs['deployment_index_id'] = kwargs['collection_name']
|
||||
logger.debug("Using collection_name as deployment_index_id: %s", kwargs['deployment_index_id'])
|
||||
elif 'deployment_index_id' in kwargs and 'collection_name' not in kwargs:
|
||||
kwargs['collection_name'] = kwargs['deployment_index_id']
|
||||
logger.debug("Using deployment_index_id as collection_name: %s", kwargs['collection_name'])
|
||||
|
||||
if "collection_name" in kwargs and "deployment_index_id" not in kwargs:
|
||||
kwargs["deployment_index_id"] = kwargs["collection_name"]
|
||||
logger.debug("Using collection_name as deployment_index_id: %s", kwargs["deployment_index_id"])
|
||||
elif "deployment_index_id" in kwargs and "collection_name" not in kwargs:
|
||||
kwargs["collection_name"] = kwargs["deployment_index_id"]
|
||||
logger.debug("Using deployment_index_id as collection_name: %s", kwargs["collection_name"])
|
||||
|
||||
try:
|
||||
config = GoogleMatchingEngineConfig(**kwargs)
|
||||
logger.debug("Config created: %s", config.model_dump())
|
||||
logger.debug("Config collection_name: %s", getattr(config, 'collection_name', None))
|
||||
logger.debug("Config collection_name: %s", getattr(config, "collection_name", None))
|
||||
except Exception as e:
|
||||
logger.error("Failed to validate config: %s", str(e))
|
||||
raise
|
||||
@@ -57,41 +57,37 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
self.deployment_index_id = config.deployment_index_id # The deployment-specific ID
|
||||
self.collection_name = config.collection_name
|
||||
self.vector_search_api_endpoint = config.vector_search_api_endpoint
|
||||
|
||||
|
||||
logger.debug("Using project=%s, location=%s", self.project_id, self.region)
|
||||
|
||||
|
||||
# Initialize Vertex AI with credentials if provided
|
||||
init_args = {
|
||||
"project": self.project_id,
|
||||
"location": self.region,
|
||||
}
|
||||
if hasattr(config, 'credentials_path') and config.credentials_path:
|
||||
if hasattr(config, "credentials_path") and config.credentials_path:
|
||||
logger.debug("Using credentials from: %s", config.credentials_path)
|
||||
credentials = service_account.Credentials.from_service_account_file(
|
||||
config.credentials_path
|
||||
)
|
||||
credentials = service_account.Credentials.from_service_account_file(config.credentials_path)
|
||||
init_args["credentials"] = credentials
|
||||
|
||||
|
||||
try:
|
||||
aiplatform.init(**init_args)
|
||||
logger.debug("Vertex AI initialized successfully")
|
||||
except Exception as e:
|
||||
logger.error("Failed to initialize Vertex AI: %s", str(e))
|
||||
raise
|
||||
|
||||
|
||||
try:
|
||||
# Format the index path properly using the configured index_id
|
||||
index_path = f"projects/{self.project_number}/locations/{self.region}/indexes/{self.index_id}"
|
||||
logger.debug("Initializing index with path: %s", index_path)
|
||||
self.index = aiplatform.MatchingEngineIndex(index_name=index_path)
|
||||
logger.debug("Index initialized successfully")
|
||||
|
||||
|
||||
# Format the endpoint name properly
|
||||
endpoint_name = self.endpoint_id
|
||||
logger.debug("Initializing endpoint with name: %s", endpoint_name)
|
||||
self.index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
|
||||
index_endpoint_name=endpoint_name
|
||||
)
|
||||
self.index_endpoint = aiplatform.MatchingEngineIndexEndpoint(index_endpoint_name=endpoint_name)
|
||||
logger.debug("Endpoint initialized successfully")
|
||||
except Exception as e:
|
||||
logger.error("Failed to initialize Matching Engine components: %s", str(e))
|
||||
@@ -119,47 +115,36 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
|
||||
def _create_restriction(self, key: str, value: Any) -> aiplatform_v1.types.index.IndexDatapoint.Restriction:
|
||||
"""Create a restriction object for the Matching Engine index.
|
||||
|
||||
|
||||
Args:
|
||||
key: The namespace/key for the restriction
|
||||
value: The value to restrict on
|
||||
|
||||
|
||||
Returns:
|
||||
Restriction object for the index
|
||||
"""
|
||||
str_value = str(value) if value is not None else ""
|
||||
return aiplatform_v1.types.index.IndexDatapoint.Restriction(
|
||||
namespace=key,
|
||||
allow_list=[str_value]
|
||||
)
|
||||
return aiplatform_v1.types.index.IndexDatapoint.Restriction(namespace=key, allow_list=[str_value])
|
||||
|
||||
def _create_datapoint(
|
||||
self,
|
||||
vector_id: str,
|
||||
vector: List[float],
|
||||
payload: Optional[Dict] = None
|
||||
self, vector_id: str, vector: List[float], payload: Optional[Dict] = None
|
||||
) -> aiplatform_v1.types.index.IndexDatapoint:
|
||||
"""Create a datapoint object for the Matching Engine index.
|
||||
|
||||
|
||||
Args:
|
||||
vector_id: The ID for the datapoint
|
||||
vector: The vector to store
|
||||
payload: Optional metadata to store with the vector
|
||||
|
||||
|
||||
Returns:
|
||||
IndexDatapoint object
|
||||
"""
|
||||
restrictions = []
|
||||
if payload:
|
||||
restrictions = [
|
||||
self._create_restriction(key, value)
|
||||
for key, value in payload.items()
|
||||
]
|
||||
|
||||
restrictions = [self._create_restriction(key, value) for key, value in payload.items()]
|
||||
|
||||
return aiplatform_v1.types.index.IndexDatapoint(
|
||||
datapoint_id=vector_id,
|
||||
feature_vector=vector,
|
||||
restricts=restrictions
|
||||
datapoint_id=vector_id, feature_vector=vector, restricts=restrictions
|
||||
)
|
||||
|
||||
def insert(
|
||||
@@ -169,41 +154,41 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
ids: Optional[List[str]] = None,
|
||||
) -> None:
|
||||
"""Insert vectors into the Matching Engine index.
|
||||
|
||||
|
||||
Args:
|
||||
vectors: List of vectors to insert
|
||||
payloads: Optional list of metadata dictionaries
|
||||
ids: Optional list of IDs for the vectors
|
||||
|
||||
|
||||
Raises:
|
||||
ValueError: If vectors is empty or lengths don't match
|
||||
GoogleAPIError: If the API call fails
|
||||
"""
|
||||
if not vectors:
|
||||
raise ValueError("No vectors provided for insertion")
|
||||
|
||||
|
||||
if payloads and len(payloads) != len(vectors):
|
||||
raise ValueError(f"Number of payloads ({len(payloads)}) does not match number of vectors ({len(vectors)})")
|
||||
|
||||
|
||||
if ids and len(ids) != len(vectors):
|
||||
raise ValueError(f"Number of ids ({len(ids)}) does not match number of vectors ({len(vectors)})")
|
||||
|
||||
|
||||
logger.debug("Starting insert of %d vectors", len(vectors))
|
||||
|
||||
|
||||
try:
|
||||
datapoints = [
|
||||
self._create_datapoint(
|
||||
vector_id=ids[i] if ids else str(uuid.uuid4()),
|
||||
vector=vector,
|
||||
payload=payloads[i] if payloads and i < len(payloads) else None
|
||||
payload=payloads[i] if payloads and i < len(payloads) else None,
|
||||
)
|
||||
for i, vector in enumerate(vectors)
|
||||
]
|
||||
|
||||
|
||||
logger.debug("Created %d datapoints", len(datapoints))
|
||||
self.index.upsert_datapoints(datapoints=datapoints)
|
||||
logger.debug("Successfully inserted datapoints")
|
||||
|
||||
|
||||
except google.api_core.exceptions.GoogleAPIError as e:
|
||||
logger.error("Failed to insert vectors: %s", str(e))
|
||||
raise
|
||||
@@ -212,21 +197,22 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
logger.error("Stack trace: %s", traceback.format_exc())
|
||||
raise
|
||||
|
||||
|
||||
def search(self, query: List[float], limit: int = 5, filters: Optional[Dict] = None) -> List[OutputData]:
|
||||
def search(
|
||||
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
|
||||
) -> List[OutputData]:
|
||||
"""
|
||||
Search for similar vectors.
|
||||
Args:
|
||||
query (List[float]): Query vector.
|
||||
query (str): Query.
|
||||
vectors (List[float]): Query vector.
|
||||
limit (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (Optional[Dict], optional): Filters to apply to the search. Defaults to None.
|
||||
Returns:
|
||||
List[OutputData]: Search results (unwrapped)
|
||||
"""
|
||||
logger.debug("Starting search")
|
||||
logger.debug("Query type: %s, length: %d", type(query), len(query))
|
||||
logger.debug("Limit: %d, Filters: %s", limit, filters)
|
||||
|
||||
|
||||
try:
|
||||
filter_namespaces = []
|
||||
if filters:
|
||||
@@ -235,53 +221,42 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
logger.debug("Processing filter %s=%s (type=%s)", key, value, type(value))
|
||||
if isinstance(value, (str, int, float)):
|
||||
logger.debug("Adding simple filter for %s", key)
|
||||
filter_namespaces.append(
|
||||
Namespace(key, [str(value)], [])
|
||||
)
|
||||
filter_namespaces.append(Namespace(key, [str(value)], []))
|
||||
elif isinstance(value, dict):
|
||||
logger.debug("Adding complex filter for %s", key)
|
||||
includes = value.get('include', [])
|
||||
excludes = value.get('exclude', [])
|
||||
filter_namespaces.append(
|
||||
Namespace(key, includes, excludes)
|
||||
)
|
||||
|
||||
includes = value.get("include", [])
|
||||
excludes = value.get("exclude", [])
|
||||
filter_namespaces.append(Namespace(key, includes, excludes))
|
||||
|
||||
logger.debug("Final filter_namespaces: %s", filter_namespaces)
|
||||
|
||||
|
||||
response = self.index_endpoint.find_neighbors(
|
||||
deployed_index_id=self.deployment_index_id,
|
||||
queries=[query],
|
||||
queries=[vectors],
|
||||
num_neighbors=limit,
|
||||
filter=filter_namespaces if filter_namespaces else None,
|
||||
return_full_datapoint=True
|
||||
return_full_datapoint=True,
|
||||
)
|
||||
|
||||
|
||||
if not response or len(response) == 0 or len(response[0]) == 0:
|
||||
logger.debug("No results found")
|
||||
return []
|
||||
|
||||
|
||||
results = []
|
||||
for neighbor in response[0]:
|
||||
logger.debug("Processing neighbor - id: %s, distance: %s",
|
||||
neighbor.id, neighbor.distance)
|
||||
|
||||
logger.debug("Processing neighbor - id: %s, distance: %s", neighbor.id, neighbor.distance)
|
||||
|
||||
payload = {}
|
||||
if hasattr(neighbor, 'restricts'):
|
||||
if hasattr(neighbor, "restricts"):
|
||||
logger.debug("Processing restricts")
|
||||
for restrict in neighbor.restricts:
|
||||
if (hasattr(restrict, 'name') and
|
||||
hasattr(restrict, 'allow_tokens') and
|
||||
restrict.allow_tokens):
|
||||
if hasattr(restrict, "name") and hasattr(restrict, "allow_tokens") and restrict.allow_tokens:
|
||||
logger.debug("Adding %s: %s", restrict.name, restrict.allow_tokens[0])
|
||||
payload[restrict.name] = restrict.allow_tokens[0]
|
||||
|
||||
output_data = OutputData(
|
||||
id=neighbor.id,
|
||||
score=neighbor.distance,
|
||||
payload=payload
|
||||
)
|
||||
|
||||
output_data = OutputData(id=neighbor.id, score=neighbor.distance, payload=payload)
|
||||
results.append(output_data)
|
||||
|
||||
|
||||
logger.debug("Returning %d results", len(results))
|
||||
return results
|
||||
|
||||
@@ -291,7 +266,6 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
logger.error("Stack trace: %s", traceback.format_exc())
|
||||
raise
|
||||
|
||||
|
||||
def delete(self, vector_id: Optional[str] = None, ids: Optional[List[str]] = None) -> bool:
|
||||
"""
|
||||
Delete vectors from the Matching Engine index.
|
||||
@@ -326,14 +300,13 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
except google.api_core.exceptions.InvalidArgument as e:
|
||||
logger.error("Invalid argument: %s", str(e))
|
||||
return False
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error occurred: %s", str(e))
|
||||
logger.error("Error type: %s", type(e))
|
||||
logger.error("Stack trace: %s", traceback.format_exc())
|
||||
return False
|
||||
|
||||
|
||||
def update(
|
||||
self,
|
||||
vector_id: str,
|
||||
@@ -341,42 +314,40 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
payload: Optional[Dict] = None,
|
||||
) -> bool:
|
||||
"""Update a vector and its payload.
|
||||
|
||||
|
||||
Args:
|
||||
vector_id: ID of the vector to update
|
||||
vector: Optional new vector values
|
||||
payload: Optional new metadata payload
|
||||
|
||||
|
||||
Returns:
|
||||
bool: True if update was successful
|
||||
|
||||
|
||||
Raises:
|
||||
ValueError: If neither vector nor payload is provided
|
||||
GoogleAPIError: If the API call fails
|
||||
"""
|
||||
logger.debug("Starting update for vector_id: %s", vector_id)
|
||||
|
||||
|
||||
if vector is None and payload is None:
|
||||
raise ValueError("Either vector or payload must be provided for update")
|
||||
|
||||
|
||||
# First check if the vector exists
|
||||
try:
|
||||
existing = self.get(vector_id)
|
||||
if existing is None:
|
||||
logger.error("Vector ID not found: %s", vector_id)
|
||||
return False
|
||||
|
||||
|
||||
datapoint = self._create_datapoint(
|
||||
vector_id=vector_id,
|
||||
vector=vector if vector is not None else [],
|
||||
payload=payload
|
||||
vector_id=vector_id, vector=vector if vector is not None else [], payload=payload
|
||||
)
|
||||
|
||||
|
||||
logger.debug("Upserting datapoint: %s", datapoint)
|
||||
self.index.upsert_datapoints(datapoints=[datapoint])
|
||||
logger.debug("Update completed successfully")
|
||||
return True
|
||||
|
||||
|
||||
except google.api_core.exceptions.GoogleAPIError as e:
|
||||
logger.error("API error during update: %s", str(e))
|
||||
return False
|
||||
@@ -385,7 +356,6 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
logger.error("Stack trace: %s", traceback.format_exc())
|
||||
raise
|
||||
|
||||
|
||||
def get(self, vector_id: str) -> Optional[OutputData]:
|
||||
"""
|
||||
Retrieve a vector by ID.
|
||||
@@ -395,24 +365,17 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
Optional[OutputData]: Retrieved vector or None if not found.
|
||||
"""
|
||||
logger.debug("Starting get for vector_id: %s", vector_id)
|
||||
|
||||
|
||||
try:
|
||||
if not self.vector_search_api_endpoint:
|
||||
raise ValueError("vector_search_api_endpoint is required for get operation")
|
||||
|
||||
vector_search_client = aiplatform_v1.MatchServiceClient(
|
||||
client_options={
|
||||
"api_endpoint": self.vector_search_api_endpoint
|
||||
},
|
||||
)
|
||||
datapoint = aiplatform_v1.IndexDatapoint(
|
||||
datapoint_id=vector_id
|
||||
client_options={"api_endpoint": self.vector_search_api_endpoint},
|
||||
)
|
||||
datapoint = aiplatform_v1.IndexDatapoint(datapoint_id=vector_id)
|
||||
|
||||
query = aiplatform_v1.FindNeighborsRequest.Query(
|
||||
datapoint=datapoint,
|
||||
neighbor_count=1
|
||||
)
|
||||
query = aiplatform_v1.FindNeighborsRequest.Query(datapoint=datapoint, neighbor_count=1)
|
||||
request = aiplatform_v1.FindNeighborsRequest(
|
||||
index_endpoint=f"projects/{self.project_number}/locations/{self.region}/indexEndpoints/{self.endpoint_id}",
|
||||
deployed_index_id=self.deployment_index_id,
|
||||
@@ -423,41 +386,36 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
try:
|
||||
response = vector_search_client.find_neighbors(request)
|
||||
logger.debug("Got response")
|
||||
|
||||
|
||||
if response and response.nearest_neighbors:
|
||||
nearest = response.nearest_neighbors[0]
|
||||
if nearest.neighbors:
|
||||
neighbor = nearest.neighbors[0]
|
||||
|
||||
|
||||
payload = {}
|
||||
if hasattr(neighbor.datapoint, 'restricts'):
|
||||
if hasattr(neighbor.datapoint, "restricts"):
|
||||
for restrict in neighbor.datapoint.restricts:
|
||||
if restrict.allow_list:
|
||||
payload[restrict.namespace] = restrict.allow_list[0]
|
||||
|
||||
return OutputData(
|
||||
id=neighbor.datapoint.datapoint_id,
|
||||
score=neighbor.distance,
|
||||
payload=payload
|
||||
)
|
||||
|
||||
|
||||
return OutputData(id=neighbor.datapoint.datapoint_id, score=neighbor.distance, payload=payload)
|
||||
|
||||
logger.debug("No results found")
|
||||
return None
|
||||
|
||||
|
||||
except google.api_core.exceptions.NotFound:
|
||||
logger.debug("Datapoint not found")
|
||||
return None
|
||||
except google.api_core.exceptions.PermissionDenied as e:
|
||||
logger.error("Permission denied: %s", str(e))
|
||||
return None
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error occurred: %s", str(e))
|
||||
logger.error("Error type: %s", type(e))
|
||||
logger.error("Stack trace: %s", traceback.format_exc())
|
||||
raise
|
||||
|
||||
|
||||
def list_cols(self) -> List[str]:
|
||||
"""
|
||||
List all collections (indexes).
|
||||
@@ -466,7 +424,6 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
"""
|
||||
return [self.deployment_index_id]
|
||||
|
||||
|
||||
def delete_col(self):
|
||||
"""
|
||||
Delete a collection (index).
|
||||
@@ -475,7 +432,6 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
logger.warning("Delete collection operation is not supported for Google Matching Engine")
|
||||
pass
|
||||
|
||||
|
||||
def col_info(self) -> Dict:
|
||||
"""
|
||||
Get information about a collection (index).
|
||||
@@ -486,17 +442,16 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
"index_id": self.index_id,
|
||||
"endpoint_id": self.endpoint_id,
|
||||
"project_id": self.project_id,
|
||||
"region": self.region
|
||||
"region": self.region,
|
||||
}
|
||||
|
||||
|
||||
def list(self, filters: Optional[Dict] = None, limit: Optional[int] = None) -> List[List[OutputData]]:
|
||||
"""List vectors matching the given filters.
|
||||
|
||||
|
||||
Args:
|
||||
filters: Optional filters to apply
|
||||
limit: Optional maximum number of results to return
|
||||
|
||||
|
||||
Returns:
|
||||
List[List[OutputData]]: List of matching vectors wrapped in an extra array
|
||||
to match the interface
|
||||
@@ -504,36 +459,31 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
logger.debug("Starting list operation")
|
||||
logger.debug("Filters: %s", filters)
|
||||
logger.debug("Limit: %s", limit)
|
||||
|
||||
|
||||
try:
|
||||
# Use a zero vector for the search
|
||||
dimension = 768 # This should be configurable based on the model
|
||||
zero_vector = [0.0] * dimension
|
||||
|
||||
|
||||
# Use a large limit if none specified
|
||||
search_limit = limit if limit is not None else 10000
|
||||
|
||||
results = self.search(
|
||||
query=zero_vector,
|
||||
limit=search_limit,
|
||||
filters=filters
|
||||
)
|
||||
|
||||
|
||||
results = self.search(query=zero_vector, limit=search_limit, filters=filters)
|
||||
|
||||
logger.debug("Found %d results", len(results))
|
||||
return [results] # Wrap in extra array to match interface
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error in list operation: %s", str(e))
|
||||
logger.error("Stack trace: %s", traceback.format_exc())
|
||||
raise
|
||||
|
||||
|
||||
def create_col(self, name=None, vector_size=None, distance=None):
|
||||
"""
|
||||
Create a new collection. For Google Matching Engine, collections (indexes)
|
||||
Create a new collection. For Google Matching Engine, collections (indexes)
|
||||
are created through the Google Cloud Console or API separately.
|
||||
This method is a no-op since indexes are pre-created.
|
||||
|
||||
|
||||
Args:
|
||||
name: Ignored for Google Matching Engine
|
||||
vector_size: Ignored for Google Matching Engine
|
||||
@@ -543,41 +493,35 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
# This method is included only to satisfy the abstract base class
|
||||
pass
|
||||
|
||||
|
||||
def add(self, text: str, metadata: Optional[Dict] = None, user_id: Optional[str] = None) -> str:
|
||||
logger.debug("Starting add operation")
|
||||
logger.debug("Text: %s", text)
|
||||
logger.debug("Metadata: %s", metadata)
|
||||
logger.debug("User ID: %s", user_id)
|
||||
|
||||
|
||||
try:
|
||||
# Generate a unique ID for this entry
|
||||
vector_id = str(uuid.uuid4())
|
||||
|
||||
|
||||
# Create the payload with all necessary fields
|
||||
payload = {
|
||||
"data": text, # Store the text in the data field
|
||||
"user_id": user_id,
|
||||
**(metadata or {})
|
||||
**(metadata or {}),
|
||||
}
|
||||
|
||||
|
||||
# Get the embedding
|
||||
vector = self.embedder.embed_query(text)
|
||||
|
||||
|
||||
# Insert using the insert method
|
||||
self.insert(
|
||||
vectors=[vector],
|
||||
payloads=[payload],
|
||||
ids=[vector_id]
|
||||
)
|
||||
|
||||
self.insert(vectors=[vector], payloads=[payload], ids=[vector_id])
|
||||
|
||||
return vector_id
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error occurred: %s", str(e))
|
||||
raise
|
||||
|
||||
|
||||
def add_texts(
|
||||
self,
|
||||
texts: List[str],
|
||||
@@ -585,47 +529,45 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
ids: Optional[List[str]] = None,
|
||||
) -> List[str]:
|
||||
"""Add texts to the vector store.
|
||||
|
||||
|
||||
Args:
|
||||
texts: List of texts to add
|
||||
metadatas: Optional list of metadata dicts
|
||||
ids: Optional list of IDs to use
|
||||
|
||||
|
||||
Returns:
|
||||
List[str]: List of IDs of the added texts
|
||||
|
||||
|
||||
Raises:
|
||||
ValueError: If texts is empty or lengths don't match
|
||||
"""
|
||||
if not texts:
|
||||
raise ValueError("No texts provided")
|
||||
|
||||
|
||||
if metadatas and len(metadatas) != len(texts):
|
||||
raise ValueError(f"Number of metadata items ({len(metadatas)}) does not match number of texts ({len(texts)})")
|
||||
|
||||
raise ValueError(
|
||||
f"Number of metadata items ({len(metadatas)}) does not match number of texts ({len(texts)})"
|
||||
)
|
||||
|
||||
if ids and len(ids) != len(texts):
|
||||
raise ValueError(f"Number of ids ({len(ids)}) does not match number of texts ({len(texts)})")
|
||||
|
||||
|
||||
logger.debug("Starting add_texts operation")
|
||||
logger.debug("Number of texts: %d", len(texts))
|
||||
logger.debug("Has metadatas: %s", metadatas is not None)
|
||||
logger.debug("Has ids: %s", ids is not None)
|
||||
|
||||
|
||||
if ids is None:
|
||||
ids = [str(uuid.uuid4()) for _ in texts]
|
||||
|
||||
|
||||
try:
|
||||
# Get embeddings
|
||||
embeddings = self.embedder.embed_documents(texts)
|
||||
|
||||
|
||||
# Add to store
|
||||
self.insert(
|
||||
vectors=embeddings,
|
||||
payloads=metadatas if metadatas else [{}] * len(texts),
|
||||
ids=ids
|
||||
)
|
||||
self.insert(vectors=embeddings, payloads=metadatas if metadatas else [{}] * len(texts), ids=ids)
|
||||
return ids
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Error in add_texts: %s", str(e))
|
||||
logger.error("Stack trace: %s", traceback.format_exc())
|
||||
@@ -657,18 +599,12 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
logger.debug("Query: %s", query)
|
||||
logger.debug("k: %d", k)
|
||||
logger.debug("Filter: %s", filter)
|
||||
|
||||
|
||||
embedding = self.embedder.embed_query(query)
|
||||
results = self.search(query=embedding, limit=k, filters=filter)
|
||||
|
||||
|
||||
docs_and_scores = [
|
||||
(
|
||||
Document(
|
||||
page_content=result.payload.get("text", ""),
|
||||
metadata=result.payload
|
||||
),
|
||||
result.score
|
||||
)
|
||||
(Document(page_content=result.payload.get("text", ""), metadata=result.payload), result.score)
|
||||
for result in results
|
||||
]
|
||||
logger.debug("Found %d results", len(docs_and_scores))
|
||||
@@ -684,4 +620,3 @@ class GoogleMatchingEngine(VectorStoreBase):
|
||||
logger.debug("Starting similarity search")
|
||||
docs_and_scores = self.similarity_search_with_score(query, k, filter)
|
||||
return [doc for doc, _ in docs_and_scores]
|
||||
|
||||
|
||||
@@ -154,7 +154,9 @@ class Weaviate(VectorStoreBase):
|
||||
|
||||
batch.add_object(collection=self.collection_name, properties=data_object, uuid=object_id, vector=vector)
|
||||
|
||||
def search(self, query: List[float], limit: int = 5, filters: Optional[Dict] = None) -> List[OutputData]:
|
||||
def search(
|
||||
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
|
||||
) -> List[OutputData]:
|
||||
"""
|
||||
Search for similar vectors.
|
||||
"""
|
||||
@@ -167,7 +169,7 @@ class Weaviate(VectorStoreBase):
|
||||
combined_filter = Filter.all_of(filter_conditions) if filter_conditions else None
|
||||
response = collection.query.hybrid(
|
||||
query="",
|
||||
vector=query,
|
||||
vector=vectors,
|
||||
limit=limit,
|
||||
filters=combined_filter,
|
||||
return_properties=["hash", "created_at", "updated_at", "user_id", "agent_id", "run_id", "data", "category"],
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "mem0ai"
|
||||
version = "0.1.73"
|
||||
version = "0.1.74"
|
||||
description = "Long-term memory for AI Agents"
|
||||
authors = ["Mem0 <founders@mem0.ai>"]
|
||||
exclude = [
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch
|
||||
from mem0.embeddings.lmstudio import LMStudioEmbedding
|
||||
from mem0.configs.embeddings.base import BaseEmbedderConfig
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_lm_studio_client():
|
||||
with patch("mem0.embeddings.lmstudio.Client") as mock_lm_studio:
|
||||
mock_client = Mock()
|
||||
mock_client.list.return_value = {"models": [{"name": "nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf"}]}
|
||||
mock_lm_studio.return_value = mock_client
|
||||
yield mock_client
|
||||
|
||||
|
||||
def test_embed_text(mock_lm_studio_client):
|
||||
config = BaseEmbedderConfig(model="nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf", embedding_dims=512)
|
||||
embedder = LMStudioEmbedding(config)
|
||||
|
||||
mock_response = {"embedding": [0.1, 0.2, 0.3, 0.4, 0.5]}
|
||||
mock_lm_studio_client.embeddings.return_value = mock_response
|
||||
|
||||
text = "Sample text to embed."
|
||||
embedding = embedder.embed(text)
|
||||
|
||||
mock_lm_studio_client.embeddings.assert_called_once_with(model="nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf", prompt=text)
|
||||
|
||||
assert embedding == [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
|
||||
|
||||
def test_ensure_model_exists(mock_lm_studio_client):
|
||||
config = BaseEmbedderConfig(model="nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf", embedding_dims=512)
|
||||
embedder = LMStudioEmbedding(config)
|
||||
|
||||
mock_lm_studio_client.pull.assert_not_called()
|
||||
|
||||
mock_lm_studio_client.list.return_value = {"models": []}
|
||||
|
||||
embedder._ensure_model_exists()
|
||||
|
||||
mock_lm_studio_client.pull.assert_called_once_with("nomic-embed-text")
|
||||
@@ -0,0 +1,34 @@
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from mem0.configs.llms.base import BaseLlmConfig
|
||||
from mem0.llms.lmstudio import LMStudioLLM
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_lm_studio_client():
|
||||
with patch("mem0.llms.lmstudio.Client") as mock_lm_studio:
|
||||
mock_client = Mock()
|
||||
mock_client.list.return_value = {"models": [{"name": "lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf"}]}
|
||||
mock_lm_studio.return_value = mock_client
|
||||
yield mock_client
|
||||
|
||||
|
||||
def test_generate_response_without_tools(mock_lm_studio_client):
|
||||
config = BaseLlmConfig(model="lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf", temperature=0.7, max_tokens=100, top_p=1.0)
|
||||
llm = LMStudioLLM(config)
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Hello, how are you?"},
|
||||
]
|
||||
|
||||
mock_response = {"message": {"content": "I'm doing well, thank you for asking!"}}
|
||||
mock_lm_studio_client.chat.return_value = mock_response
|
||||
|
||||
response = llm.generate_response(messages)
|
||||
|
||||
mock_lm_studio_client.chat.assert_called_once_with(
|
||||
model="lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf", messages=messages, options={"temperature": 0.7, "num_predict": 100, "top_p": 1.0}
|
||||
)
|
||||
assert response == "I'm doing well, thank you for asking!"
|
||||
@@ -35,12 +35,11 @@ def test_search_vectors(chromadb_instance, mock_chromadb_client):
|
||||
}
|
||||
chromadb_instance.collection.query.return_value = mock_result
|
||||
|
||||
query = [[0.1, 0.2, 0.3]]
|
||||
results = chromadb_instance.search(query=query, limit=2)
|
||||
vectors = [[0.1, 0.2, 0.3]]
|
||||
results = chromadb_instance.search(query="", vectors=vectors, limit=2)
|
||||
|
||||
chromadb_instance.collection.query.assert_called_once_with(query_embeddings=query, where=None, n_results=2)
|
||||
chromadb_instance.collection.query.assert_called_once_with(query_embeddings=vectors, where=None, n_results=2)
|
||||
|
||||
print(results, type(results))
|
||||
assert len(results) == 2
|
||||
assert results[0].id == "id1"
|
||||
assert results[0].score == 0.1
|
||||
|
||||
@@ -196,8 +196,8 @@ class TestElasticsearchDB(unittest.TestCase):
|
||||
self.client_mock.search.return_value = mock_response
|
||||
|
||||
# Perform search
|
||||
query_vector = [0.1] * 1536
|
||||
results = self.es_db.search(query=query_vector, limit=5)
|
||||
vectors = [[0.1] * 1536]
|
||||
results = self.es_db.search(query="", vectors=vectors, limit=5)
|
||||
|
||||
# Verify search call
|
||||
self.client_mock.search.assert_called_once()
|
||||
@@ -210,7 +210,7 @@ class TestElasticsearchDB(unittest.TestCase):
|
||||
# Verify KNN query structure
|
||||
self.assertIn("knn", body)
|
||||
self.assertEqual(body["knn"]["field"], "vector")
|
||||
self.assertEqual(body["knn"]["query_vector"], query_vector)
|
||||
self.assertEqual(body["knn"]["query_vector"], vectors)
|
||||
self.assertEqual(body["knn"]["k"], 5)
|
||||
self.assertEqual(body["knn"]["num_candidates"], 10)
|
||||
|
||||
@@ -226,13 +226,13 @@ class TestElasticsearchDB(unittest.TestCase):
|
||||
self.es_db.custom_search_query.return_value = {"custom_key": "custom_value"}
|
||||
|
||||
# Perform search
|
||||
query_vector = [0.1] * 1536
|
||||
vectors = [[0.1] * 1536]
|
||||
limit = 5
|
||||
filters = {"key1": "value1"}
|
||||
self.es_db.search(query=query_vector, limit=limit, filters=filters)
|
||||
self.es_db.search(query="", vectors=vectors, limit=limit, filters=filters)
|
||||
|
||||
# Verify custom search query function was called
|
||||
self.es_db.custom_search_query.assert_called_once_with(query_vector, limit, filters)
|
||||
self.es_db.custom_search_query.assert_called_once_with(vectors, limit, filters)
|
||||
|
||||
# Verify custom search query was used
|
||||
self.client_mock.search.assert_called_once_with(index=self.es_db.collection_name, body={"custom_key": "custom_value"})
|
||||
|
||||
@@ -126,15 +126,15 @@ class TestOpenSearchDB(unittest.TestCase):
|
||||
def test_search(self):
|
||||
mock_response = {"hits": {"hits": [{"_id": "id1", "_score": 0.8, "_source": {"vector": [0.1] * 1536, "metadata": {"key1": "value1"}}}]}}
|
||||
self.client_mock.search.return_value = mock_response
|
||||
query_vector = [0.1] * 1536
|
||||
results = self.os_db.search(query=query_vector, limit=5)
|
||||
vectors = [[0.1] * 1536]
|
||||
results = self.os_db.search(query="", vectors=vectors, limit=5)
|
||||
self.client_mock.search.assert_called_once()
|
||||
search_args = self.client_mock.search.call_args[1]
|
||||
self.assertEqual(search_args["index"], "test_collection")
|
||||
body = search_args["body"]
|
||||
self.assertIn("knn", body["query"])
|
||||
self.assertIn("vector", body["query"]["knn"])
|
||||
self.assertEqual(body["query"]["knn"]["vector"]["vector"], query_vector)
|
||||
self.assertEqual(body["query"]["knn"]["vector"]["vector"], vectors)
|
||||
self.assertEqual(body["query"]["knn"]["vector"]["k"], 5)
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertEqual(results[0].id, "id1")
|
||||
|
||||
@@ -50,15 +50,15 @@ class TestQdrant(unittest.TestCase):
|
||||
self.assertEqual(points[0].payload, payloads[0])
|
||||
|
||||
def test_search(self):
|
||||
query_vector = [0.1, 0.2]
|
||||
vectors = [[0.1, 0.2]]
|
||||
mock_point = MagicMock(id=str(uuid.uuid4()), score=0.95, payload={"key": "value"})
|
||||
self.client_mock.query_points.return_value = MagicMock(points=[mock_point])
|
||||
|
||||
results = self.qdrant.search(query=query_vector, limit=1)
|
||||
results = self.qdrant.search(query="", vectors=vectors, limit=1)
|
||||
|
||||
self.client_mock.query_points.assert_called_once_with(
|
||||
collection_name="test_collection",
|
||||
query=query_vector,
|
||||
query=vectors,
|
||||
query_filter=None,
|
||||
limit=1,
|
||||
)
|
||||
|
||||
@@ -77,12 +77,12 @@ def test_search_vectors(supabase_instance, mock_collection):
|
||||
]
|
||||
mock_collection.query.return_value = mock_results
|
||||
|
||||
query = [0.1, 0.2, 0.3]
|
||||
vectors = [[0.1, 0.2, 0.3]]
|
||||
filters = {"category": "test"}
|
||||
results = supabase_instance.search(query=query, limit=2, filters=filters)
|
||||
results = supabase_instance.search(query="", vectors=vectors, limit=2, filters=filters)
|
||||
|
||||
mock_collection.query.assert_called_once_with(
|
||||
data=query,
|
||||
data=vectors,
|
||||
limit=2,
|
||||
filters={"category": {"$eq": "test"}},
|
||||
include_metadata=True,
|
||||
|
||||
@@ -73,12 +73,12 @@ def test_insert_vectors(vector_store, mock_vertex_ai):
|
||||
|
||||
def test_search_vectors(vector_store, mock_vertex_ai):
|
||||
"""Test searching vectors with filters"""
|
||||
query = [0.1, 0.2, 0.3]
|
||||
vectors = [[0.1, 0.2, 0.3]]
|
||||
filters = {"user_id": "test_user"}
|
||||
|
||||
mock_datapoint = Mock()
|
||||
mock_datapoint.datapoint_id = "test-id"
|
||||
mock_datapoint.feature_vector = query
|
||||
mock_datapoint.feature_vector = vectors
|
||||
|
||||
mock_restrict = Mock()
|
||||
mock_restrict.namespace = "user_id"
|
||||
@@ -96,11 +96,11 @@ def test_search_vectors(vector_store, mock_vertex_ai):
|
||||
|
||||
mock_vertex_ai['endpoint'].find_neighbors.return_value = [[mock_neighbor]]
|
||||
|
||||
results = vector_store.search(query=query, filters=filters, limit=1)
|
||||
results = vector_store.search(query="", vectors=vectors, filters=filters, limit=1)
|
||||
|
||||
mock_vertex_ai['endpoint'].find_neighbors.assert_called_once_with(
|
||||
deployed_index_id=vector_store.deployment_index_id,
|
||||
queries=[query],
|
||||
queries=[vectors],
|
||||
num_neighbors=1,
|
||||
filter=[Namespace("user_id", ["test_user"], [])],
|
||||
return_full_datapoint=True
|
||||
|
||||
@@ -147,8 +147,8 @@ class TestWeaviateDB(unittest.TestCase):
|
||||
self.client_mock.collections.get.return_value.query.hybrid = mock_hybrid
|
||||
mock_hybrid.return_value = mock_response
|
||||
|
||||
query_vector = [0.1] * 1536
|
||||
results = self.weaviate_db.search(query=query_vector, limit=5)
|
||||
vectors = [[0.1] * 1536]
|
||||
results = self.weaviate_db.search(query="", vectors=vectors, limit=5)
|
||||
|
||||
mock_hybrid.assert_called_once()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user