forked from SpeedyFoxAi/jarvis-memory
Initial commit: Jarvis Memory system
This commit is contained in:
189
skills/qdrant-memory/scripts/search_memories.py
Executable file
189
skills/qdrant-memory/scripts/search_memories.py
Executable file
@@ -0,0 +1,189 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Search memories by semantic similarity in Qdrant
|
||||
Usage: search_memories.py "Query text" [--limit 5] [--filter-tag tag] [--track-access]
|
||||
|
||||
Now with access tracking - updates access_count and last_accessed when memories are retrieved.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import urllib.request
|
||||
from datetime import datetime
|
||||
|
||||
import os
|
||||
|
||||
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")
|
||||
|
||||
def get_embedding(text):
|
||||
"""Generate embedding using snowflake-arctic-embed2 via Ollama"""
|
||||
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"Error generating embedding: {e}", file=sys.stderr)
|
||||
return None
|
||||
|
||||
def update_access_stats(point_id, current_payload):
|
||||
"""Update access_count and last_accessed for a memory"""
|
||||
|
||||
# Get current values or defaults
|
||||
access_count = current_payload.get("access_count", 0) + 1
|
||||
last_accessed = datetime.now().isoformat()
|
||||
|
||||
# Prepare update payload
|
||||
update_body = {
|
||||
"points": [
|
||||
{
|
||||
"id": point_id,
|
||||
"payload": {
|
||||
"access_count": access_count,
|
||||
"last_accessed": last_accessed
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
req = urllib.request.Request(
|
||||
f"{QDRANT_URL}/collections/{COLLECTION_NAME}/points/payload?wait=true",
|
||||
data=json.dumps(update_body).encode(),
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="PUT"
|
||||
)
|
||||
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=5) as response:
|
||||
result = json.loads(response.read().decode())
|
||||
return result.get("status") == "ok"
|
||||
except Exception as e:
|
||||
# Silently fail - don't break search if update fails
|
||||
return False
|
||||
|
||||
def search_memories(query_vector, limit=5, tag_filter=None, track_access=True):
|
||||
"""Search memories in Qdrant with optional access tracking"""
|
||||
|
||||
search_body = {
|
||||
"vector": query_vector,
|
||||
"limit": limit,
|
||||
"with_payload": True,
|
||||
"with_vector": False
|
||||
}
|
||||
|
||||
# Add filter if tag specified
|
||||
if tag_filter:
|
||||
search_body["filter"] = {
|
||||
"must": [
|
||||
{
|
||||
"key": "tags",
|
||||
"match": {
|
||||
"value": tag_filter
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
req = urllib.request.Request(
|
||||
f"{QDRANT_URL}/collections/{COLLECTION_NAME}/points/search",
|
||||
data=json.dumps(search_body).encode(),
|
||||
headers={"Content-Type": "application/json"}
|
||||
)
|
||||
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=10) as response:
|
||||
result = json.loads(response.read().decode())
|
||||
results = result.get("result", [])
|
||||
|
||||
# Track access for retrieved memories
|
||||
if track_access and results:
|
||||
for r in results:
|
||||
point_id = r.get("id")
|
||||
payload = r.get("payload", {})
|
||||
if point_id:
|
||||
update_access_stats(point_id, payload)
|
||||
|
||||
return results
|
||||
except Exception as e:
|
||||
print(f"Error searching memories: {e}", file=sys.stderr)
|
||||
return []
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Search memories by semantic similarity")
|
||||
parser.add_argument("query", help="Search query text")
|
||||
parser.add_argument("--limit", type=int, default=5, help="Number of results (default: 5)")
|
||||
parser.add_argument("--filter-tag", help="Filter by tag")
|
||||
parser.add_argument("--json", action="store_true", help="Output as JSON")
|
||||
parser.add_argument("--no-track", action="store_true", help="Don't update access stats")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f"Generating query embedding...", file=sys.stderr)
|
||||
query_vector = get_embedding(args.query)
|
||||
|
||||
if query_vector is None:
|
||||
print("❌ Failed to generate embedding", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
print(f"Searching Qdrant...", file=sys.stderr)
|
||||
results = search_memories(query_vector, args.limit, args.filter_tag, track_access=not args.no_track)
|
||||
|
||||
if not results:
|
||||
print("No matching memories found.")
|
||||
sys.exit(0)
|
||||
|
||||
if args.json:
|
||||
# JSON output with all metadata
|
||||
output = []
|
||||
for r in results:
|
||||
payload = r["payload"]
|
||||
output.append({
|
||||
"id": r.get("id"),
|
||||
"score": r["score"],
|
||||
"text": payload.get("text", ""),
|
||||
"date": payload.get("date", ""),
|
||||
"tags": payload.get("tags", []),
|
||||
"importance": payload.get("importance", "medium"),
|
||||
"confidence": payload.get("confidence", "medium"),
|
||||
"verified": payload.get("verified", False),
|
||||
"source_type": payload.get("source_type", "inferred"),
|
||||
"access_count": payload.get("access_count", 0),
|
||||
"last_accessed": payload.get("last_accessed", ""),
|
||||
"expires_at": payload.get("expires_at", None)
|
||||
})
|
||||
print(json.dumps(output, indent=2))
|
||||
else:
|
||||
# Human-readable output
|
||||
print(f"\n🔍 Found {len(results)} similar memories:\n")
|
||||
for i, r in enumerate(results, 1):
|
||||
payload = r["payload"]
|
||||
score = r["score"]
|
||||
text = payload.get("text", "")[:200]
|
||||
if len(payload.get("text", "")) > 200:
|
||||
text += "..."
|
||||
date = payload.get("date", "unknown")
|
||||
tags = ", ".join(payload.get("tags", []))
|
||||
importance = payload.get("importance", "medium")
|
||||
access_count = payload.get("access_count", 0)
|
||||
verified = "✓" if payload.get("verified", False) else "?"
|
||||
|
||||
print(f"{i}. [{date}] (score: {score:.3f}) [{importance}] {verified}")
|
||||
print(f" {text}")
|
||||
if tags:
|
||||
print(f" Tags: {tags}")
|
||||
if access_count > 0:
|
||||
print(f" Accessed: {access_count} times")
|
||||
print()
|
||||
Reference in New Issue
Block a user