noesium 0.1.0__py3-none-any.whl

noesium/core/utils/logging.py

@@ -0,0 +1,172 @@
+"""
+Simplified logging configuration for noesium with color support.
+
+Usage:
+    from noesium.core.utils.logging import setup_logging, get_logger
+
+    # Basic usage - automatically uses environment variables if available
+    setup_logging()
+    logger = get_logger(__name__)
+
+    # Explicit configuration
+    setup_logging(level="DEBUG", enable_colors=True)
+"""
+
+import logging
+import os
+import sys
+from pathlib import Path
+from typing import Dict, List, Optional
+
+try:
+    import colorlog
+except ImportError:
+    colorlog = None  # fallback if colorlog is not installed
+
+# Configure logging
+log_level = os.getenv("LOG_LEVEL", "INFO").upper()
+logging.basicConfig(level=getattr(logging, log_level, logging.INFO))
+logger = logging.getLogger(__name__)
+
+
+def setup_logging(
+    level: Optional[str] = None,
+    log_file: Optional[str] = None,
+    log_file_level: Optional[str] = None,
+    enable_colors: Optional[bool] = None,
+    log_format: Optional[str] = None,
+    custom_colors: Optional[Dict[str, str]] = None,
+    third_party_level: Optional[str] = None,
+    clear_existing: bool = True,
+) -> None:
+    """Initialize logging for noesium and third-party libs.
+
+    Args:
+        level: Log level (INFO, DEBUG, etc.)
+        log_file: Path to log file
+        log_file_level: Log level for file handler
+        enable_colors: Whether to enable colored output
+        log_format: Custom log format string
+        custom_colors: Custom color mapping
+        third_party_level: Log level for third-party libraries
+        clear_existing: Whether to clear existing handlers
+    """
+    # Set defaults for any remaining None values
+    level = level or "INFO"
+    log_file_level = log_file_level or "DEBUG"
+    enable_colors = enable_colors if enable_colors is not None else True
+    third_party_level = third_party_level or "WARNING"
+
+    root_logger = logging.getLogger()
+    if clear_existing:
+        root_logger.handlers.clear()
+
+    log_level = getattr(logging, level.upper(), logging.INFO)
+    root_logger.setLevel(log_level)
+
+    fmt = log_format or "%(log_color)s%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+    datefmt = "%H:%M:%S"
+
+    if enable_colors and colorlog:
+        formatter = colorlog.ColoredFormatter(
+            fmt=fmt,
+            datefmt=datefmt,
+            log_colors=custom_colors
+            or {
+                "DEBUG": "cyan",
+                "INFO": "green",
+                "WARNING": "yellow",
+                "ERROR": "red",
+                "CRITICAL": "bold_red",
+            },
+            style="%",
+        )
+    else:
+        formatter = logging.Formatter(fmt.replace("%(log_color)s", ""), datefmt=datefmt)
+
+    # Console handler
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setFormatter(formatter)
+    console_handler.setLevel(log_level)
+    root_logger.addHandler(console_handler)
+
+    # Optional file handler
+    if log_file:
+        Path(log_file).parent.mkdir(parents=True, exist_ok=True)
+        file_handler = logging.FileHandler(log_file)
+        file_handler.setFormatter(formatter)
+        file_handler.setLevel(getattr(logging, log_file_level.upper(), logging.DEBUG))
+        root_logger.addHandler(file_handler)
+
+
+def get_logger(name: str) -> logging.Logger:
+    return logging.getLogger(name)
+
+
+# --- Per-message color utilities -------------------------------------------------
+
+_ANSI_COLORS: Dict[str, str] = {
+    "black": "\033[30m",
+    "red": "\033[31m",
+    "green": "\033[32m",
+    "yellow": "\033[33m",
+    "blue": "\033[34m",
+    "magenta": "\033[35m",
+    "cyan": "\033[36m",
+    "white": "\033[37m",
+}
+
+_ANSI_ATTRS: Dict[str, str] = {
+    "bold": "\033[1m",
+    "dim": "\033[2m",
+    "underline": "\033[4m",
+    "blink": "\033[5m",
+    "reverse": "\033[7m",
+}
+
+_ANSI_RESET = "\033[0m"
+
+
+def color_text(message: str, color: Optional[str] = None, attrs: Optional[List[str]] = None) -> str:
+    """Wrap message with ANSI color/attribute codes.
+
+    Works regardless of whether colorlog is installed. Use to color a specific
+    log invocation, e.g. `logger.info(color_text("done", "cyan"))`.
+    """
+    if not color and not attrs:
+        return message
+
+    parts: List[str] = []
+    if attrs:
+        for attr in attrs:
+            code = _ANSI_ATTRS.get(attr.lower())
+            if code:
+                parts.append(code)
+    if color:
+        parts.append(_ANSI_COLORS.get(color.lower(), ""))
+
+    return f"{''.join(parts)}{message}{_ANSI_RESET}"
+
+
+def info_color(logger_obj: logging.Logger, message: str, color: Optional[str] = None, *, bold: bool = False) -> None:
+    """Log INFO with optional per-message color and bold attribute."""
+    attrs = ["bold"] if bold else None
+    logger_obj.info(color_text(message, color=color, attrs=attrs))
+
+
+def debug_color(logger_obj: logging.Logger, message: str, color: Optional[str] = None, *, bold: bool = False) -> None:
+    """Log DEBUG with optional per-message color and bold attribute."""
+    attrs = ["bold"] if bold else None
+    logger_obj.debug(color_text(message, color=color, attrs=attrs))
+
+
+def warning_color(logger_obj: logging.Logger, message: str, color: Optional[str] = None, *, bold: bool = False) -> None:
+    """Log WARNING with optional per-message color and bold attribute."""
+    attrs = ["bold"] if bold else None
+    logger_obj.warning(color_text(message, color=color, attrs=attrs))
+
+
+def error_color(logger_obj: logging.Logger, message: str, color: Optional[str] = None, *, bold: bool = False) -> None:
+    """Log ERROR with optional per-message color and bold attribute."""
+    attrs = ["bold"] if bold else None
+    logger_obj.error(color_text(message, color=color, attrs=attrs))

