pgvector-template 0.3.3__tar.gz → 0.3.5__tar.gz

{pgvector_template-0.3.3 → pgvector_template-0.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pgvector-template
-Version: 0.3.3
+Version: 0.3.5
 Summary: Template library for flexible PGVector RAG implementations
 Author-email: DL <v49t9zpqd@mozmail.com>
 License: MIT
@@ -21,7 +21,8 @@ Requires-Dist: pytest>=7.0.0; extra == "test"
 Requires-Dist: pytest-cov>=4.0.0; extra == "test"
 Requires-Dist: python-dotenv>=1.0.0; extra == "test"
 Provides-Extra: dev
-Requires-Dist: black>=23.0.0; extra == "dev"
+Requires-Dist: ruff>=0.9.0; extra == "dev"
+Requires-Dist: ty>=0.0.1a0; extra == "dev"
 Provides-Extra: dist
 Requires-Dist: build>=1.2.2; extra == "dist"
 Requires-Dist: twine>=6.1.0; extra == "dist"

{pgvector_template-0.3.3 → pgvector_template-0.3.5}/pgvector_template/__init__.py

@@ -3,4 +3,4 @@ from pgvector_template.db.connection import DatabaseManager
 from pgvector_template.db.document_db import DocumentDatabaseManager
 
 # Re-export at the top level
-__all__ = ["DatabaseManager", "DocumentDatabaseManager"]
+__all__ = ["DatabaseManager", "DocumentDatabaseManager"]

{pgvector_template-0.3.3 → pgvector_template-0.3.5}/pgvector_template/core/__init__.py

@@ -1,4 +1,8 @@
-from pgvector_template.core.document import BaseDocument, BaseDocumentMetadata, BaseDocumentOptionalProps
+from pgvector_template.core.document import (
+    BaseDocument,
+    BaseDocumentMetadata,
+    BaseDocumentOptionalProps,
+)
 from pgvector_template.core.embedder import BaseEmbeddingProvider
 from pgvector_template.core.manager import BaseCorpusManager, BaseCorpusManagerConfig, Corpus
 from pgvector_template.core.search import BaseSearchClientConfig, BaseSearchClient

{pgvector_template-0.3.3 → pgvector_template-0.3.5}/pgvector_template/core/document.py

@@ -2,7 +2,7 @@ from datetime import datetime
 from typing import Any, Type, TypeVar
 from uuid import uuid4, UUID as UuidLiteral
 
-from pydantic import BaseModel, Field, field_validator, model_validator
+from pydantic import BaseModel, Field, field_validator
 from sqlalchemy import (
     Column,
     String,

{pgvector_template-0.3.3 → pgvector_template-0.3.5}/pgvector_template/core/manager.py

@@ -1,7 +1,7 @@
-from abc import ABC, abstractmethod
+from abc import ABC
 from dataclasses import dataclass
 from logging import getLogger
-from typing import Any, Literal, Type
+from typing import Any, Type
 from uuid import UUID, uuid4
 
 from pydantic import BaseModel, Field
@@ -88,7 +88,7 @@ class BaseCorpusManager(ABC):
             self.session.query(self.config.document_cls)
             .filter(
                 self.config.document_cls.corpus_id == corpus_id,
-                self.config.document_cls.is_deleted == False,
+                self.config.document_cls.is_deleted == False,  # noqa: E712
             )
             .order_by(self.config.document_cls.chunk_index)
             .all()
@@ -105,6 +105,15 @@ class BaseCorpusManager(ABC):
             documents=chunks,
         )
 
+    def _generate_corpus_id(
+        self, corpus_id: UUID | str | None = None, content: str | None = None
+    ) -> UUID | str:
+        """
+        Helper to generate a corpus_id if not provided.
+        Override this method to provide a consistent way to generate.
+        """
+        return corpus_id if corpus_id else uuid4()
+
     def insert_corpus(
         self,
         content: str,
@@ -126,8 +135,7 @@ class BaseCorpusManager(ABC):
         Returns:
             int: The number of **documents** inserted for the provided corpus
         """
-        if not corpus_id:
-            corpus_id = uuid4()
+        corpus_id = self._generate_corpus_id(corpus_id)
         document_contents = self._split_corpus(content)
         document_embeddings = self.embedding_provider.embed_batch(document_contents)
         return self.insert_documents(
@@ -214,9 +222,12 @@ class BaseCorpusManager(ABC):
 
     def _delete_existing_corpus(self, corpus_id: UUID | str) -> None:
         """Delete all existing documents for the given corpus_id"""
-        self.session.query(self.config.document_cls).filter(
+        query = self.session.query(self.config.document_cls).filter(
             self.config.document_cls.corpus_id == corpus_id
-        ).delete()
+        )
+        if count := query.count():
+            logger.warning(f"Deleting {count} existing documents for corpus_id: {corpus_id}")
+            query.delete()
 
     def _split_corpus(self, content: str, **kwargs) -> list[str]:
         """
@@ -244,9 +255,9 @@ class BaseCorpusManager(ABC):
             logger.warning(
                 "Using default _join_documents. Override this method to improve functionality."
             )
-        documents.sort(key=lambda d: d.chunk_index)  # type: ignore
+        documents.sort(key=lambda d: d.chunk_index)
         # since _split_corpus performs a simple split on every 1000 chars, we can simply call `join`
-        corpus_content = "".join(d.content for d in documents)  # type: ignore
+        corpus_content = "".join(d.content for d in documents)  # ty: ignore[no-matching-overload]
         corpus_metadata = self._infer_corpus_metadata(documents)
         return corpus_content, corpus_metadata
 

{pgvector_template-0.3.3 → pgvector_template-0.3.5}/pgvector_template/db/connection.py

@@ -2,7 +2,7 @@ from contextlib import contextmanager
 from logging import getLogger
 from typing import Generator, Type
 
-from sqlalchemy import create_engine, text
+from sqlalchemy import Engine, create_engine, text
 from sqlalchemy.orm import sessionmaker, Session
 from sqlalchemy.ext.declarative import DeclarativeMeta
 
@@ -12,21 +12,29 @@ class DatabaseManager:
 
     def __init__(self, database_url: str):
         self.database_url = database_url
-        self.engine = None
+        self.engine: Engine | None = None
         self.SessionLocal = None
         self.logger = getLogger(self.__class__.__name__)
 
+    def _require_engine(self) -> Engine:
+        if self.engine is None:
+            raise RuntimeError("DatabaseManager not initialized. Call initialize() first.")
+        return self.engine
+
     def initialize(self):
         """Initialize database connection and session factory"""
         self.engine = create_engine(
-            self.database_url, pool_pre_ping=True, pool_recycle=300, echo=False  # Set to True for SQL debugging
+            self.database_url,
+            pool_pre_ping=True,
+            pool_recycle=300,
+            echo=False,  # Set to True for SQL debugging
         )
         self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine)
         self._ensure_pgvector_extension()
 
     def create_schema(self, schema_name: str) -> None:
         """Create a new schema for a collection type"""
-        with self.engine.connect() as conn:
+        with self._require_engine().connect() as conn:
             conn.execute(text(f"CREATE SCHEMA IF NOT EXISTS {schema_name}"))
             conn.commit()
 
@@ -37,12 +45,12 @@ class DatabaseManager:
         # Ensure all tables in this base class use the specified schema
         for table in base_class.metadata.tables.values():
             table.schema = schema_name
-        base_class.metadata.create_all(self.engine, checkfirst=True)
+        base_class.metadata.create_all(self._require_engine(), checkfirst=True)
         self.logger.info(f"Created tables for schema: {schema_name}")
 
     def _ensure_pgvector_extension(self):
         """Ensure pgvector extension is available"""
-        with self.engine.connect() as conn:
+        with self._require_engine().connect() as conn:
             conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
             conn.commit()
 
@@ -56,6 +64,8 @@ class DatabaseManager:
             # do stuff with the Session here
         ```
         """
+        if self.SessionLocal is None:
+            raise RuntimeError("DatabaseManager not initialized. Call initialize() first.")
         session = self.SessionLocal()
         try:
             yield session

{pgvector_template-0.3.3 → pgvector_template-0.3.5}/pgvector_template/db/document_db.py

@@ -23,7 +23,9 @@ class DocumentDatabaseManager(DatabaseManager):
 
     SCHEMA_PREFIX = "knowledge_base_"
 
-    def __init__(self, database_url: str, schema_suffix: str, document_classes: List[Type[BaseDocument]]):
+    def __init__(
+        self, database_url: str, schema_suffix: str, document_classes: List[Type[BaseDocument]]
+    ):
         """
         Initialize a document-oriented database manager
 
@@ -79,7 +81,7 @@ class TempDocumentDatabaseManager(DocumentDatabaseManager):
         if not schema_name.startswith("temp_"):
             raise ValueError("Can only drop schemas with 'temp_' prefix for safety")
 
-        with self.engine.connect() as conn:
+        with self._require_engine().connect() as conn:
             conn.execute(text(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE"))
             conn.commit()
 

pgvector_template-0.3.5/pgvector_template/models/__init__.py

@@ -0,0 +1,5 @@
+from pgvector_template.models.search import (
+    MetadataFilter as MetadataFilter,
+    RetrievalResult as RetrievalResult,
+    SearchQuery as SearchQuery,
+)

{pgvector_template-0.3.3 → pgvector_template-0.3.5}/pgvector_template/models/search.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass, asdict
 from datetime import datetime
-from typing import Any, Literal, Type
+from typing import Any, Literal
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 

{pgvector_template-0.3.3 → pgvector_template-0.3.5}/pgvector_template/service/document_service.py

@@ -3,8 +3,7 @@ Document service layer combining corpus management and search capabilities.
 """
 
 from logging import getLogger
-from typing import Any, Generic, Type, TypeVar
-from uuid import UUID
+from typing import Generic, Type, TypeVar
 
 from pydantic import BaseModel, Field
 from sqlalchemy.orm import Session

{pgvector_template-0.3.3 → pgvector_template-0.3.5}/pgvector_template/utils/metadata_filter.py

@@ -66,9 +66,9 @@ def validate_condition_compatibility(field_type: Type, condition: str) -> None:
         if origin is list:
             field_type = list
         elif len(args) > 0:
-            field_type = args[0]  # First non-None type
+            field_type = args[0]  # ty: ignore[index-out-of-bounds]
 
-    valid_conditions = {
+    valid_conditions: dict[type, set[str]] = {
         str: {"eq", "gt", "gte", "lt", "lte", "in", "exists"},
         int: {"eq", "gt", "gte", "lt", "lte", "exists"},
         float: {"eq", "gt", "gte", "lt", "lte", "exists"},
@@ -76,6 +76,6 @@ def validate_condition_compatibility(field_type: Type, condition: str) -> None:
         list: {"contains", "in", "exists"},
     }
 
-    allowed = valid_conditions.get(field_type, {"eq", "exists"})
+    allowed: set[str] = valid_conditions.get(field_type, {"eq", "exists"})
     if condition not in allowed:
         raise ValueError(f"Condition '{condition}' not valid for field type {field_type.__name__}")

{pgvector_template-0.3.3 → pgvector_template-0.3.5}/pgvector_template.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pgvector-template
-Version: 0.3.3
+Version: 0.3.5
 Summary: Template library for flexible PGVector RAG implementations
 Author-email: DL <v49t9zpqd@mozmail.com>
 License: MIT
@@ -21,7 +21,8 @@ Requires-Dist: pytest>=7.0.0; extra == "test"
 Requires-Dist: pytest-cov>=4.0.0; extra == "test"
 Requires-Dist: python-dotenv>=1.0.0; extra == "test"
 Provides-Extra: dev
-Requires-Dist: black>=23.0.0; extra == "dev"
+Requires-Dist: ruff>=0.9.0; extra == "dev"
+Requires-Dist: ty>=0.0.1a0; extra == "dev"
 Provides-Extra: dist
 Requires-Dist: build>=1.2.2; extra == "dist"
 Requires-Dist: twine>=6.1.0; extra == "dist"

{pgvector_template-0.3.3 → pgvector_template-0.3.5}/pgvector_template.egg-info/requires.txt

@@ -4,7 +4,8 @@ sqlalchemy>=2.0.0
 typing-extensions>=4.0.0
 
 [dev]
-black>=23.0.0
+ruff>=0.9.0
+ty>=0.0.1a0
 
 [dist]
 build>=1.2.2

{pgvector_template-0.3.3 → pgvector_template-0.3.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pgvector-template"
-version = "0.3.3"
+version = "0.3.5"
 description = "Template library for flexible PGVector RAG implementations"
 authors = [{ name="DL", email="v49t9zpqd@mozmail.com" }]
 license = { text = "MIT" }
@@ -30,7 +30,8 @@ test = [
     "python-dotenv>=1.0.0",
 ]
 dev = [
-    "black>=23.0.0",
+    "ruff>=0.9.0",
+    "ty>=0.0.1a0",
 ]
 dist = [
     "build>=1.2.2",
@@ -46,6 +47,9 @@ where = ["."]  # Look for packages in the root directory
 include = ["pgvector_template*"]  # Include all packages starting with pgvector_template
 
 
-[tool.black]
+[tool.ruff]
 line-length = 100
-target-version = ["py311", "py312", "py313"]
+target-version = "py311"
+
+[tool.ty.src]
+exclude = []

pgvector_template-0.3.3/pgvector_template/models/__init__.py

@@ -1 +0,0 @@
-from pgvector_template.models.search import *