wizit-context-ingestor 0.2.5b3__py3-none-any.whl → 0.3.0b2__py3-none-any.whl

wizit_context_ingestor/infra/persistence/local_storage.py

@@ -3,22 +3,27 @@ from ...domain.models import ParsedDoc
 from typing import Optional
 import logging
 import os
+
 logger = logging.getLogger(__name__)
 
+
 class LocalStorageService(PersistenceService):
     """Persistence service for local storage."""
 
-    def __init__(self):
-        self.tmp_folder = "tmp"
+    def __init__(self, source_storage_route: str, target_storage_route: str):
+        self.source_storage_route = source_storage_route
+        self.target_storage_route = target_storage_route
+        self.supports_tagging = hasattr(self, "retrieve_file_tags")
 
     def load_markdown_file_content(self, file_key: str) -> str:
         """Load markdown file content from local storage."""
         file_content = None
-        with open(f"{self.tmp_folder}/{file_key}", "r", encoding="utf-8") as file:
+        with open(
+            f"{self.source_storage_route}/{file_key}", "r", encoding="utf-8"
+        ) as file:
             file_content = file.read()
         return file_content
 
-
     def retrieve_raw_file(self, file_key: str) -> str:
         """Retrieve file path in tmp folder from local storage.
 
@@ -32,16 +37,21 @@ class LocalStorageService(PersistenceService):
             ClientError: If there's an error retrieving the object from local storage
         """
         try:
-            tmp_file_path = f"{self.tmp_folder}/{file_key}"
+            tmp_file_path = f"{self.source_storage_route}/{file_key}"
             if not os.path.exists(tmp_file_path):
                 raise FileNotFoundError(f"File {file_key} not found in local storage")
             return tmp_file_path
         except Exception as e:
-            logger.error(f"Unexpected error retrieving file {file_key} from local storage: {str(e)}")
+            logger.error(
+                f"Unexpected error retrieving file {file_key} from local storage: {str(e)}"
+            )
             raise
 
-
-    def save_parsed_document(self, file_key: str, parsed_document: ParsedDoc, file_tags: Optional[dict] = {}):
+    def save_parsed_document(
+        self, file_key: str, parsed_document: ParsedDoc, file_tags: Optional[dict] = {}
+    ):
         """Save a parsed document."""
-        with open(f"{self.tmp_folder}/{file_key}", "w", encoding="utf-8") as f:
+        with open(
+            f"{self.target_storage_route}/{file_key}", "w", encoding="utf-8"
+        ) as f:
             f.write(parsed_document.document_text)

wizit_context_ingestor/infra/persistence/s3_storage.py

@@ -11,12 +11,19 @@ logger = logging.getLogger(__name__)
 
 class S3StorageService(PersistenceService):
     """Persistence service for S3 storage."""
-    __slots__ = ('origin_bucket_name', 'target_bucket_name', 'region_name')
-    def __init__(self, origin_bucket_name: str, target_bucket_name: str, region_name: str = 'us-east-1'):
-        self.s3 = boto3_client('s3', region_name=region_name)
+
+    __slots__ = ("origin_bucket_name", "target_bucket_name", "region_name")
+
+    def __init__(
+        self,
+        origin_bucket_name: str,
+        target_bucket_name: str,
+        region_name: str = "us-east-1",
+    ):
+        self.s3 = boto3_client("s3", region_name=region_name)
         self.origin_bucket_name = origin_bucket_name
         self.target_bucket_name = target_bucket_name
-
+        self.supports_tagging = hasattr(self, "retrieve_file_tags")
 
     def load_markdown_file_content(self, file_key: str) -> str:
         """Load markdown file content from S3 storage.
@@ -36,9 +43,9 @@ class S3StorageService(PersistenceService):
             response = self.s3.get_object(Bucket=self.target_bucket_name, Key=file_key)
             tmp_file_key = f"/tmp/{file_key}"
             os.makedirs(os.path.dirname(tmp_file_key), exist_ok=True)
-            with open(tmp_file_key, 'wb') as f:
-                f.write(response['Body'].read())
-            with open(tmp_file_key, 'r', encoding='utf-8') as f:
+            with open(tmp_file_key, "wb") as f:
+                f.write(response["Body"].read())
+            with open(tmp_file_key, "r", encoding="utf-8") as f:
                 file_content = f.read()
             return file_content
         except ClientError as e:
@@ -48,7 +55,6 @@ class S3StorageService(PersistenceService):
             logger.error(f"Unexpected error loading file {file_key} from S3: {str(e)}")
             raise
 
-
     def retrieve_raw_file(self, file_key: str) -> str:
         """Retrieve file path in tmp folder from S3 storage.
 
