89 lines
2.8 KiB
Python
Executable File
89 lines
2.8 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Auto-Memory Recall — query the auto-memory ChromaDB collection.
|
|
Usage:
|
|
python3 auto-memory-recall.py "how did we fix nginx"
|
|
python3 auto-memory-recall.py --query "chromadb setup" --limit 3 --topic infrastructure
|
|
"""
|
|
|
|
import sys, json, os, argparse
|
|
import requests
|
|
import chromadb
|
|
|
|
CHROMADB_HOST = os.environ.get("CHROMADB_HOST", "192.168.86.25")
|
|
CHROMADB_PORT = int(os.environ.get("CHROMADB_PORT", "8000"))
|
|
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://192.168.86.40:11434")
|
|
EMBED_MODEL = "nomic-embed-text"
|
|
COLLECTION = "auto-memory"
|
|
|
|
|
|
def ollama_embed(text: str) -> list[float]:
|
|
r = requests.post(
|
|
f"{OLLAMA_URL}/api/embed",
|
|
json={"model": EMBED_MODEL, "input": text},
|
|
timeout=10,
|
|
)
|
|
r.raise_for_status()
|
|
return r.json()["embeddings"][0]
|
|
|
|
|
|
def recall(query: str, limit: int = 5, topic: str = None, agent_id: str = None) -> list[dict]:
|
|
client = chromadb.HttpClient(host=CHROMADB_HOST, port=CHROMADB_PORT)
|
|
collection = client.get_or_create_collection(name=COLLECTION, metadata={"hnsw:space": "cosine"})
|
|
|
|
where = {}
|
|
if topic:
|
|
where["topic"] = topic
|
|
if agent_id:
|
|
where["agent_id"] = agent_id
|
|
|
|
try:
|
|
embedding = ollama_embed(query)
|
|
results = collection.query(
|
|
query_embeddings=[embedding],
|
|
n_results=limit,
|
|
where=where if where else None,
|
|
)
|
|
except Exception:
|
|
# Fallback to text query if embedding fails
|
|
results = collection.query(
|
|
query_texts=[query],
|
|
n_results=limit,
|
|
where=where if where else None,
|
|
)
|
|
|
|
memories = []
|
|
if results and results["documents"]:
|
|
for i, doc in enumerate(results["documents"][0]):
|
|
meta = results["metadatas"][0][i] if results["metadatas"] else {}
|
|
dist = results["distances"][0][i] if results["distances"] else None
|
|
memories.append({
|
|
"memory": doc,
|
|
"similarity": round(1 - dist, 3) if dist is not None else None,
|
|
"id": results["ids"][0][i] if results["ids"] else None,
|
|
**meta,
|
|
})
|
|
return memories
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Recall memories from auto-memory collection")
|
|
parser.add_argument("query", nargs="?", help="Search query")
|
|
parser.add_argument("--query", "-q", dest="query_flag")
|
|
parser.add_argument("--limit", "-n", type=int, default=5)
|
|
parser.add_argument("--topic", "-t")
|
|
parser.add_argument("--agent-id", "-a")
|
|
args = parser.parse_args()
|
|
|
|
query = args.query or args.query_flag
|
|
if not query:
|
|
parser.print_help()
|
|
sys.exit(1)
|
|
|
|
results = recall(query, args.limit, args.topic, args.agent_id)
|
|
print(json.dumps(results, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|