1# --- Configuration ---
2# Customize these values for your migration
3PINECONE_INDEX_NAME = "my-pinecone-index" # Your Pinecone index name
4TABLE_NAME = PINECONE_INDEX_NAME          # Antfly table name (using same name)
5INDEX_NAME = "nomic_index"                # Name for the vector index in Antfly
6TEMPLATE = "{{text}}"                     # Handlebars template - which field(s) to embed
7QUERY_TEXT = "example search"             # Test query for verification
8
9ANTFLY_BASE_URL = os.getenv("ANTFLY_BASE_URL", "http://localhost:8080/api/v1")
10
11# Embedder configuration (optional)
12# Set these if you have a local embedder matching your Pinecone vectors' dimension
13# Common options: ("ollama", "nomic-embed-text") for 768d, ("openai", "text-embedding-3-small") for 1536d
14EMBEDDER_PROVIDER = "ollama"              # "ollama", "termite", or "openai"
15EMBEDDER_MODEL = "nomic-embed-text"       # Model name

1def main():
2    load_dotenv()
3
4    # --- Initialize Pinecone ---
5    pinecone_api_key = os.getenv("PINECONE_API_KEY")
6    if not pinecone_api_key:
7        print("FATAL: PINECONE_API_KEY not found in environment.", file=sys.stderr)
8        sys.exit(1)
9
10    pc = Pinecone(api_key=pinecone_api_key)
11    pinecone_index = pc.Index(PINECONE_INDEX_NAME)
12    print(f"Connected to Pinecone index '{PINECONE_INDEX_NAME}'.")
13
14    # --- Initialize Antfly ---
15    client = AntflyClient(base_url=ANTFLY_BASE_URL)
16    try:
17        print(f"Connecting to Antfly at {ANTFLY_BASE_URL}...")
18        client.list_tables()
19        print("Connected to Antfly.")
20    except Exception as e:
21        print(f"Could not connect to Antfly: {e}")
22        print("Ensure the Antfly server is running.")
23        sys.exit(1)

1    # --- Create Antfly table ---
2    # Tables are schemaless - no schema definition required
3    try:
4        client.drop_table(TABLE_NAME)
5        print(f"Deleted existing table '{TABLE_NAME}'")
6    except AntflyException:
7        pass
8
9    try:
10        client.create_table(name=TABLE_NAME)
11        print(f"Created table '{TABLE_NAME}'")
12    except Exception as e:
13        print(f"Error creating table: {e}")
14        sys.exit(1)
15
16    print("Waiting for table shards to initialize...")
17    time.sleep(5)

1    # --- Create vector index BEFORE inserting data ---
2    # This tells Antfly where to store pre-computed embeddings
3    try:
4        print(f"\nCreating vector index '{INDEX_NAME}' with dimension {dimension}...")
5        index_config = {
6            "name": INDEX_NAME,
7            "type": "aknn_v0",
8            "dimension": dimension,
9            "template": TEMPLATE,
10            "embedder": {"provider": EMBEDDER_PROVIDER, "model": EMBEDDER_MODEL}
11        }
12        print(f"  -> Using embedder: {EMBEDDER_PROVIDER}/{EMBEDDER_MODEL}")
13
14        resp = httpx.post(
15            f"{ANTFLY_BASE_URL}/tables/{TABLE_NAME}/indexes/{INDEX_NAME}",
16            json=index_config,
17            timeout=30.0
18        )
19        if resp.status_code == 200:
20            print("Index created successfully.")
21        else:
22            print(f"Index creation response: {resp.status_code} - {resp.text}")
23    except Exception as e:
24        print(f"Error creating index: {e}")
25

1    def fetch_all_vectors(pc_index, namespace: str):
2        """
3        Fetch all vectors from a Pinecone namespace.
4
5        Uses a zero-vector query with high top_k - a common pattern
6        since Pinecone doesn't have a direct "list all" API.
7        """
8        print(f"Fetching vectors from namespace: '{namespace}'...")
9
10        stats = pc_index.describe_index_stats()
11        dim = stats['dimension']
12
13        # Query with zero vector to get all vectors
14        query_response = pc_index.query(
15            namespace=namespace,
16            vector=[0.0] * dim,
17            top_k=1000,
18            include_values=True,
19            include_metadata=True
20        )
21
22        vectors = query_response.get('matches', [])
23        if len(vectors) == 1000:
24            print("  WARNING: Hit 1000 limit - may have more vectors.", file=sys.stderr)
25

1    # --- Upsert into Antfly ---
2    for namespace, items in all_vectors.items():
3        try:
4            print(f"Upserting {len(items)} vectors for namespace '{namespace}'...")
5
6            inserts = {}
7            for item in items:
8                key = item['id']
9                properties = {k: v for k, v in item.items() if k != 'id'}
10
11                # Use _embeddings format for pre-computed embeddings
12                # Format: {"_embeddings": {"<index_name>": [vector values]}}
13                if 'values' in properties:
14                    embedding_vector = properties.pop('values')
15                    properties['_embeddings'] = {INDEX_NAME: embedding_vector}
16
17                inserts[key] = properties
18
19            resp = httpx.post(
20                f"{ANTFLY_BASE_URL}/tables/{TABLE_NAME}/batch",
21                json={"inserts": inserts},
22                timeout=60.0
23            )
24            if resp.status_code in (200, 201):
25                print(f"  -> Successfully upserted {len(items)} vectors.")

1    # --- Verify migration ---
2    print("\n=== Verifying Migration ===")
3
4    try:
5        resp = httpx.post(
6            f"{ANTFLY_BASE_URL}/tables/{TABLE_NAME}/query",
7            json={"limit": 200},
8            timeout=30.0
9        )
10        if resp.status_code == 200:
11            data = resp.json()
12            hits = data.get('responses', [{}])[0].get('hits', {}).get('hits', [])
13            print(f"Total records in Antfly: {len(hits)}")
14
15            if hits:
16                print("\nSample records:")
17                for hit in hits[:3]:
18                    key = hit.get('_id', 'unknown')
19                    source = hit.get('_source', {})
20                    text = source.get('text', '')[:80] if source else ''
21                    print(f"  - {key}: {text}...")
22    except Exception as e:
23        print(f"Verification failed: {e}")
24
25    # --- Test semantic search ---