103 lines
2.7 KiB
Python
103 lines
2.7 KiB
Python
"""
|
|
Memory Vector Search CLI
|
|
========================
|
|
Search memories using semantic similarity.
|
|
|
|
Usage:
|
|
python search_memories.py "your query here"
|
|
python search_memories.py --interactive
|
|
"""
|
|
|
|
import sys
|
|
import sys
|
|
import os
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
import requests
|
|
from memory_vector import search_memories
|
|
|
|
OLLAMA_URL = "http://localhost:11434"
|
|
EMBED_MODEL = "nomic-embed-text"
|
|
|
|
|
|
def get_embedding(text: str) -> list:
|
|
"""Generate embedding for query text."""
|
|
response = requests.post(
|
|
f"{OLLAMA_URL}/api/embeddings",
|
|
json={"model": EMBED_MODEL, "prompt": text},
|
|
timeout=30
|
|
)
|
|
response.raise_for_status()
|
|
return response.json()["embedding"]
|
|
|
|
|
|
def search(query: str, k: int = 5):
|
|
"""Search memories by query."""
|
|
print(f"Searching: '{query}'\n")
|
|
|
|
# Generate embedding for query
|
|
print("Generating embedding...")
|
|
query_embedding = get_embedding(query)
|
|
|
|
# Search database
|
|
print(f"Searching {k} nearest neighbors...\n")
|
|
results = search_memories(query_embedding, k=k)
|
|
|
|
if not results:
|
|
print("No results found.")
|
|
return
|
|
|
|
print("=" * 60)
|
|
print("RESULTS")
|
|
print("=" * 60)
|
|
|
|
for i, (path, content, distance) in enumerate(results, 1):
|
|
# Clean up unicode for console
|
|
path_clean = path.encode('ascii', 'ignore').decode() if path else "Unknown"
|
|
content_clean = content.encode('ascii', 'ignore').decode() if content else ""
|
|
print(f"\n{i}. [{path_clean}]")
|
|
print(f" Distance: {distance:.4f}")
|
|
print(f" Content: {content_clean[:150]}...")
|
|
|
|
print("\n" + "=" * 60)
|
|
print(f"Found {len(results)} result(s)")
|
|
|
|
|
|
def interactive():
|
|
"""Interactive search mode."""
|
|
print("Memory Vector Search")
|
|
print("Type 'quit' to exit\n")
|
|
|
|
while True:
|
|
query = input("Search: ").strip()
|
|
if query.lower() in ('quit', 'exit', 'q'):
|
|
break
|
|
if not query:
|
|
continue
|
|
|
|
try:
|
|
search(query)
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|
|
print()
|
|
|
|
|
|
def main():
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description='Search memories with vector similarity')
|
|
parser.add_argument('query', nargs='?', help='Search query')
|
|
parser.add_argument('-n', '--num', type=int, default=5, help='Number of results (default: 5)')
|
|
parser.add_argument('-i', '--interactive', action='store_true', help='Interactive mode')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.interactive or not args.query:
|
|
interactive()
|
|
else:
|
|
search(args.query, k=args.num)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|