noesium/core/utils/statistics.py

@@ -0,0 +1,12 @@
+import math
+from typing import List
+
+
+def cosine_similarity(a: List[float], b: List[float]) -> float:
+    """Calculate cosine similarity between two vectors"""
+    dot_product = sum(x * y for x, y in zip(a, b))
+    magnitude_a = math.sqrt(sum(x * x for x in a))
+    magnitude_b = math.sqrt(sum(x * x for x in b))
+    if magnitude_a == 0 or magnitude_b == 0:
+        return 0
+    return dot_product / (magnitude_a * magnitude_b)

noesium/core/utils/typing.py

@@ -0,0 +1,17 @@
+"""
+Compatibility module for typing features not available in older Python versions.
+"""
+
+import sys
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import override
+else:
+    if sys.version_info >= (3, 12):
+        from typing import override
+    else:
+        # Fallback for Python < 3.12
+        def override(func):
+            """Fallback override decorator for Python < 3.12."""
+            return func

noesium/core/vector_store/__init__.py

@@ -0,0 +1,79 @@
+import os
+
+from .base import BaseVectorStore, OutputData
+
+# Optional imports - vector store providers might not be available
+try:
+    from .weaviate import WeaviateVectorStore
+
+    WEAVIATE_AVAILABLE = True
+except ImportError:
+    WeaviateVectorStore = None
+    WEAVIATE_AVAILABLE = False
+
+try:
+    from .pgvector import PGVectorStore
+
+    PGVECTOR_AVAILABLE = True
+except ImportError:
+    PGVectorStore = None
+    PGVECTOR_AVAILABLE = False
+
+__all__ = [
+    "BaseVectorStore",
+    "OutputData",
+    "WeaviateVectorStore",
+    "PGVectorStore",
+    "get_vector_store",
+]
+
+#############################
+# Common Vector Store helper functions
+#############################
+
+
+def get_vector_store(
+    provider: str = os.getenv("COGENTS_VECTOR_STORE_PROVIDER", "weaviate"),
+    collection_name: str = "default_collection",
+    embedding_model_dims: int = int(os.getenv("COGENTS_EMBEDDING_DIMS", "768")),
+    **kwargs,
+):
+    """
+    Get a vector store instance based on the specified provider.
+
+    Args:
+        provider: Vector store provider to use ("weaviate", "pgvector")
+        collection_name: Name of the collection/table to use
+        embedding_model_dims: Dimensions of the embedding model
+        **kwargs: Additional provider-specific arguments:
+            - weaviate: cluster_url, auth_client_secret, additional_headers
+            - pgvector: dbname, user, password, host, port, diskann, hnsw
+
+    Returns:
+        BaseVectorStore instance for the specified provider
+
+    Raises:
+        ValueError: If provider is not supported or not available
+    """
+    if provider == "weaviate":
+        if not WEAVIATE_AVAILABLE:
+            raise ValueError(
+                "weaviate provider is not available. Please install the required dependencies: pip install weaviate-client"
+            )
+        return WeaviateVectorStore(
+            collection_name=collection_name,
+            embedding_model_dims=embedding_model_dims,
+            **kwargs,
+        )
+    elif provider == "pgvector":
+        if not PGVECTOR_AVAILABLE:
+            raise ValueError(
+                "pgvector provider is not available. Please install the required dependencies: pip install psycopg2"
+            )
+        return PGVectorStore(
+            collection_name=collection_name,
+            embedding_model_dims=embedding_model_dims,
+            **kwargs,
+        )
+    else:
+        raise ValueError(f"Unsupported provider: {provider}. Supported providers: weaviate, pgvector")

