search_tana.py

import os
import lancedb
from fastembed import TextEmbedding
from dotenv import load_dotenv
import sys

load_dotenv()

# Get current directory for portability
BASE_DIR = "/Users/krshirkoohi/Documents/AI Workspace/projects/MCP Servers/tana-embeddings"
LANCE_DB_PATH = os.path.join(BASE_DIR, "vector_store")
TABLE_NAME = "tana_nodes"

# Setup Local BGE Model (BAAI/bge-small-en-v1.5)
model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")

def search(query, limit=5):
    db = lancedb.connect(LANCE_DB_PATH)
    if TABLE_NAME not in db.table_names():
        print("Vector store not initialized yet.")
        return
    
    table = db.open_table(TABLE_NAME)
    
    # Generate query embedding locally
    query_vector = list(model.embed([query]))[0]
    
    results = table.search(query_vector).limit(limit).to_pandas()
    
    print(f"\nSemantic search results (LOCAL BGE) for: '{query}'\n")
    for _, row in results.iterrows():
        print(f"[{row['id']}] {row['name']}")
        if row['description']:
            desc = row['description'][:100] + "..." if len(row['description']) > 100 else row['description']
            print(f"   Description: {desc}")
        print(f"   Last Updated: {row['last_updated']}")
        print("-" * 40)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python3 search_tana.py 'your query'")
    else:
        search(" ".join(sys.argv[1:]))