@@ -67,18 +73,21 @@ class S3StorageService(PersistenceService):
             tmp_file_key = f"/tmp/{file_key}"
             # Create parent directories if they don't exist
             os.makedirs(os.path.dirname(tmp_file_key), exist_ok=True)
-            with open(tmp_file_key, 'wb') as f:
-                f.write(response['Body'].read())
+            with open(tmp_file_key, "wb") as f:
+                f.write(response["Body"].read())
             return tmp_file_key
         except ClientError as e:
             logger.error(f"Error retrieving file {file_key} from S3: {str(e)}")
             raise
         except Exception as e:
-            logger.error(f"Unexpected error retrieving file {file_key} from S3: {str(e)}")
+            logger.error(
+                f"Unexpected error retrieving file {file_key} from S3: {str(e)}"
+            )
             raise
 
-
-    def save_parsed_document(self, file_key: str, parsed_document: ParsedDoc, file_tags: Optional[dict] = {}):
+    def save_parsed_document(
+        self, file_key: str, parsed_document: ParsedDoc, file_tags: Optional[dict] = {}
+    ):
         """Save a parsed document to S3.
 
         Args:
@@ -91,21 +100,21 @@ class S3StorageService(PersistenceService):
         """
         try:
             # Convert document content to bytes
-            content_bytes = parsed_document.document_text.encode('utf-8')
+            content_bytes = parsed_document.document_text.encode("utf-8")
             # Upload the file to S3
             if not file_tags:
                 self.s3.put_object(
-                    Bucket=self.target_bucket_name,
-                    Key=file_key,
-                    Body=content_bytes
+                    Bucket=self.target_bucket_name, Key=file_key, Body=content_bytes
                 )
             else:
-                tagging_string = "&".join([f"{key}={value}" for key, value in file_tags.items()])
+                tagging_string = "&".join(
+                    [f"{key}={value}" for key, value in file_tags.items()]
+                )
                 self.s3.put_object(
                     Bucket=self.target_bucket_name,
                     Key=file_key,
                     Body=content_bytes,
-                    Tagging=tagging_string
+                    Tagging=tagging_string,
                 )
 
             logger.info(f"Successfully saved document to S3 as {file_key}")
@@ -122,8 +131,5 @@ class S3StorageService(PersistenceService):
         Args:
             file_key: The key (path) to retrieve tags
         """
-        response = self.s3.get_object_tagging(
-            Bucket=bucket_name,
-            Key=file_key
-        )
+        response = self.s3.get_object_tagging(Bucket=bucket_name, Key=file_key)
         return {item["Key"]: item["Value"] for item in response["TagSet"]}

wizit_context_ingestor/infra/rag/chroma_embeddings.py

@@ -1,25 +1,21 @@
-from typing_extensions import Sequence
-from test.test_typing import CoolEmployee
 from langchain_core.documents import Document
 from langchain_chroma import Chroma
-from typing import List
 import logging
-from uuid import uuid4
 from ...application.interfaces import EmbeddingsManager
 
 # load_dotenv()
 
 logger = logging.getLogger(__name__)
 
+
 class ChromaEmbeddingsManager(EmbeddingsManager):
+    __slots__ = ("embeddings_model", "collection_name")
 
-    __slots__ = ("embeddings_model", "chroma_host", "collection_name", "metadata_tags")
     def __init__(
         self,
         embeddings_model,
-        chroma_host,
-        collection_name: str,
-        metadata_tags: dict
+        chroma_host=None,
+        **chroma_conn_kwargs,
     ):
         """
         Initialize the ChromaEmbeddingsManager.
@@ -27,33 +23,28 @@ class ChromaEmbeddingsManager(EmbeddingsManager):
             embeddings_model: The embeddings model to use for generating vector embeddings
                               (typically a LangChain embeddings model instance)
             chroma_host: The Chroma host URL