noesium/core/vector_store/base.py

@@ -0,0 +1,94 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel
+
+
+class OutputData(BaseModel):
+    """Standard output data structure for vector store operations."""
+
+    id: str
+    score: Optional[float] = None
+    payload: Dict[str, Any]
+
+
+class BaseVectorStore(ABC):
+    def __init__(self, embedding_model_dims: int):
+        """
+        Initialize the vector store.
+
+        Args:
+            embedding_model_dims: Expected dimensions for embedding vectors
+        """
+        self.embedding_model_dims = embedding_model_dims
+
+    def _validate_vector_dimensions(self, vectors: List[List[float]]) -> None:
+        """
+        Validate that all vectors have the expected dimensions.
+
+        Args:
+            vectors: List of vectors to validate
+
+        Raises:
+            ValueError: If any vector has incorrect dimensions
+        """
+        for i, vector in enumerate(vectors):
+            if len(vector) != self.embedding_model_dims:
+                raise ValueError(
+                    f"Vector at index {i} has {len(vector)} dimensions, "
+                    f"expected {self.embedding_model_dims}. "
+                    f"Check that your embedding model matches COGENTS_EMBEDDING_DIMS={self.embedding_model_dims}"
+                )
+
+    @abstractmethod
+    def create_collection(self, vector_size: int, distance: str = "cosine") -> None:
+        """Create a new collection."""
+
+    @abstractmethod
+    def insert(
+        self,
+        vectors: List[List[float]],
+        payloads: Optional[List[Dict[str, Any]]] = None,
+        ids: Optional[List[str]] = None,
+    ) -> None:
+        """Insert vectors into a collection."""
+
+    @abstractmethod
+    def search(
+        self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict[str, Any]] = None
+    ) -> List[OutputData]:
+        """Search for similar vectors."""
+
+    @abstractmethod
+    def delete(self, vector_id: str) -> None:
+        """Delete a vector by ID."""
+
+    @abstractmethod
+    def update(
+        self, vector_id: str, vector: Optional[List[float]] = None, payload: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Update a vector and its payload."""
+
+    @abstractmethod
+    def get(self, vector_id: str) -> Optional[OutputData]:
+        """Retrieve a vector by ID."""
+
+    @abstractmethod
+    def list_collections(self) -> List[str]:
+        """List all collections."""
+
+    @abstractmethod
+    def delete_collection(self) -> None:
+        """Delete a collection."""
+
+    @abstractmethod
+    def collection_info(self) -> Dict[str, Any]:
+        """Get information about a collection."""
+
+    @abstractmethod
+    def list(self, filters: Optional[Dict[str, Any]] = None, limit: Optional[int] = None) -> List[OutputData]:
+        """List all memories."""
+
+    @abstractmethod
+    def reset(self) -> None:
+        """Reset by delete the collection and recreate it."""

noesium/core/vector_store/pgvector.py

@@ -0,0 +1,304 @@
+import json
+import logging
+from typing import Any, Dict, List, Optional
+
+try:
+    import psycopg2
+    from psycopg2.extras import execute_values
+except ImportError:
+    raise ImportError("The 'psycopg2' library is required. Please install it using 'pip install psycopg2'.")
+
+from .base import BaseVectorStore, OutputData
+
+logger = logging.getLogger(__name__)
+
+
+class PGVectorStore(BaseVectorStore):
+    def __init__(
+        self,
+        dbname,
+        collection_name,
+        embedding_model_dims,
+        user,
+        password,
+        host,
+        port,
+        diskann,
+        hnsw,
+    ):
+        """
+        Initialize the PGVector database.
+
+        Args:
+            dbname (str): Database name
+            collection_name (str): Collection name
+            embedding_model_dims (int): Dimension of the embedding vector
+            user (str): Database user
+            password (str): Database password
+            host (str, optional): Database host
+            port (int, optional): Database port
+            diskann (bool, optional): Use DiskANN for faster search
+            hnsw (bool, optional): Use HNSW for faster search
+        """
+        super().__init__(embedding_model_dims)
+
+        self.collection_name = collection_name
+        self.use_diskann = diskann
+        self.use_hnsw = hnsw
+
+        self.conn = psycopg2.connect(dbname=dbname, user=user, password=password, host=host, port=port)
+        self.cur = self.conn.cursor()
+
+        collections = self.list_collections()
+        if collection_name not in collections:
+            self.create_collection(embedding_model_dims)
+
+    def create_collection(self, vector_size: int, distance: str = "cosine") -> None:
+        """
+        Create a new collection (table in PostgreSQL).
+        Will also initialize vector search index if specified.
+
+        Args:
+            vector_size (int): Dimension of the embedding vector.
+            distance (str): Distance metric (not used in PGVector, for compatibility).
+        """
+        self.cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
+        self.cur.execute(
+            f"""
+            CREATE TABLE IF NOT EXISTS {self.collection_name} (
+                id UUID PRIMARY KEY,
+                vector vector({vector_size}),
+                payload JSONB
+            );
+        """
+        )
+
+        if self.use_diskann and vector_size < 2000:
+            # Check if vectorscale extension is installed
+            self.cur.execute("SELECT * FROM pg_extension WHERE extname = 'vectorscale'")
+            if self.cur.fetchone():
+                # Create DiskANN index if extension is installed for faster search
+                self.cur.execute(
+                    f"""
+                    CREATE INDEX IF NOT EXISTS {self.collection_name}_diskann_idx
+                    ON {self.collection_name}
+                    USING diskann (vector);
+                """
+                )
+        elif self.use_hnsw:
+            self.cur.execute(
+                f"""
+                CREATE INDEX IF NOT EXISTS {self.collection_name}_hnsw_idx
+                ON {self.collection_name}
+                USING hnsw (vector vector_cosine_ops)
+            """
+            )
+
+        self.conn.commit()
+
+    def insert(
+        self,
+        vectors: List[List[float]],
+        payloads: Optional[List[Dict[str, Any]]] = None,
+        ids: Optional[List[str]] = None,
+    ) -> None:
+        """
+        Insert vectors into a collection.
+
+        Args:
+            vectors (List[List[float]]): List of vectors to insert.
+            payloads (List[Dict], optional): List of payloads corresponding to vectors.
+            ids (List[str], optional): List of IDs corresponding to vectors.
+        """
+        # Validate vector dimensions
+        self._validate_vector_dimensions(vectors)
+
+        logger.info(f"Inserting {len(vectors)} vectors into collection {self.collection_name}")
+        json_payloads = [json.dumps(payload) for payload in payloads]
+
+        data = [(id, vector, payload) for id, vector, payload in zip(ids, vectors, json_payloads)]
+        execute_values(
+            self.cur,
+            f"INSERT INTO {self.collection_name} (id, vector, payload) VALUES %s",
+            data,
+        )
+        self.conn.commit()
+
+    def search(
+        self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict[str, Any]] = None
+    ) -> List[OutputData]:
+        """
+        Search for similar vectors.
+
+        Args:
+            query (str): Query.
+            vectors (List[float]): Query vector.
+            limit (int, optional): Number of results to return. Defaults to 5.
+            filters (Dict, optional): Filters to apply to the search. Defaults to None.
+
+        Returns:
+            list: Search results.
+        """
+        filter_conditions = []
+        filter_params = []
+
+        if filters:
+            for k, v in filters.items():
+                filter_conditions.append("payload->>%s = %s")
+                filter_params.extend([k, str(v)])
+
+        filter_clause = "WHERE " + " AND ".join(filter_conditions) if filter_conditions else ""
+
+        self.cur.execute(
+            f"""
+            SELECT id, vector <=> %s::vector AS distance, payload
+            FROM {self.collection_name}
+            {filter_clause}
+            ORDER BY distance
+            LIMIT %s
+        """,
+            (vectors, *filter_params, limit),
+        )
+
+        results = self.cur.fetchall()
+        return [OutputData(id=str(r[0]), score=float(r[1]), payload=r[2]) for r in results]
+
+    def delete(self, vector_id: str) -> None:
+        """
+        Delete a vector by ID.
+
+        Args:
+            vector_id (str): ID of the vector to delete.
+        """
+        self.cur.execute(f"DELETE FROM {self.collection_name} WHERE id = %s", (vector_id,))
+        self.conn.commit()
+
+    def update(
+        self, vector_id: str, vector: Optional[List[float]] = None, payload: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """
+        Update a vector and its payload.
+
+        Args:
+            vector_id (str): ID of the vector to update.
+            vector (List[float], optional): Updated vector.
+            payload (Dict, optional): Updated payload.
+        """
+        if vector:
+            self.cur.execute(
+                f"UPDATE {self.collection_name} SET vector = %s WHERE id = %s",
+                (vector, vector_id),
+            )
+        if payload:
+            self.cur.execute(
+                f"UPDATE {self.collection_name} SET payload = %s WHERE id = %s",
+                (psycopg2.extras.Json(payload), vector_id),
+            )
+        self.conn.commit()
+
+    def get(self, vector_id: str) -> Optional[OutputData]:
+        """
+        Retrieve a vector by ID.
+
+        Args:
+            vector_id (str): ID of the vector to retrieve.
+
+        Returns:
+            OutputData: Retrieved vector.
+        """
+        self.cur.execute(
+            f"SELECT id, vector, payload FROM {self.collection_name} WHERE id = %s",
+            (vector_id,),
+        )
+        result = self.cur.fetchone()
+        if not result:
+            return None
+        return OutputData(id=str(result[0]), score=None, payload=result[2])
+
+    def list_collections(self) -> List[str]:
+        """
+        List all collections.
+
+        Returns:
+            List[str]: List of collection names.
+        """
+        self.cur.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'")
+        return [row[0] for row in self.cur.fetchall()]
+
+    def delete_collection(self) -> None:
+        """Delete a collection."""
+        self.cur.execute(f"DROP TABLE IF EXISTS {self.collection_name}")
+        self.conn.commit()
+
+    def collection_info(self) -> Dict[str, Any]:
+        """
+        Get information about a collection.
+
+        Returns:
+            Dict[str, Any]: Collection information.
+        """
+        self.cur.execute(
+            f"""
+            SELECT
+                table_name,
+                (SELECT COUNT(*) FROM {self.collection_name}) as row_count,
+                (SELECT pg_size_pretty(pg_total_relation_size('{self.collection_name}'))) as total_size
+            FROM information_schema.tables
+            WHERE table_schema = 'public' AND table_name = %s
+        """,
+            (self.collection_name,),
+        )
+        result = self.cur.fetchone()
+        return {"name": result[0], "count": result[1], "size": result[2]}
+
+    def list(self, filters: Optional[Dict[str, Any]] = None, limit: Optional[int] = None) -> List[OutputData]:
+        """
+        List all vectors in a collection.
+
+        Args:
+            filters (Dict, optional): Filters to apply to the list.
+            limit (int, optional): Number of vectors to return. Defaults to 10000 if None.
+
+        Returns:
+            List[OutputData]: List of vectors.
+        """
+        filter_conditions = []
+        filter_params = []
+
+        if filters:
+            for k, v in filters.items():
+                filter_conditions.append("payload->>%s = %s")
+                filter_params.extend([k, str(v)])
+
+        filter_clause = "WHERE " + " AND ".join(filter_conditions) if filter_conditions else ""
+
+        # Handle None limit by using a large default value
+        if limit is None:
+            limit = 10000  # Default large limit
+
+        query = f"""
+            SELECT id, vector, payload
+            FROM {self.collection_name}
+            {filter_clause}
+            LIMIT %s
+        """
+
+        self.cur.execute(query, (*filter_params, limit))
+
+        results = self.cur.fetchall()
+        return [OutputData(id=str(r[0]), score=None, payload=r[2]) for r in results]
+
+    def __del__(self) -> None:
+        """
+        Close the database connection when the object is deleted.
+        """
+        if hasattr(self, "cur"):
+            self.cur.close()
+        if hasattr(self, "conn"):
+            self.conn.close()
+
+    def reset(self) -> None:
+        """Reset the index by deleting and recreating it."""
+        logger.warning(f"Resetting index {self.collection_name}...")
+        self.delete_collection()
+        self.create_collection(self.embedding_model_dims)