kodit 0.1.10__tar.gz → 0.1.11__tar.gz

{kodit-0.1.10 → kodit-0.1.11}/.gitignore

@@ -174,4 +174,6 @@ cython_debug/
 .pypirc
 .kodit/
 .DS_Store
-.kodit.db
+.kodit.db
+benchmark.db
+profile.prof

{kodit-0.1.10 → kodit-0.1.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kodit
-Version: 0.1.10
+Version: 0.1.11
 Summary: Code indexing for better AI code generation
 Project-URL: Homepage, https://docs.helixml.tech/kodit/
 Project-URL: Documentation, https://docs.helixml.tech/kodit/

{kodit-0.1.10 → kodit-0.1.11}/pyproject.toml

@@ -56,6 +56,7 @@ dev = [
     "pytest>=8.3.5",
     "pytest-cov>=6.1.1",
     "ruff>=0.11.8",
+    "snakeviz>=2.2.2",
 ]
 
 [project.urls]

{kodit-0.1.10 → kodit-0.1.11}/src/kodit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.10'
-__version_tuple__ = version_tuple = (0, 1, 10)
+__version__ = version = '0.1.11'
+__version_tuple__ = version_tuple = (0, 1, 11)

{kodit-0.1.10 → kodit-0.1.11}/src/kodit/retreival/repository.py

@@ -5,21 +5,14 @@ related to searching and retrieving code snippets, including string-based search
 and their associated file information.
 """
 
-import math
-from typing import Any, TypeVar
+from typing import TypeVar
 
+import numpy as np
 import pydantic
 from sqlalchemy import (
-    ColumnElement,
-    Float,
-    cast,
-    desc,
-    func,
-    literal,
     select,
 )
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import Mapped
 
 from kodit.embedding.models import Embedding, EmbeddingType
 from kodit.indexing.models import Snippet
@@ -129,55 +122,110 @@ class RetrievalRepository:
         # Return results in the same order as input IDs
         return [id_to_result[i] for i in ids]
 
-    async def list_semantic_results(
-        self, embedding_type: EmbeddingType, embedding: list[float], top_k: int = 10
-    ) -> list[tuple[int, float]]:
-        """List semantic results."""
-        cosine_similarity = cosine_similarity_json(Embedding.embedding, embedding)
+    async def fetch_embeddings(
+        self, embedding_type: EmbeddingType
+    ) -> list[tuple[int, list[float]]]:
+        """Fetch all embeddings of a given type from the database.
 
-        query = (
-            select(Embedding, cosine_similarity)
-            .where(Embedding.type == embedding_type)
-            .order_by(desc(cosine_similarity))
-            .limit(top_k)
+        Args:
+            embedding_type: The type of embeddings to fetch
+
+        Returns:
+            List of (snippet_id, embedding) tuples
+
+        """
+        # Only select the fields we need and use a more efficient query
+        query = select(Embedding.snippet_id, Embedding.embedding).where(
+            Embedding.type == embedding_type
         )
         rows = await self.session.execute(query)
-        return [(embedding.snippet_id, distance) for embedding, distance in rows.all()]
+        return [tuple(row) for row in rows.all()]  # Convert Row objects to tuples
+
+    def prepare_vectors(
+        self, embeddings: list[tuple[int, list[float]]], query_embedding: list[float]
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Convert embeddings to numpy arrays.
 
+        Args:
+            embeddings: List of (snippet_id, embedding) tuples
+            query_embedding: Query embedding vector
 
-def cosine_similarity_json(
-    col: Mapped[Any], query_vec: list[float]
-) -> ColumnElement[Any]:
-    """Calculate the cosine similarity using pure sqlalchemy.
+        Returns:
+            Tuple of (stored_vectors, query_vector) as numpy arrays
 
-    Works for a *fixed-length* vector stored as a JSON array in SQLite.
-    The calculation is done entirely in SQL using SQLite's JSON functions.
+        """
+        stored_vecs = np.array(
+            [emb[1] for emb in embeddings]
+        )  # Use index 1 to get embedding
+        query_vec = np.array(query_embedding)
+        return stored_vecs, query_vec
 
-    Args:
-        col: The column containing the JSON array of floats
-        query_vec: The query vector to compare against
+    def compute_similarities(
+        self, stored_vecs: np.ndarray, query_vec: np.ndarray
+    ) -> np.ndarray:
+        """Compute cosine similarities between stored vectors and query vector.
 
-    Returns:
-        A SQLAlchemy expression that computes the cosine similarity
+        Args:
+            stored_vecs: Array of stored embedding vectors
+            query_vec: Query embedding vector
 
-    """
-    # Pre-compute query norm
-    q_norm = math.sqrt(sum(x * x for x in query_vec))
-
-    # Calculate dot product using JSON array functions
-    dot = sum(
-        cast(func.json_extract(col, f"$[{i}]"), Float) * literal(float(q))
-        for i, q in enumerate(query_vec)
-    )
-
-    # Calculate row norm on the fly
-    row_norm = func.sqrt(
-        sum(
-            cast(func.json_extract(col, f"$[{i}]"), Float)
-            * cast(func.json_extract(col, f"$[{i}]"), Float)
-            for i in range(len(query_vec))
-        )
-    )
+        Returns:
+            Array of similarity scores
+
+        """
+        stored_norms = np.linalg.norm(stored_vecs, axis=1)
+        query_norm = np.linalg.norm(query_vec)
+        return np.dot(stored_vecs, query_vec) / (stored_norms * query_norm)
+
+    def get_top_k_results(
+        self,
+        similarities: np.ndarray,
+        embeddings: list[tuple[int, list[float]]],
+        top_k: int,
+    ) -> list[tuple[int, float]]:
+        """Get top-k results by similarity score.
+
+        Args:
+            similarities: Array of similarity scores
+            embeddings: List of (snippet_id, embedding) tuples
+            top_k: Number of results to return
+
+        Returns:
+            List of (snippet_id, similarity_score) tuples
+
+        """
+        top_indices = np.argsort(similarities)[::-1][:top_k]
+        return [
+            (embeddings[i][0], float(similarities[i])) for i in top_indices
+        ]  # Use index 0 to get snippet_id
+
+    async def list_semantic_results(
+        self, embedding_type: EmbeddingType, embedding: list[float], top_k: int = 10
+    ) -> list[tuple[int, float]]:
+        """List semantic results using cosine similarity.
+
+        This implementation fetches all embeddings of the given type and computes
+        cosine similarity in Python using NumPy for better performance.
+
+        Args:
+            embedding_type: The type of embeddings to search
+            embedding: The query embedding vector
+            top_k: Number of results to return
+
+        Returns:
+            List of (snippet_id, similarity_score) tuples, sorted by similarity
+
+        """
+        # Step 1: Fetch embeddings from database
+        embeddings = await self.fetch_embeddings(embedding_type)
+        if not embeddings:
+            return []
+
+        # Step 2: Convert to numpy arrays
+        stored_vecs, query_vec = self.prepare_vectors(embeddings, embedding)
+
+        # Step 3: Compute similarities
+        similarities = self.compute_similarities(stored_vecs, query_vec)
 
-    # Calculate cosine similarity
-    return (dot / (row_norm * literal(q_norm))).label("cosine_similarity")
+        # Step 4: Get top-k results
+        return self.get_top_k_results(similarities, embeddings, top_k)

kodit-0.1.11/tests/performance/similarity.py

@@ -0,0 +1,139 @@
+"""Benchmark script for semantic similarity search performance."""
+
+import asyncio
+from pathlib import Path
+import random
+import time
+from typing import List
+
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
+
+from kodit.embedding.models import Embedding, EmbeddingType
+from kodit.indexing.models import Index, Snippet
+from kodit.retreival.repository import RetrievalRepository
+from kodit.sources.models import File, Source
+
+
+def generate_random_embedding(dim: int = 750) -> List[float]:
+    """Generate a random embedding vector of specified dimension."""
+    return [random.uniform(-1, 1) for _ in range(dim)]
+
+
+async def setup_test_data(session: AsyncSession, num_embeddings: int = 5000) -> None:
+    """Set up test data with random embeddings."""
+    # Create a test index
+    source = Source(uri="test", cloned_path="test")
+    session.add(source)
+    await session.commit()
+    index = Index(source_id=source.id)
+    session.add(index)
+    await session.commit()
+    file = File(
+        uri="test",
+        cloned_path="test",
+        source_id=source.id,
+    )
+    session.add(file)
+    await session.commit()
+    snippet = Snippet(
+        file_id=file.id,
+        index_id=index.id,
+        content="This is a test snippet",
+    )
+    session.add(snippet)
+    await session.commit()
+
+    # Create test embeddings
+    embeddings = []
+    for i in range(num_embeddings):
+        embedding = Embedding(
+            snippet_id=snippet.id,
+            type=EmbeddingType.CODE,
+            embedding=generate_random_embedding(),
+        )
+        embeddings.append(embedding)
+
+    session.add_all(embeddings)
+    await session.commit()
+
+
+async def run_benchmark(session: AsyncSession) -> None:
+    """Run the semantic search benchmark."""
+    # Setup test data
+    print("Setting up test data...")
+    await setup_test_data(session)
+
+    # Create repository instance
+    repo = RetrievalRepository(session)
+
+    # Generate a test query embedding
+    query_embedding = generate_random_embedding()
+
+    # Run the benchmark
+    num_runs = 10
+    total_time = 0
+    results = []  # Initialize results list
+
+    print("Running warm-up query...")
+    # Warm up
+    await repo.list_semantic_results(
+        embedding_type=EmbeddingType.CODE, embedding=query_embedding, top_k=10
+    )
+
+    print(f"\nRunning {num_runs} benchmark queries...")
+
+    # Actual benchmark
+    for i in range(num_runs):
+        start_time = time.perf_counter()
+        results = await repo.list_semantic_results(
+            embedding_type=EmbeddingType.CODE, embedding=query_embedding, top_k=10
+        )
+        end_time = time.perf_counter()
+        run_time = end_time - start_time
+        total_time += run_time
+        print(f"\nRun {i + 1}/{num_runs}: {run_time * 1000:.2f}ms")
+
+    # Calculate average time per run
+    avg_time = total_time / num_runs
+
+    print(f"\nSemantic Search Performance Results:")
+    print(f"Number of runs: {num_runs}")
+    print(f"Total execution time: {total_time:.2f} seconds")
+    print(f"Average time per query: {avg_time * 1000:.2f} ms")
+
+    # Print sample results
+    print(f"\nSample query returned {len(results)} results")
+    if results:  # Add safety check
+        print(f"First result score: {results[0][1]:.4f}")
+
+
+async def main():
+    """Main entry point for the benchmark."""
+    # Remove the database file if it exists
+    if Path("benchmark.db").exists():
+        Path("benchmark.db").unlink()
+
+    # Create async engine and session
+    engine = create_async_engine("sqlite+aiosqlite:///benchmark.db")
+
+    # Create tables
+    async with engine.begin() as conn:
+        await conn.run_sync(Source.metadata.create_all)
+        await conn.run_sync(File.metadata.create_all)
+        await conn.run_sync(Index.metadata.create_all)
+        await conn.run_sync(Snippet.metadata.create_all)
+        await conn.run_sync(Embedding.metadata.create_all)
+
+    # Create session factory
+    async_session = async_sessionmaker(engine, expire_on_commit=False)
+
+    # Run benchmark
+    async with async_session() as session:
+        await run_benchmark(session)
+
+    # Cleanup
+    await engine.dispose()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

{kodit-0.1.10 → kodit-0.1.11}/uv.lock

@@ -803,6 +803,7 @@ dev = [
     { name = "pytest-asyncio" },
     { name = "pytest-cov" },
     { name = "ruff" },
+    { name = "snakeviz" },
 ]
 
 [package.metadata]
@@ -842,6 +843,7 @@ dev = [
     { name = "pytest-asyncio", specifier = ">=0.26.0" },
     { name = "pytest-cov", specifier = ">=6.1.1" },
     { name = "ruff", specifier = ">=0.11.8" },
+    { name = "snakeviz", specifier = ">=2.2.2" },
 ]
 
 [[package]]
@@ -1939,6 +1941,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/04/be/d09147ad1ec7934636ad912901c5fd7667e1c858e19d355237db0d0cd5e4/smmap-5.0.2-py3-none-any.whl", hash = "sha256:b30115f0def7d7531d22a0fb6502488d879e75b260a9db4d0819cfb25403af5e", size = 24303, upload-time = "2025-01-02T07:14:38.724Z" },
 ]
 
+[[package]]
+name = "snakeviz"
+version = "2.2.2"
+source = { registry = "https://pypi.org/simple/" }
+dependencies = [
+    { name = "tornado" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/04/06/82f56563b16d33c2586ac2615a3034a83a4ff1969b84c8d79339e5d07d73/snakeviz-2.2.2.tar.gz", hash = "sha256:08028c6f8e34a032ff14757a38424770abb8662fb2818985aeea0d9bc13a7d83", size = 182039, upload-time = "2024-11-09T22:03:58.99Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/cd/f7/83b00cdf4f114f10750a18b64c27dc34636d0ac990ccac98282f5c0fbb43/snakeviz-2.2.2-py3-none-any.whl", hash = "sha256:77e7b9c82f6152edc330040319b97612351cd9b48c706434c535c2df31d10ac5", size = 183477, upload-time = "2024-11-09T22:03:57.049Z" },
+]
+
 [[package]]
 name = "sniffio"
 version = "1.3.1"