-            collection_name: The Chroma collection name
-            metadata_tags: Tags to add as metadata to Chroma vector store
 
         Raises:
             Exception: If there's an error initializing the RedisEmbeddingsManager
         """
-        self.collection_name = collection_name
         self.embeddings_model = embeddings_model
         self.chroma_host = chroma_host
-        self.metadata_tags_schema = []
-
-        for tag_key in metadata_tags:
-          self.metadata_tags_schema.append({
-              "type": "tag",
-              "name": tag_key
-          })
-
         try:
-            self.chroma = Chroma(
-                collection_name=self.collection_name,
-                embedding_function=self.embeddings_model,
-                host=self.chroma_host,
-            )
-            logger.info("ChromaEmbeddingsManager initialized")
+            if chroma_host:
+                self.chroma = Chroma(
+                    embedding_function=self.embeddings_model,
+                    host=chroma_host,
+                    **chroma_conn_kwargs,
+                )
+                logger.info("ChromaEmbeddingsManager initialized")
+            else:
+                self.chroma = Chroma(
+                    embedding_function=self.embeddings_model, **chroma_conn_kwargs
+                )
+                logger.info("ChromaEmbeddingsManager initialized")
         except Exception as e:
-          logger.error(f"Failed to initialize ChromaEmbeddingsManager: {str(e)}")
-          raise
+            logger.error(f"Failed to initialize ChromaEmbeddingsManager: {str(e)}")
+            raise
 
     def configure_vector_store(
         self,
@@ -61,8 +52,6 @@ class ChromaEmbeddingsManager(EmbeddingsManager):
         vector_size: int = 768,
         content_column: str = "document",
         id_column: str = "id",
-        metadata_json_column: str = "cmetadata",
-        pg_record_manager: str = ""
     ):
         """Configure the vector store."""
         pass
@@ -71,13 +60,11 @@ class ChromaEmbeddingsManager(EmbeddingsManager):
         self,
         table_name: str = "",
         content_column: str = "document",
-        metadata_json_column: str = "cmetadata",
         id_column: str = "id",
     ):
         """Initialize the vector store."""
         pass
 
-
     def index_documents(self, documents: list[Document]):
         """
         Add documents to the vector store with their embeddings.
@@ -123,6 +110,18 @@ class ChromaEmbeddingsManager(EmbeddingsManager):
             logger.error(f"Error deleting documents by ID: {str(e)}")
             raise
 
+    def delete_documents_by_metadata_key(self, metadata_key: str, metadata_value: str):
+        """
+        Delete documents by filter from the vector store.
+        """
+        try:
+            self.chroma.delete(where={metadata_key: metadata_value})
+        except Exception as error:
+            logger.error(
+                f"Error deleting documents by filter: {str(filter)}, error: {error} "
+            )
+            raise
+
     def get_documents_keys_by_source_id(self, source_id: str):
         """Get documents keys by source ID."""
         pass

wizit_context_ingestor/infra/rag/pg_embeddings.py

@@ -6,6 +6,7 @@ from langchain_postgres import PGVectorStore, PGEngine
 from sqlalchemy import create_engine
 from dotenv import load_dotenv
 from wizit_context_ingestor.application.interfaces import EmbeddingsManager
+
 load_dotenv()
 
 logger = logging.getLogger(__name__)
@@ -38,19 +39,21 @@ class PgEmbeddingsManager(EmbeddingsManager):
       ... )
       >>> documents = [Document(page_content="Sample text", metadata={"source": "example"})]
     """
+
     __slots__ = ("embeddings_model", "pg_connection")
+
     def __init__(self, embeddings_model, pg_connection: str):
         """
-          Initialize the PgEmbeddingsManager.
+        Initialize the PgEmbeddingsManager.
 
