Store user_id in vectordb (#2466)
This commit is contained in:
@@ -3,6 +3,7 @@ import concurrent
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
import warnings
|
||||
from datetime import datetime
|
||||
@@ -18,7 +19,7 @@ from mem0.configs.prompts import (
|
||||
get_update_memory_messages,
|
||||
)
|
||||
from mem0.memory.base import MemoryBase
|
||||
from mem0.memory.setup import setup_config
|
||||
from mem0.memory.setup import mem0_dir, setup_config
|
||||
from mem0.memory.storage import SQLiteManager
|
||||
from mem0.memory.telemetry import capture_event
|
||||
from mem0.memory.utils import (
|
||||
@@ -62,6 +63,15 @@ class Memory(MemoryBase):
|
||||
self.graph = MemoryGraph(self.config)
|
||||
self.enable_graph = True
|
||||
|
||||
self.config.vector_store.config.collection_name = "mem0_migrations"
|
||||
if self.config.vector_store.provider in ["faiss", "qdrant"]:
|
||||
provider_path = f"migrations_{self.config.vector_store.provider}"
|
||||
self.config.vector_store.config.path = os.path.join(mem0_dir, provider_path)
|
||||
os.makedirs(self.config.vector_store.config.path, exist_ok=True)
|
||||
|
||||
self._telemetry_vector_store = VectorStoreFactory.create(
|
||||
self.config.vector_store.provider, self.config.vector_store.config
|
||||
)
|
||||
capture_event("mem0.init", self)
|
||||
|
||||
@classmethod
|
||||
|
||||
@@ -3,6 +3,7 @@ import os
|
||||
import uuid
|
||||
|
||||
# Set up the directory path
|
||||
VECTOR_ID = str(uuid.uuid4())
|
||||
home_dir = os.path.expanduser("~")
|
||||
mem0_dir = os.environ.get("MEM0_DIR") or os.path.join(home_dir, ".mem0")
|
||||
os.makedirs(mem0_dir, exist_ok=True)
|
||||
@@ -29,3 +30,29 @@ def get_user_id():
|
||||
return user_id
|
||||
except Exception:
|
||||
return "anonymous_user"
|
||||
|
||||
|
||||
def get_or_create_user_id(vector_store):
|
||||
"""Store user_id in vector store and return it."""
|
||||
user_id = get_user_id()
|
||||
|
||||
# Try to get existing user_id from vector store
|
||||
try:
|
||||
existing = vector_store.get(vector_id=VECTOR_ID)
|
||||
if existing and hasattr(existing, "payload") and existing.payload and "user_id" in existing.payload:
|
||||
return existing.payload["user_id"]
|
||||
except:
|
||||
pass
|
||||
|
||||
# If we get here, we need to insert the user_id
|
||||
try:
|
||||
dims = getattr(vector_store, "embedding_model_dims", 1)
|
||||
vector_store.insert(
|
||||
vectors=[[0.0] * dims],
|
||||
payloads=[{"user_id": user_id, "type": "user_identity"}],
|
||||
ids=[VECTOR_ID]
|
||||
)
|
||||
except:
|
||||
pass
|
||||
|
||||
return user_id
|
||||
|
||||
@@ -6,7 +6,7 @@ import sys
|
||||
from posthog import Posthog
|
||||
|
||||
import mem0
|
||||
from mem0.memory.setup import get_user_id, setup_config
|
||||
from mem0.memory.setup import get_or_create_user_id
|
||||
|
||||
MEM0_TELEMETRY = os.environ.get("MEM0_TELEMETRY", "True")
|
||||
|
||||
@@ -21,11 +21,11 @@ logging.getLogger("urllib3").setLevel(logging.CRITICAL + 1)
|
||||
|
||||
|
||||
class AnonymousTelemetry:
|
||||
def __init__(self, project_api_key, host):
|
||||
def __init__(self, project_api_key, host, vector_store=None):
|
||||
self.posthog = Posthog(project_api_key=project_api_key, host=host)
|
||||
# Call setup config to ensure that the user_id is generated
|
||||
setup_config()
|
||||
self.user_id = get_user_id()
|
||||
|
||||
self.user_id = get_or_create_user_id(vector_store)
|
||||
|
||||
if not MEM0_TELEMETRY:
|
||||
self.posthog.disabled = True
|
||||
|
||||
@@ -50,14 +50,12 @@ class AnonymousTelemetry:
|
||||
self.posthog.shutdown()
|
||||
|
||||
|
||||
# Initialize AnonymousTelemetry
|
||||
telemetry = AnonymousTelemetry(
|
||||
project_api_key="phc_hgJkUVJFYtmaJqrvf6CYN67TIQ8yhXAkWzUn9AMU4yX",
|
||||
host="https://us.i.posthog.com",
|
||||
)
|
||||
|
||||
|
||||
def capture_event(event_name, memory_instance, additional_data=None):
|
||||
global telemetry
|
||||
|
||||
# For OSS, we use the telemetry vector store to store the user_id
|
||||
telemetry = AnonymousTelemetry(project_api_key="phc_hgJkUVJFYtmaJqrvf6CYN67TIQ8yhXAkWzUn9AMU4yX", host="https://us.i.posthog.com", vector_store=memory_instance._telemetry_vector_store if hasattr(memory_instance, "_telemetry_vector_store") else None)
|
||||
|
||||
event_data = {
|
||||
"collection": memory_instance.collection_name,
|
||||
"vector_size": memory_instance.embedding_model.config.embedding_dims,
|
||||
|
||||
Reference in New Issue
Block a user