# Simple example of generating an embedding
from sentence_transformers import SentenceTransformer

model = SentenceTransformer('mixedbread-ai/mxbai-embed-large-v1')
text = "Hello world"
embedding = model.encode(text)  # This creates a numerical representation

print(f"Embedding shape: {embedding.shape}")  # Typically (1024,) or similar

# Using SentenceTransformer to compare two sentences
from sentence_transformers import SentenceTransformer, util

model = SentenceTransformer('mixedbread-ai/mxbai-embed-large-v1')

# Generate embeddings for two sentences
embedding1 = model.encode("I love programming")
embedding2 = model.encode("I enjoy coding")

# Calculate similarity (cosine similarity ranges from -1 to 1)
similarity = util.cos_sim(embedding1, embedding2)
print(f"Similarity: {similarity.item():.4f}")  # Should be high (close to 1)

# Basic example of SQLite-Vec setup
import sqlite3
import sqlite_vec

# Connect to SQLite and load the extension
conn = sqlite3.connect("example.db")
conn.enable_load_extension(True)
sqlite_vec.load(conn)
conn.enable_load_extension(False)

# Create a vector table
conn.execute("""
    CREATE VIRTUAL TABLE IF NOT EXISTS vectors 
    USING vec0(
        id INTEGER PRIMARY KEY,
        embedding bit[1024]  -- Binary vectors with 1024 dimensions
    )
""")

from sentence_transformers.quantization import quantize_embeddings

# Generate a standard embedding
embedding = model.encode("Binary quantization makes embeddings tiny!", normalize_embeddings=True)

# Quantize to binary (1 bit per dimension)
binary_embedding = quantize_embeddings(embedding, precision="ubinary")

print(f"Original size: {embedding.nbytes} bytes")
print(f"Binary size: {binary_embedding.nbytes} bytes")
print(f"Compression ratio: {embedding.nbytes / binary_embedding.nbytes}x")

# Comparing binary embeddings with Hamming distance
import numpy as np

# Create two binary vectors (already quantized)
binary_vec1 = np.array([0, 1, 1, 0, 1], dtype=np.bool_)
binary_vec2 = np.array([1, 1, 0, 0, 1], dtype=np.bool_)

# Hamming distance calculation (XOR + bit count)
hamming_distance = np.sum(binary_vec1 ^ binary_vec2)
print(f"Hamming distance: {hamming_distance}")  # Number of differing bits

pip install sentence-transformers
pip install sqlite-vec

# Import Needed libraries
from sentence_transformers import SentenceTransformer
import sqlite3
import sqlite_vec

# Load model
model = SentenceTransformer('mixedbread-ai/mxbai-embed-large-v1')

# Test SQLite-Vec
conn = sqlite3.connect(":memory:")
conn.enable_load_extension(True)
sqlite_vec.load(conn)
conn.enable_load_extension(False)

print("Setup complete! Everything is working correctly.")

from sentence_transformers import SentenceTransformer
from sentence_transformers.quantization import quantize_embeddings

# Load the model
model = SentenceTransformer('mixedbread-ai/mxbai-embed-large-v1')

# Generate embeddings for some text
text = "This is an example sentence."
embedding = model.encode(text, normalize_embeddings=True)

# Optionally quantize to binary
binary_embedding = quantize_embeddings(embedding, precision="ubinary")

import os
import sqlite3
import sqlite_vec
from sentence_transformers import SentenceTransformer
from typing import List, Tuple
import json