-          Args:
-              embeddings_model: The embeddings model to use for generating vector embeddings
-                                (typically a LangChain embeddings model instance)
-              pg_connection: The PostgreSQL connection string
-                            (format: postgresql://user:password@host:port/database)
+        Args:
+            embeddings_model: The embeddings model to use for generating vector embeddings
+                              (typically a LangChain embeddings model instance)
+            pg_connection: The PostgreSQL connection string
+                          (format: postgresql://user:password@host:port/database)
 
-          Raises:
-              Exception: If there's an error initializing the vector store
+        Raises:
+            Exception: If there's an error initializing the vector store
         """
         self.pg_connection = pg_connection
         self.embeddings_model = embeddings_model
@@ -58,65 +61,65 @@ class PgEmbeddingsManager(EmbeddingsManager):
         self.vector_store = None
         self.record_manager = None
         try:
-          self.pg_engine = PGEngine.from_connection_string(url=pg_connection)
-          logger.info("PgEmbeddingsManager initialized")
+            self.pg_engine = PGEngine.from_connection_string(url=pg_connection)
+            logger.info("PgEmbeddingsManager initialized")
         except Exception as e:
             logger.error(f"Failed to initialize PgEmbeddingsManager: {str(e)}")
             raise
 
     def configure_vector_store(
-      self,
-      table_name: str = "langchain_pg_embedding",
-      vector_size: int = 768,
-      content_column: str = "document",
-      id_column: str = "id",
-      metadata_json_column: str = "cmetadata",
-      pg_record_manager: str = "postgres/langchain_pg_collection"
+        self,
+        table_name: str = "langchain_pg_embedding",
+        vector_size: int = 768,
+        content_column: str = "document",
+        id_column: str = "id",
+        metadata_json_column: str = "cmetadata",
+        pg_record_manager: str = "postgres/langchain_pg_collection",
     ):
-      self.pg_engine.init_vectorstore_table(
-        table_name=table_name,
-        vector_size=vector_size,
-        content_column=content_column,
-        id_column=id_column,
-        metadata_json_column=metadata_json_column,
-      )
-      self.record_manager = SQLRecordManager(
-          pg_record_manager,
-          engine=create_engine(url=self.pg_connection)
-      )
-      # TODO move this from here
-      self.record_manager.create_schema()
+        self.pg_engine.init_vectorstore_table(
+            table_name=table_name,
+            vector_size=vector_size,
+            content_column=content_column,
+            id_column=id_column,
+            metadata_json_column=metadata_json_column,
+        )
+        self.record_manager = SQLRecordManager(
+            pg_record_manager, engine=create_engine(url=self.pg_connection)
+        )
+        # TODO move this from here
+        self.record_manager.create_schema()
 
     def init_vector_store(
-      self,
-      table_name: str = "langchain_pg_embedding",
-      content_column: str = "document",
-      metadata_json_column: str = "cmetadata",
-      id_column: str = "id",
-      pg_record_manager: str = "postgres/langchain_pg_collection"
+        self,
+        table_name: str = "langchain_pg_embedding",
+        content_column: str = "document",
+        metadata_json_column: str = "cmetadata",
+        id_column: str = "id",
+        pg_record_manager: str = "postgres/langchain_pg_collection",
     ):
-      self.vector_store = PGVectorStore.create_sync(
-        embedding_service=self.embeddings_model,
-        engine=self.pg_engine,
-        table_name=table_name,
-        content_column=content_column,
-        metadata_json_column=metadata_json_column,
-        id_column=id_column,
-      )
-      self.record_manager = SQLRecordManager(
-          pg_record_manager,
-          engine=create_engine(url=self.pg_connection)
-      )
+        self.vector_store = PGVectorStore.create_sync(
+            embedding_service=self.embeddings_model,
+            engine=self.pg_engine,
+            table_name=table_name,
+            content_column=content_column,
+            metadata_json_column=metadata_json_column,
+            id_column=id_column,
+        )
+        self.record_manager = SQLRecordManager(
+            pg_record_manager, engine=create_engine(url=self.pg_connection)
+        )
 
     def vector_store_initialized(func):
         """validate vector store initialization"""
+
         def wrapper(self, *args, **kwargs):
-          # Common validation logic
-          if self.vector_store is None:
-            raise Exception("Vector store not initialized")
-          if self.record_manager is None:
-            raise Exception("Record manager not initialized")
-          return func(self, *args, **kwargs)
+            # Common validation logic
+            if self.vector_store is None:
+                raise Exception("Vector store not initialized")
+            if self.record_manager is None:
+                raise Exception("Record manager not initialized")
+            return func(self, *args, **kwargs)
+
         return wrapper
 
     @vector_store_initialized

wizit_context_ingestor/infra/rag/redis_embeddings.py

@@ -2,6 +2,7 @@ from langchain_core.documents import Document
 from langchain_redis import RedisConfig, RedisVectorStore
 from typing import List
 import logging
+
 # from dotenv import load_dotenv
 from ...application.interfaces import EmbeddingsManager
 
@@ -9,10 +10,13 @@ from ...application.interfaces import EmbeddingsManager
 
 logger = logging.getLogger(__name__)
 
-class RedisEmbeddingsManager(EmbeddingsManager):
 
+class RedisEmbeddingsManager(EmbeddingsManager):
     __slots__ = ("embeddings_model", "redis_conn_string", "metadata_tags")
-    def __init__(self, embeddings_model, redis_conn_string: str, metadata_tags: dict):
+
+    def __init__(
+        self, embeddings_model, redis_conn_string: str, metadata_tags: List[str] = []
+    ):
         """
         Initialize the RedisEmbeddingsManager.
         Args:
@@ -27,27 +31,23 @@ class RedisEmbeddingsManager(EmbeddingsManager):
         """
         self.redis_conn_string = redis_conn_string
         self.embeddings_model = embeddings_model
-        self.metadata_tags_schema = []
-
+        self.metadata_tags_schema = [{"type": "text", "name": "context"}]
         for tag_key in metadata_tags:
-          self.metadata_tags_schema.append({
-              "type": "tag",
-              "name": tag_key
-          })
+            self.metadata_tags_schema.append({"type": "text", "name": tag_key})
 
         try:
-          self.redis_config = RedisConfig(
-            index_name="vector_store",
-            redis_url=self.redis_conn_string,
-            metadata_schema=[
-              {"type": "text", "name": "context"}
-            ]+self.metadata_tags_schema,
-          )
-          self.vector_store = RedisVectorStore(self.embeddings_model, config=self.redis_config)
-          logger.info("RedisEmbeddingsManager initialized")
+            self.redis_config = RedisConfig(
+                index_name="vector_store",
+                redis_url=self.redis_conn_string,
+                metadata_schema=self.metadata_tags_schema,
+            )
+            self.vector_store = RedisVectorStore(
+                self.embeddings_model, config=self.redis_config
+            )
+            logger.info("RedisEmbeddingsManager initialized")
         except Exception as e:
-          logger.error(f"Failed to initialize RedisEmbeddingsManager: {str(e)}")
-          raise
+            logger.error(f"Failed to initialize RedisEmbeddingsManager: {str(e)}")
+            raise
 
     def configure_vector_store(
         self,
@@ -56,7 +56,7 @@ class RedisEmbeddingsManager(EmbeddingsManager):
         content_column: str = "document",
         id_column: str = "id",
         metadata_json_column: str = "cmetadata",
-        pg_record_manager: str = "postgres/langchain_pg_collection"
+        pg_record_manager: str = "postgres/langchain_pg_collection",
     ):
         """Configure the vector store."""
         pass
@@ -73,13 +73,14 @@ class RedisEmbeddingsManager(EmbeddingsManager):
 
     def vector_store_initialized(func):
         """validate vector store initialization"""
+
         def wrapper(self, *args, **kwargs):
-          # Common validation logic
-          if self.vector_store is None:
-            raise Exception("Vector store not initialized")
-          return func(self, *args, **kwargs)
-        return wrapper
+            # Common validation logic
+            if self.vector_store is None:
+                raise Exception("Vector store not initialized")
+            return func(self, *args, **kwargs)
 
+        return wrapper
 
     @vector_store_initialized
     def index_documents(self, docs: List[Document]):
@@ -129,6 +130,14 @@ class RedisEmbeddingsManager(EmbeddingsManager):
             logger.error(f"Error deleting documents by ID: {str(e)}")
             raise
 
+    @vector_store_initialized
+    def delete_documents_by_metadata_key(self, metadata_key: str, metadata_value: str):
+        """
+        Delete documents by filter from the vector store.
+        """
+        # TODO investigate how to do this
+        pass
+
     def get_documents_keys_by_source_id(self, source_id: str):
         """Get documents keys by source ID."""
         pass

wizit_context_ingestor/infra/rag/semantic_chunks.py

@@ -1,3 +1,5 @@
+from posix import fork
+
 # check this documentation
 # https://python.langchain.com/docs/how_to/semantic-chunker/
 # https://github.com/FullStackRetrieval-com/RetrievalTutorials/blob/main/tutorials/LevelsOfTextSplitting/5_Levels_Of_Text_Splitting.ipynb
@@ -16,7 +18,9 @@ class SemanticChunks(RagChunker):
     Class for semantically chunking documents into smaller pieces based on semantic similarity.
     Uses LangChain's SemanticChunker to create semantically coherent document chunks.
     """
+
     __slots__ = ("embeddings_model",)
+
     def __init__(self, embeddings_model: Any):
         """
         Initialize a document chunker with an embeddings model.
@@ -35,7 +39,7 @@ class SemanticChunks(RagChunker):
             add_start_index=True,
             breakpoint_threshold_type="percentile",
             breakpoint_threshold_amount=95,
-            min_chunk_size=200
+            min_chunk_size=200,
         )
 
     def gen_chunks_for_document(self, document: Document) -> List[Document]:
@@ -53,6 +57,10 @@ class SemanticChunks(RagChunker):
         """
         try:
             chunks = self.text_splitter.split_documents([document])
+            source = document.metadata["source"]
+            for i, chunk in enumerate(chunks):
+                if document.metadata["source"]:
+                    chunk.id = f"{source}-{i}"
             logger.info(f"{len(chunks)} chunks generated successfully")
             return chunks
         except Exception as e: