forked from SpeedyFoxAi/jarvis-memory
Initial commit: Jarvis Memory system
This commit is contained in:
388
skills/qdrant-memory/scripts/auto_store.py
Executable file
388
skills/qdrant-memory/scripts/auto_store.py
Executable file
@@ -0,0 +1,388 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Auto Conversation Memory - TRUE Mem0-style Full Context Storage
|
||||
|
||||
User-centric memory - all conversations link to persistent user_id.
|
||||
NOT session/chat-centric like old version.
|
||||
|
||||
Features:
|
||||
- Persistent user_id (e.g., "rob") across all conversations
|
||||
- Cross-conversation retrieval (find memories from any chat)
|
||||
- Automatic conversation threading
|
||||
- Deduplication
|
||||
- Mem0-style: memories belong to USER, not to session
|
||||
|
||||
Usage:
|
||||
python3 scripts/auto_store.py "user_message" "ai_response" \
|
||||
--user-id "rob" \
|
||||
--conversation-id <uuid> \
|
||||
--turn <n>
|
||||
|
||||
Mem0 Architecture:
|
||||
- user_id: "rob" (persistent across all your chats)
|
||||
- conversation_id: Groups turns within one conversation
|
||||
- session_id: Optional - tracks specific chat instance
|
||||
- Retrieved by: user_id + semantic similarity (NOT session_id)
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import urllib.request
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from typing import List, Optional, Dict, Any
|
||||
|
||||
QDRANT_URL = os.getenv("QDRANT_URL", "http://127.0.0.1:6333")
|
||||
COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "kimi_memories")
|
||||
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://127.0.0.1:11434/v1")
|
||||
|
||||
# In-memory cache for deduplication (per process)
|
||||
_recent_hashes = set()
|
||||
|
||||
def get_content_hash(user_msg: str, ai_response: str) -> str:
|
||||
"""Generate hash for deduplication (stable across platforms)."""
|
||||
content = f"{user_msg.strip()}::{ai_response.strip()}".encode("utf-8", errors="replace")
|
||||
return hashlib.sha256(content).hexdigest()
|
||||
|
||||
def is_duplicate(user_id: str, user_msg: str, ai_response: str) -> bool:
|
||||
"""
|
||||
Check if this conversation turn already exists for this user.
|
||||
Uses: user_id + content_hash
|
||||
"""
|
||||
content_hash = get_content_hash(user_msg, ai_response)
|
||||
|
||||
# Check in-memory cache first
|
||||
if content_hash in _recent_hashes:
|
||||
return True
|
||||
|
||||
# Check Qdrant for existing entry with this user_id + content_hash
|
||||
try:
|
||||
search_body = {
|
||||
"filter": {
|
||||
"must": [
|
||||
{"key": "user_id", "match": {"value": user_id}},
|
||||
{"key": "content_hash", "match": {"value": content_hash}}
|
||||
]
|
||||
},
|
||||
"limit": 1,
|
||||
"with_payload": False
|
||||
}
|
||||
|
||||
req = urllib.request.Request(
|
||||
f"{QDRANT_URL}/collections/{COLLECTION_NAME}/points/scroll",
|
||||
data=json.dumps(search_body).encode(),
|
||||
headers={"Content-Type": "application/json"}
|
||||
)
|
||||
|
||||
with urllib.request.urlopen(req, timeout=10) as response:
|
||||
result = json.loads(response.read().decode())
|
||||
points = result.get("result", {}).get("points", [])
|
||||
if len(points) > 0:
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return False
|
||||
|
||||
def mark_stored(user_msg: str, ai_response: str):
|
||||
"""Mark content as stored in memory cache"""
|
||||
content_hash = get_content_hash(user_msg, ai_response)
|
||||
_recent_hashes.add(content_hash)
|
||||
if len(_recent_hashes) > 1000:
|
||||
_recent_hashes.clear()
|
||||
|
||||
def get_embedding(text: str) -> Optional[List[float]]:
|
||||
"""Generate embedding using snowflake-arctic-embed2"""
|
||||
data = json.dumps({
|
||||
"model": "snowflake-arctic-embed2",
|
||||
"input": text[:8192]
|
||||
}).encode()
|
||||
|
||||
req = urllib.request.Request(
|
||||
f"{OLLAMA_URL}/embeddings",
|
||||
data=data,
|
||||
headers={"Content-Type": "application/json"}
|
||||
)
|
||||
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=30) as response:
|
||||
result = json.loads(response.read().decode())
|
||||
return result["data"][0]["embedding"]
|
||||
except Exception as e:
|
||||
print(f"[AutoMemory] Embedding error: {e}", file=sys.stderr)
|
||||
return None
|
||||
|
||||
def generate_conversation_summary(user_msg: str, ai_response: str) -> str:
|
||||
"""Generate a searchable summary of the conversation turn"""
|
||||
summary = f"Q: {user_msg[:200]} A: {ai_response[:300]}"
|
||||
return summary
|
||||
|
||||
def store_memory_point(
|
||||
user_id: str,
|
||||
text: str,
|
||||
speaker: str,
|
||||
date_str: str,
|
||||
conversation_id: str,
|
||||
turn_number: int,
|
||||
session_id: Optional[str],
|
||||
tags: List[str],
|
||||
importance: str = "medium",
|
||||
content_hash: Optional[str] = None
|
||||
) -> Optional[str]:
|
||||
"""Store a single memory point to Qdrant with user_id"""
|
||||
|
||||
embedding = get_embedding(text)
|
||||
if embedding is None:
|
||||
return None
|
||||
|
||||
point_id = str(uuid.uuid4())
|
||||
|
||||
payload = {
|
||||
# MEM0-STYLE: user_id is PRIMARY key
|
||||
"user_id": user_id,
|
||||
"text": text,
|
||||
"date": date_str,
|
||||
"tags": tags,
|
||||
"importance": importance,
|
||||
"source": "conversation_auto",
|
||||
"source_type": "user" if speaker == "user" else "assistant",
|
||||
"category": "Full Conversation",
|
||||
"confidence": "high",
|
||||
"verified": True,
|
||||
"created_at": datetime.now().isoformat(),
|
||||
"access_count": 0,
|
||||
"last_accessed": datetime.now().isoformat(),
|
||||
"conversation_id": conversation_id,
|
||||
"turn_number": turn_number,
|
||||
"session_id": session_id or ""
|
||||
}
|
||||
|
||||
if content_hash:
|
||||
payload["content_hash"] = content_hash
|
||||
|
||||
upsert_data = {
|
||||
"points": [{
|
||||
"id": point_id,
|
||||
"vector": embedding,
|
||||
"payload": payload
|
||||
}]
|
||||
}
|
||||
|
||||
req = urllib.request.Request(
|
||||
f"{QDRANT_URL}/collections/{COLLECTION_NAME}/points?wait=true",
|
||||
data=json.dumps(upsert_data).encode(),
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="PUT"
|
||||
)
|
||||
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=10) as response:
|
||||
result = json.loads(response.read().decode())
|
||||
if result.get("status") == "ok":
|
||||
return point_id
|
||||
except Exception as e:
|
||||
print(f"[AutoMemory] Storage error: {e}", file=sys.stderr)
|
||||
|
||||
return None
|
||||
|
||||
def store_conversation_turn(
|
||||
user_id: str,
|
||||
user_message: str,
|
||||
ai_response: str,
|
||||
conversation_id: Optional[str] = None,
|
||||
turn_number: Optional[int] = None,
|
||||
session_id: Optional[str] = None,
|
||||
date_str: Optional[str] = None,
|
||||
skip_if_duplicate: bool = True
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Store a full conversation turn to Qdrant (Mem0-style)
|
||||
|
||||
Args:
|
||||
user_id: PERSISTENT user identifier (e.g., "rob") - REQUIRED
|
||||
user_message: User's message
|
||||
ai_response: AI's response
|
||||
conversation_id: Groups related turns (auto-generated if None)
|
||||
turn_number: Sequential turn number
|
||||
session_id: Optional chat session identifier
|
||||
date_str: Date in YYYY-MM-DD format
|
||||
|
||||
Returns:
|
||||
dict with success status and memory IDs
|
||||
"""
|
||||
if not user_id:
|
||||
raise ValueError("user_id is required for Mem0-style storage")
|
||||
|
||||
if date_str is None:
|
||||
date_str = datetime.now().strftime("%Y-%m-%d")
|
||||
|
||||
# Check for duplicates (per user)
|
||||
if skip_if_duplicate and is_duplicate(user_id, user_message, ai_response):
|
||||
return {
|
||||
"user_point_id": None,
|
||||
"ai_point_id": None,
|
||||
"user_id": user_id,
|
||||
"conversation_id": conversation_id or "",
|
||||
"turn_number": turn_number or 1,
|
||||
"success": True,
|
||||
"skipped": True
|
||||
}
|
||||
|
||||
if conversation_id is None:
|
||||
conversation_id = str(uuid.uuid4())
|
||||
|
||||
if turn_number is None:
|
||||
turn_number = 1
|
||||
|
||||
# Tags include user_id for easy filtering
|
||||
tags = [
|
||||
"conversation",
|
||||
f"user:{user_id}",
|
||||
date_str
|
||||
]
|
||||
|
||||
if session_id:
|
||||
tags.append(f"session:{session_id[:8]}")
|
||||
|
||||
# Determine importance
|
||||
importance = "high" if any(kw in (user_message + ai_response).lower()
|
||||
for kw in ["remember", "important", "always", "never", "rule"]) else "medium"
|
||||
|
||||
content_hash = get_content_hash(user_message, ai_response)
|
||||
|
||||
# Store user message
|
||||
user_text = f"[{user_id}]: {user_message}"
|
||||
user_id_point = store_memory_point(
|
||||
user_id=user_id,
|
||||
text=user_text,
|
||||
speaker="user",
|
||||
date_str=date_str,
|
||||
conversation_id=conversation_id,
|
||||
turn_number=turn_number,
|
||||
session_id=session_id,
|
||||
tags=tags + ["user-message"],
|
||||
importance=importance,
|
||||
content_hash=content_hash
|
||||
)
|
||||
|
||||
# Store AI response
|
||||
ai_text = f"[Kimi]: {ai_response}"
|
||||
ai_id_point = store_memory_point(
|
||||
user_id=user_id,
|
||||
text=ai_text,
|
||||
speaker="assistant",
|
||||
date_str=date_str,
|
||||
conversation_id=conversation_id,
|
||||
turn_number=turn_number,
|
||||
session_id=session_id,
|
||||
tags=tags + ["ai-response"],
|
||||
importance=importance,
|
||||
content_hash=content_hash
|
||||
)
|
||||
|
||||
# Store summary
|
||||
summary = generate_conversation_summary(user_message, ai_response)
|
||||
summary_text = f"[Turn {turn_number}] {summary}"
|
||||
|
||||
summary_embedding = get_embedding(summary_text)
|
||||
if summary_embedding:
|
||||
summary_id = str(uuid.uuid4())
|
||||
summary_payload = {
|
||||
"user_id": user_id,
|
||||
"text": summary_text,
|
||||
"date": date_str,
|
||||
"tags": tags + ["summary", "combined"],
|
||||
"importance": importance,
|
||||
"source": "conversation_summary",
|
||||
"source_type": "system",
|
||||
"category": "Conversation Summary",
|
||||
"confidence": "high",
|
||||
"verified": True,
|
||||
"created_at": datetime.now().isoformat(),
|
||||
"access_count": 0,
|
||||
"last_accessed": datetime.now().isoformat(),
|
||||
"conversation_id": conversation_id,
|
||||
"turn_number": turn_number,
|
||||
"session_id": session_id or "",
|
||||
"content_hash": content_hash,
|
||||
"user_message": user_message[:500],
|
||||
"ai_response": ai_response[:800]
|
||||
}
|
||||
|
||||
upsert_data = {
|
||||
"points": [{
|
||||
"id": summary_id,
|
||||
"vector": summary_embedding,
|
||||
"payload": summary_payload
|
||||
}]
|
||||
}
|
||||
|
||||
req = urllib.request.Request(
|
||||
f"{QDRANT_URL}/collections/{COLLECTION_NAME}/points?wait=true",
|
||||
data=json.dumps(upsert_data).encode(),
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="PUT"
|
||||
)
|
||||
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=10) as response:
|
||||
json.loads(response.read().decode())
|
||||
except Exception as e:
|
||||
print(f"[AutoMemory] Summary storage error: {e}", file=sys.stderr)
|
||||
|
||||
# Mark as stored
|
||||
if user_id_point and ai_id_point:
|
||||
mark_stored(user_message, ai_response)
|
||||
|
||||
return {
|
||||
"user_point_id": user_id_point,
|
||||
"ai_point_id": ai_id_point,
|
||||
"user_id": user_id,
|
||||
"conversation_id": conversation_id,
|
||||
"turn_number": turn_number,
|
||||
"success": bool(user_id_point and ai_id_point),
|
||||
"skipped": False
|
||||
}
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Auto-store conversation turns to Qdrant (TRUE Mem0-style with user_id)"
|
||||
)
|
||||
parser.add_argument("user_message", help="The user's message")
|
||||
parser.add_argument("ai_response", help="The AI's response")
|
||||
parser.add_argument("--user-id", required=True,
|
||||
help="REQUIRED: Persistent user ID (e.g., 'rob')")
|
||||
parser.add_argument("--conversation-id",
|
||||
help="Conversation ID for threading (auto-generated if not provided)")
|
||||
parser.add_argument("--turn", type=int, help="Turn number in conversation")
|
||||
parser.add_argument("--session-id",
|
||||
help="Optional: Session/chat instance ID")
|
||||
parser.add_argument("--date", default=datetime.now().strftime("%Y-%m-%d"),
|
||||
help="Date in YYYY-MM-DD format")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
result = store_conversation_turn(
|
||||
user_id=args.user_id,
|
||||
user_message=args.user_message,
|
||||
ai_response=args.ai_response,
|
||||
conversation_id=args.conversation_id,
|
||||
turn_number=args.turn,
|
||||
session_id=args.session_id,
|
||||
date_str=args.date
|
||||
)
|
||||
|
||||
if result.get("skipped"):
|
||||
print(f"⚡ Skipped duplicate (already stored for user {result['user_id']})")
|
||||
elif result["success"]:
|
||||
print(f"✅ Stored for user '{result['user_id']}' turn {result['turn_number']}")
|
||||
print(f" Conversation: {result['conversation_id'][:8]}...")
|
||||
else:
|
||||
print("❌ Failed to store conversation", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user