python-sdk/pkg/hanzo-tools-vector/hanzo_tools/vector/mock_infinity.py

"""Mock implementation of infinity_embedded for testing on unsupported platforms."""

import random
from typing import Any, Dict, List
from pathlib import Path


class MockTable:
    """Mock implementation of an Infinity table."""

    def __init__(self, name: str, schema: Dict[str, Any]):
        self.name = name
        self.schema = schema
        self.data = []
        self._id_counter = 0

    def insert(self, records: List[Dict[str, Any]]):
        """Insert records into the table."""
        for record in records:
            # Add an internal ID if not present
            if "id" not in record:
                record["_internal_id"] = self._id_counter
                self._id_counter += 1
            self.data.append(record)

    def delete(self, condition: str):
        """Delete records matching condition."""
        # Simple implementation - just clear for now
        self.data = [r for r in self.data if not self._eval_condition(r, condition)]

    def output(self, columns: List[str]):
        """Start a query chain."""
        return MockQuery(self, columns)

    def _eval_condition(self, record: Dict[str, Any], condition: str) -> bool:
        """Evaluate a simple condition."""
        # Very basic implementation
        if "=" in condition:
            field, value = condition.split("=", 1)
            field = field.strip()
            value = value.strip().strip("'\"")
            return str(record.get(field, "")) == value
        return False


class MockQuery:
    """Mock query builder."""

    def __init__(self, table: MockTable, columns: List[str]):
        self.table = table
        self.columns = columns
        self.filters = []
        self.vector_search = None
        self.limit_value = None

    def filter(self, condition: str):
        """Add a filter condition."""
        self.filters.append(condition)
        return self

    def match_dense(self, column: str, vector: List[float], dtype: str, metric: str, limit: int):
        """Add vector search."""
        self.vector_search = {
            "column": column,
            "vector": vector,
            "dtype": dtype,
            "metric": metric,
            "limit": limit,
        }
        self.limit_value = limit
        return self

    def to_pl(self):
        """Execute query and return polars-like result."""
        results = self.table.data.copy()

        # Apply filters
        for condition in self.filters:
            results = [r for r in results if self.table._eval_condition(r, condition)]

        # Apply vector search (mock similarity)
        if self.vector_search:
            # Add mock scores
            for r in results:
                r["score"] = random.uniform(0.5, 1.0)
            # Sort by score
            results.sort(key=lambda x: x.get("score", 0), reverse=True)
            # Limit results
            if self.limit_value:
                results = results[: self.limit_value]

        # Return mock polars DataFrame
        return MockDataFrame(results)


class MockDataFrame:
    """Mock polars DataFrame."""

    def __init__(self, data: List[Dict[str, Any]]):
        self.data = data

    def __len__(self):
        return len(self.data)

    def iter_rows(self, named: bool = False):
        """Iterate over rows."""
        if named:
            return iter(self.data)
        else:
            # Return tuples
            if not self.data:
                return iter([])
            keys = list(self.data[0].keys())
            return iter([tuple(row.get(k) for k in keys) for row in self.data])


class MockDatabase:
    """Mock implementation of an Infinity database."""

    def __init__(self, name: str):
        self.name = name
        self.tables = {}

    def create_table(self, name: str, schema: Dict[str, Any]) -> MockTable:
        """Create a new table."""
        table = MockTable(name, schema)
        self.tables[name] = table
        return table

    def get_table(self, name: str) -> MockTable:
        """Get an existing table."""
        if name not in self.tables:
            raise KeyError(f"Table {name} not found")
        return self.tables[name]


class MockInfinity:
    """Mock implementation of Infinity connection."""

    def __init__(self, path: str):
        self.path = Path(path)
        self.databases = {}
        # Ensure directory exists
        self.path.mkdir(parents=True, exist_ok=True)

    def get_database(self, name: str) -> MockDatabase:
        """Get or create a database."""
        if name not in self.databases:
            self.databases[name] = MockDatabase(name)
        return self.databases[name]

    def disconnect(self):
        """Disconnect from Infinity."""
        pass


def connect(path: str) -> MockInfinity:
    """Connect to Infinity (mock implementation)."""
    return MockInfinity(path)