Initial commit: Jarvis Memory system
This commit is contained in:
242
skills/mem-redis/scripts/search_mem.py
Executable file
242
skills/mem-redis/scripts/search_mem.py
Executable file
@@ -0,0 +1,242 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Search memory: First Redis (exact), then Qdrant (semantic).
|
||||
|
||||
Usage: python3 search_mem.py "your search query" [--limit 10] [--user-id rob]
|
||||
|
||||
Searches:
|
||||
1. Redis (mem:{user_id}) - exact text match in recent buffer
|
||||
2. Qdrant (kimi_memories) - semantic similarity search
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import redis
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
|
||||
# Config
|
||||
REDIS_HOST = os.getenv("REDIS_HOST", "10.0.0.36")
|
||||
REDIS_PORT = int(os.getenv("REDIS_PORT", "6379"))
|
||||
USER_ID = os.getenv("USER_ID", "yourname")
|
||||
|
||||
QDRANT_URL = os.getenv("QDRANT_URL", "http://10.0.0.40:6333")
|
||||
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://10.0.0.10:11434/v1")
|
||||
|
||||
def search_redis(query, user_id, limit=20):
|
||||
"""Search Redis buffer for exact text matches."""
|
||||
try:
|
||||
r = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, decode_responses=True)
|
||||
key = f"mem:{user_id}"
|
||||
|
||||
# Get all items from list
|
||||
items = r.lrange(key, 0, -1)
|
||||
if not items:
|
||||
return []
|
||||
|
||||
query_lower = query.lower()
|
||||
matches = []
|
||||
|
||||
for item in items:
|
||||
try:
|
||||
turn = json.loads(item)
|
||||
content = turn.get('content', '').lower()
|
||||
if query_lower in content:
|
||||
matches.append({
|
||||
'source': 'redis',
|
||||
'turn': turn.get('turn'),
|
||||
'role': turn.get('role'),
|
||||
'content': turn.get('content'),
|
||||
'timestamp': turn.get('timestamp'),
|
||||
'score': 'exact'
|
||||
})
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
# Sort by turn number descending (newest first)
|
||||
matches.sort(key=lambda x: x.get('turn', 0), reverse=True)
|
||||
return matches[:limit]
|
||||
except Exception as e:
|
||||
print(f"Redis search error: {e}", file=sys.stderr)
|
||||
return []
|
||||
|
||||
def get_embedding(text):
|
||||
"""Get embedding from Ollama."""
|
||||
import urllib.request
|
||||
|
||||
payload = json.dumps({
|
||||
"model": "snowflake-arctic-embed2",
|
||||
"input": text
|
||||
}).encode()
|
||||
|
||||
req = urllib.request.Request(
|
||||
f"{OLLAMA_URL}/embeddings",
|
||||
data=payload,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST"
|
||||
)
|
||||
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||||
result = json.loads(resp.read().decode())
|
||||
return result.get('data', [{}])[0].get('embedding')
|
||||
except Exception as e:
|
||||
print(f"Embedding error: {e}", file=sys.stderr)
|
||||
return None
|
||||
|
||||
def search_qdrant(query, user_id, limit=10):
|
||||
"""Search Qdrant for semantic similarity."""
|
||||
import urllib.request
|
||||
|
||||
embedding = get_embedding(query)
|
||||
if not embedding:
|
||||
return []
|
||||
|
||||
payload = json.dumps({
|
||||
"vector": embedding,
|
||||
"limit": limit,
|
||||
"with_payload": True,
|
||||
"filter": {
|
||||
"must": [
|
||||
{"key": "user_id", "match": {"value": user_id}}
|
||||
]
|
||||
}
|
||||
}).encode()
|
||||
|
||||
req = urllib.request.Request(
|
||||
f"{QDRANT_URL}/collections/kimi_memories/points/search",
|
||||
data=payload,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST"
|
||||
)
|
||||
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||||
result = json.loads(resp.read().decode())
|
||||
points = result.get('result', [])
|
||||
|
||||
matches = []
|
||||
for point in points:
|
||||
payload = point.get('payload', {})
|
||||
matches.append({
|
||||
'source': 'qdrant',
|
||||
'score': round(point.get('score', 0), 3),
|
||||
'turn': payload.get('turn_number'),
|
||||
'role': payload.get('role'),
|
||||
'content': payload.get('user_message') or payload.get('content', ''),
|
||||
'ai_response': payload.get('ai_response', ''),
|
||||
'timestamp': payload.get('timestamp'),
|
||||
'conversation_id': payload.get('conversation_id')
|
||||
})
|
||||
return matches
|
||||
except Exception as e:
|
||||
print(f"Qdrant search error: {e}", file=sys.stderr)
|
||||
return []
|
||||
|
||||
def format_result(result, index):
|
||||
"""Format a single search result."""
|
||||
source = result.get('source', 'unknown')
|
||||
role = result.get('role', 'unknown')
|
||||
turn = result.get('turn', '?')
|
||||
score = result.get('score', '?')
|
||||
|
||||
content = result.get('content', '')
|
||||
if len(content) > 200:
|
||||
content = content[:200] + "..."
|
||||
|
||||
# Role emoji
|
||||
role_emoji = "👤" if role == "user" else "🤖"
|
||||
|
||||
# Source indicator
|
||||
source_icon = "🔴" if source == "redis" else "🔵"
|
||||
|
||||
lines = [
|
||||
f"{source_icon} [{index}] Turn {turn} ({role}):",
|
||||
f" {role_emoji} {content}"
|
||||
]
|
||||
|
||||
if source == "qdrant" and result.get('ai_response'):
|
||||
ai_resp = result['ai_response'][:150]
|
||||
if len(result['ai_response']) > 150:
|
||||
ai_resp += "..."
|
||||
lines.append(f" 💬 AI: {ai_resp}")
|
||||
|
||||
if score != 'exact':
|
||||
lines.append(f" 📊 Score: {score}")
|
||||
else:
|
||||
lines.append(f" 📊 Match: exact (Redis)")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Search memory: Redis first, then Qdrant")
|
||||
parser.add_argument("query", help="Search query")
|
||||
parser.add_argument("--limit", type=int, default=10, help="Results per source (default: 10)")
|
||||
parser.add_argument("--user-id", default=USER_ID, help="User ID")
|
||||
parser.add_argument("--redis-only", action="store_true", help="Only search Redis")
|
||||
parser.add_argument("--qdrant-only", action="store_true", help="Only search Qdrant")
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f"🔍 Searching for: \"{args.query}\"\n")
|
||||
|
||||
all_results = []
|
||||
|
||||
# Search Redis first (unless qdrant-only)
|
||||
if not args.qdrant_only:
|
||||
print("📍 Searching Redis (exact match)...")
|
||||
redis_results = search_redis(args.query, args.user_id, limit=args.limit)
|
||||
if redis_results:
|
||||
print(f"✅ Found {len(redis_results)} matches in Redis\n")
|
||||
all_results.extend(redis_results)
|
||||
else:
|
||||
print("❌ No exact matches in Redis\n")
|
||||
|
||||
# Search Qdrant (unless redis-only)
|
||||
if not args.redis_only:
|
||||
print("🧠 Searching Qdrant (semantic similarity)...")
|
||||
qdrant_results = search_qdrant(args.query, args.user_id, limit=args.limit)
|
||||
if qdrant_results:
|
||||
print(f"✅ Found {len(qdrant_results)} matches in Qdrant\n")
|
||||
all_results.extend(qdrant_results)
|
||||
else:
|
||||
print("❌ No semantic matches in Qdrant\n")
|
||||
|
||||
# Display results
|
||||
if not all_results:
|
||||
print("No results found in either Redis or Qdrant.")
|
||||
sys.exit(0)
|
||||
|
||||
print(f"=== Search Results ({len(all_results)} total) ===\n")
|
||||
|
||||
# Sort: Redis first (chronological), then Qdrant (by score)
|
||||
redis_sorted = [r for r in all_results if r['source'] == 'redis']
|
||||
qdrant_sorted = sorted(
|
||||
[r for r in all_results if r['source'] == 'qdrant'],
|
||||
key=lambda x: x.get('score', 0),
|
||||
reverse=True
|
||||
)
|
||||
|
||||
# Display Redis results first
|
||||
if redis_sorted:
|
||||
print("🔴 FROM REDIS (Recent Buffer):\n")
|
||||
for i, result in enumerate(redis_sorted, 1):
|
||||
print(format_result(result, i))
|
||||
print()
|
||||
|
||||
# Then Qdrant results
|
||||
if qdrant_sorted:
|
||||
print("🔵 FROM QDRANT (Long-term Memory):\n")
|
||||
for i, result in enumerate(qdrant_sorted, len(redis_sorted) + 1):
|
||||
print(format_result(result, i))
|
||||
print()
|
||||
|
||||
print(f"=== {len(all_results)} results ===")
|
||||
if redis_sorted:
|
||||
print(f" 🔴 Redis: {len(redis_sorted)} (exact, recent)")
|
||||
if qdrant_sorted:
|
||||
print(f" 🔵 Qdrant: {len(qdrant_sorted)} (semantic, long-term)")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user