class RAG:
    """Semantic text search for Retrieval Augmented Generation"""
    
    def __init__(self, db_path: str = None):
        """Initialize the RAG system.
        
        Args:
            db_path: Path to the SQLite database file
        """
        self.db_path = os.path.join(os.path.dirname(__file__), "rag.db") if db_path is None else db_path
        self._init_db()
        # Initialize the embedding model
        self.model = SentenceTransformer('mixedbread-ai/mxbai-embed-large-v1')
    
    def _init_db(self) -> None:
        """Initialize the SQLite database with required tables."""
        conn = sqlite3.connect(self.db_path)
        conn.enable_load_extension(True)
        sqlite_vec.load(conn)
        conn.enable_load_extension(False)
        
        # Create virtual table for vector storage with binary embeddings
        conn.execute("""
            CREATE VIRTUAL TABLE IF NOT EXISTS segments 
            USING vec0(
                id TEXT PRIMARY KEY,       -- Segment ID
                text TEXT,                 -- Segment text
                embedding bit[1024],       -- Binary quantized embedding vector
                metadata TEXT              -- Optional metadata in JSON format
            )
        """)
        
        conn.commit()
        conn.close()
    
    def _get_db_connection(self) -> sqlite3.Connection:
        """Get a database connection with sqlite-vec extension loaded."""
        conn = sqlite3.connect(self.db_path)
        conn.enable_load_extension(True)
        sqlite_vec.load(conn)
        conn.enable_load_extension(False)
        return conn
    
    def load_segment(self, id: str, text: str, metadata: str = "{}") -> None:
        """Load a segment into the database.
        
        Args:
            id: Unique identifier for the segment
            text: The text content
            metadata: Optional JSON metadata string
        """
        if not text or len(text.strip()) == 0:
            return

        # Generate embedding using SentenceTransformer
        embedding = self.model.encode(text, normalize_embeddings=True)
        # Convert embedding to JSON string for SQLite-Vec's vec_quantize_binary function
        embedding_json = json.dumps(embedding.tolist())
        
        # Store in database using vec_quantize_binary
        conn = self._get_db_connection()
        
        conn.execute("""
            INSERT OR REPLACE INTO segments (id, text, embedding, metadata)
            VALUES (?, ?, vec_quantize_binary(?), ?)
        """, (id, text, embedding_json, metadata))
        
        conn.commit()
        conn.close()
    
    def search_segments(self, text: str, limit: int = 10) -> List[Tuple[str, str, float, str]]:
        """Search for segments similar to the given text.
        
        Args:
            text: The text to search for
            limit: Maximum number of results to return
            
        Returns:
            List of tuples containing (id, text, similarity_score, metadata)
        """
        # Generate embedding for query text
        query_embedding = self.model.encode(text, normalize_embeddings=True)
        # Convert query embedding to JSON string
        query_embedding_json = str(query_embedding.tolist())
        
        # Search database using vec_quantize_binary for the query
        conn = self._get_db_connection()
        
        cursor = conn.execute("""
            SELECT id, text, distance, metadata 
            FROM segments 
            WHERE embedding MATCH vec_quantize_binary(?) AND k = ?
            ORDER BY distance
        """, (query_embedding_json, limit))
        
        results = []
        for id, text, distance, metadata in cursor.fetchall():
            # Convert distance to similarity score (normalize to [0,1] range)
            # The distance is a Hamming distance between binary vectors
            similarity = 1.0 - (distance / 1024.0)  # 1024 is our vector dimension
            results.append((id, text, similarity, metadata))
        
        conn.close()
        
        return results
    
    def remove_segment(self, id: str) -> None:
        """Remove a segment from the database.
        
        Args:
            id: The ID of the segment to remove
        """
        conn = self._get_db_connection()
        conn.execute("DELETE FROM segments WHERE id = ?", (id,))
        conn.commit()
        conn.close()

# Initialize the RAG system
rag = RAG()

# Add some documents
rag.load_segment("doc1", "Tokyo is the capital of Japan.")
rag.load_segment("doc2", "Paris is the capital of France.")
rag.load_segment("doc3", "Berlin is the capital of Germany.")
rag.load_segment("doc4", "Rome is the capital of Italy.")

# Search for similar documents
results = rag.search_segments("What is the capital city of Japan?")

# Print results
for id, text, similarity, metadata in results:
    print(f"{id}: '{text}' (similarity: {similarity:.4f})")