wizit-context-ingestor 0.3.0b2__py3-none-any.whl → 0.4.0__py3-none-any.whl

wizit_context_ingestor/infra/rag/pg_embeddings.py

@@ -1,13 +1,13 @@
-from langchain_core.documents import Document
-from langchain.indexes import index, SQLRecordManager
-from typing import List
 import logging
-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()
+from langchain.indexes import IndexingResult, SQLRecordManager, aindex, index
+from langchain_core.documents import Document
+from langchain_postgres import PGEngine, PGVectorStore
+from langchain_postgres.v2.indexes import HNSWIndex
+from sqlalchemy.ext.asyncio import create_async_engine
+from typing_extensions import Literal
+
+from wizit_context_ingestor.application.interfaces import EmbeddingsManager
 
 logger = logging.getLogger(__name__)
 
@@ -42,7 +42,17 @@ class PgEmbeddingsManager(EmbeddingsManager):
 
     __slots__ = ("embeddings_model", "pg_connection")
 
-    def __init__(self, embeddings_model, pg_connection: str):
+    def __init__(
+        self,
+        embeddings_model,
+        pg_connection: str,
+        embeddings_vectors_table_name: str = "langchain_pg_embedding",
+        vector_size: int = 768,
+        content_column: str = "document",
+        id_column: str = "id",
+        metadata_json_column: str = "cmetadata",
+        records_manager_table_name: str = "langchain_record_manager",
+    ):
         """
         Initialize the PgEmbeddingsManager.
 
@@ -57,155 +67,135 @@ class PgEmbeddingsManager(EmbeddingsManager):
         """
         self.pg_connection = pg_connection
         self.embeddings_model = embeddings_model
-        self.pg_engine = None
         self.vector_store = None
         self.record_manager = None
-        try:
-            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
+        self.pg_engine = PGEngine.from_connection_string(pg_connection)
+        self.embeddings_vectors_table_name = embeddings_vectors_table_name
+        self.vector_size = vector_size
+        self.content_column = content_column
+        self.id_column = id_column
+        self.metadata_json_column = metadata_json_column
+        self.records_manager_table_name = records_manager_table_name
+        # self.async_engine = create_async_engine(pg_connection)
+        # self.pg_engine = PGEngine.from_engine(
+        #     self.async_engine
+        # )
+        logger.info("PgEmbeddingsManager initialized")
 
     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.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(
+        try:
+            self.pg_engine.init_vectorstore_table(
+                table_name=self.embeddings_vectors_table_name,
+                vector_size=self.vector_size,
+                content_column=self.content_column,
+                id_column=self.id_column,
+                metadata_json_column=self.metadata_json_column,
+            )
+            record_manager = SQLRecordManager(
+                self.records_manager_table_name,
+                db_url=self.pg_connection,
+                async_mode=False,
+            )
+            record_manager.create_schema()
+        except Exception as e:
+            logger.error(f"Error configure_vector_store: {e}")
+            raise
+
+    def retrieve_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.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):
+    ) -> tuple[PGVectorStore, SQLRecordManager]:
+        try:
+            self.vector_store = PGVectorStore.create_sync(
+                embedding_service=self.embeddings_model,
+                engine=self.pg_engine,
+                table_name=self.embeddings_vectors_table_name,
+                content_column=self.content_column,
+                metadata_json_column=self.metadata_json_column,
+                id_column=self.id_column,
+            )
+            self.record_manager = SQLRecordManager(
+                self.records_manager_table_name, db_url=self.pg_connection
+            )
+            return (self.vector_store, self.record_manager)
+        except Exception as e:
+            logger.error(f"Error retrieve vector store: ", e)
+            raise e
+
+    def check_vector_store_init(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")
+            if self.vector_store is None or self.record_manager is None:
+                self.retrieve_vector_store()
             return func(self, *args, **kwargs)
 
         return wrapper
 
-    @vector_store_initialized
-    def index_documents(self, docs: List[Document]):
+    @check_vector_store_init
+    def create_index(self):
+        try:
+            if self.vector_size < 2000:
+                index = HNSWIndex()
+                self.vector_store.apply_vector_index(index)
+            else:
+                raise NotImplementedError(
+                    "Indexing for vector size > 2000 is not supported"
+                )
+        except Exception as e:
+            logger.info(f"Error creating index: {e}")
+            raise e
+
+    @check_vector_store_init
+    def index_documents(
+        self,
+        docs: list[Document],
+        cleanup: Literal["incremental", "full", "scoped_full"] | None = "incremental",
+        source_id_key: str = "source",
+    ) -> IndexingResult:
         """
-        Add documents to the vector store with their embeddings.
+        Index documents in the vector store with their embeddings.
 
-        This method takes a list of Document objects, generates embeddings for them
-        using the embeddings model, and stores both the documents and their
-        embeddings in the PostgreSQL database.
+        This method takes a list of Document objects and indexes them using LangChain's
+        aindex function with incremental cleanup. The documents are processed through
+        the embeddings model and stored in the PostgreSQL database with pgvector.
 
         Args:
-          docs: A list of LangChain Document objects to add to the vector store
-                Each Document should have page_content and metadata attributes
-                from langchain_core.documents import Document
+            vector_store: The PGVectorStore instance to use for storage
+            record_manager: The SQLRecordManager instance for tracking indexed documents
+            docs: A list of LangChain Document objects to index in the vector store.
+                  Each Document should have page_content and metadata attributes.
+
         Returns:
-          None
+            IndexingResult: Result object containing information about the indexing operation
 
         Raises:
-          Exception: If there's an error adding documents to the vector store
+            Exception: If there's an error during the document indexing process
         """
         try:
             logger.info(f"Indexing {len(docs)} documents in vector store")
+            # await self.vector_store.aadd_documents(docs)
             return index(
                 docs,
                 self.record_manager,
                 self.vector_store,
-                cleanup="incremental",
-                source_id_key="source",
+                cleanup=cleanup,
+                source_id_key=source_id_key,
             )
         except Exception as e:
             logger.error(f"Error indexing documents: {str(e)}")
-            raise
+            raise e
 
-    @vector_store_initialized
-    def get_documents_keys_by_source_id(self, source_id: str):
-        """
-        Get document keys by source ID from the vector store.
-        """
-        try:
-            return self.record_manager.list_keys(group_ids=[source_id])
-        except Exception as e:
-            logger.error(f"Error getting documents keys by source ID: {str(e)}")
-            raise
-
-    @vector_store_initialized
-    def delete_documents_by_source_id(self, source_id: str):
-        """
-        Delete documents by source ID from the vector store.
-        """
+    @check_vector_store_init
+    def search_records(
+        self,
+        query: str,
+    ) -> list[Document]:
         try:
-            objects_keys = self.get_documents_keys_by_source_id(source_id)
-            self.record_manager.delete_keys(objects_keys)
-            self.vector_store.delete(ids=objects_keys)
+            logger.info(f"Searching for '{query}' in vector store")
+            reply = self.vector_store.search(query=query, search_type="similarity", k=1)
+            return reply
         except Exception as e:
-            logger.error(f"Error deleting documents by source ID: {str(e)}")
-            raise
-
-    # def get_retriever(self, search_type: str = "mmr", k: int = 20):
-    #     """
-    #     Get a retriever interface to the vector store for semantic search.
-
-    #     This method returns a LangChain retriever object that can be used in retrieval
-    #     pipelines, retrieval-augmented generation, and other LangChain chains.
-
-    #     Args:
-    #       search_type: The search algorithm to use. Options include:
-    #                    - "similarity" (standard cosine similarity)
-    #                    - "mmr" (Maximum Marginal Relevance, balances relevance with diversity)
-    #                    - "similarity_score_threshold" (filters by minimum similarity)
-    #       k: The number of documents to retrieve (default: 20)
-
-    #     Returns:
-    #       Retriever: A LangChain Retriever object that can be used in chains and pipelines
-
-    #     Raises:
-    #       Exception: If there's an error creating the retriever
-
-    #     Example:
-    #       >>> retriever = pg_manager.get_retriever(search_type="mmr", k=5)
-    #       >>> docs = retriever.get_relevant_documents("quantum computing")
-    #     """
-    #     try:
-    #         return self.vector_store.as_retriever(
-    #             search_type=search_type, search_kwargs={"k": k}
-    #         )
-    #     except Exception as e:
-    #         logger.info(f"failed to get vector store as retriever {str(e)}")
-    #         raise
+            logger.error(f"Error indexing documents: {str(e)}")
+            raise e

wizit_context_ingestor/infra/rag/semantic_chunks.py

@@ -1,10 +1,9 @@
-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
 # https://python.langchain.com/docs/how_to/embed_text/
 import logging
+import uuid
 from typing import List, Any
 from langchain_core.documents import Document
 from langchain_experimental.text_splitter import SemanticChunker
@@ -60,7 +59,7 @@ class SemanticChunks(RagChunker):
             source = document.metadata["source"]
             for i, chunk in enumerate(chunks):
                 if document.metadata["source"]:
-                    chunk.id = f"{source}-{i}"
+                    chunk.id = f"{uuid.uuid4()}"
             logger.info(f"{len(chunks)} chunks generated successfully")
             return chunks
         except Exception as e:

wizit_context_ingestor/infra/secrets/aws_secrets_manager.py

@@ -3,11 +3,10 @@ import logging
 
 logger = logging.getLogger(__name__)
 
-class AwsSecretsManager:
-
-    def __init__(self):
-        self.client = boto3_client('secretsmanager')
 
+class AwsSecretsManager:
+    def __init__(self, aws_region="us-east-1"):
+        self.client = boto3_client("secretsmanager", region_name=aws_region)
 
     def get_secret(self, secret_name):
         """

wizit_context_ingestor/main.py

@@ -11,7 +11,7 @@ from .infra.rag.chroma_embeddings import ChromaEmbeddingsManager
 from .infra.secrets.aws_secrets_manager import AwsSecretsManager
 from .data.storage import storage_services, StorageServices
 from .data.kdb import kdb_services, KdbServices
-from .utils.file_utils import has_invalid_file_name_format
+from .utils.file_utils import validate_file_name_format
 from langsmith import Client, tracing_context
 
 
@@ -78,7 +78,7 @@ class TranscriptionManager:
         llm_model_id: str = "claude-sonnet-4@20250514",
         target_language: str = "es",
         transcription_additional_instructions: str = "",
-        transcription_accuracy_threshold: int = 90,
+        transcription_accuracy_threshold: float = 0.90,
         max_transcription_retries: int = 2,
     ):
         self.gcp_project_id = gcp_project_id
@@ -116,18 +116,18 @@ class TranscriptionManager:
         return vertex_model
 
     def tracing(func):
-        def gen_tracing_context(self, *args, **kwargs):
+        async def gen_tracing_context(self, *args, **kwargs):
             with tracing_context(
                 enabled=True,
                 project_name=self.langsmith_project_name,
                 client=self.langsmith_client,
             ):
-                return func(self, *args, **kwargs)
+                return await func(self, *args, **kwargs)
 
         return gen_tracing_context
 
     @tracing
-    def transcribe_document(self, file_key: str):
+    async def transcribe_document(self, file_key: str):
         """Transcribe a document from source storage to target storage.
         This method serves as a generic interface for transcribing documents from
         various storage sources to target destinations. The specific implementation
@@ -143,7 +143,7 @@ class TranscriptionManager:
             Exception: If an error occurs during the transcription process.
         """
         try:
-            if has_invalid_file_name_format(file_key):
+            if not validate_file_name_format(file_key):
                 raise ValueError(
                     "Invalid file name format, do not provide special characters or spaces (instead use underscores or hyphens)"
                 )
@@ -162,9 +162,10 @@ class TranscriptionManager:
                 transcription_accuracy_threshold=self.transcription_accuracy_threshold,
                 max_transcription_retries=self.max_transcription_retries,
             )
-            parsed_pages, parsed_document = (
-                transcribe_document_service.process_document(file_key)
-            )
+            (
+                parsed_pages,
+                parsed_document,
+            ) = await transcribe_document_service.process_document(file_key)
             source_storage_file_tags = {}
             if persistence_service.supports_tagging:
                 # source_storage_file_tags.tag_file(file_key, {"status": "transcribed"})
@@ -231,18 +232,18 @@ class ChunksManager:
         return vertex_model
 
     def tracing(func):
-        def gen_tacing_context(self, *args, **kwargs):
+        async def gen_tracing_context(self, *args, **kwargs):
             with tracing_context(
                 enabled=True,
                 project_name=self.langsmith_project_name,
                 client=self.langsmith_client,
             ):
-                return func(self, *args, **kwargs)
+                return await func(self, *args, **kwargs)
 
-        return gen_tacing_context
+        return gen_tracing_context
 
     @tracing
-    def gen_context_chunks(
+    async def gen_context_chunks(
         self, file_key: str, source_storage_route: str, target_storage_route: str
     ):
         try:
@@ -272,7 +273,7 @@ class ChunksManager:
                 target_language=self.target_language,
             )
             context_chunks = (
-                context_chunks_in_document_service.get_context_chunks_in_document(
+                await context_chunks_in_document_service.get_context_chunks_in_document(
                     file_key, target_bucket_file_tags
                 )
             )

wizit_context_ingestor/main_chunks.py

@@ -0,0 +1,173 @@
+import json
+from logging import getLogger
+from typing import Any, Dict, Literal
+
+from langchain_core.documents import Document
+from langsmith import Client, tracing_context
+
+from .application.context_chunk_service import ContextChunksInDocumentService
+from .application.kdb_service import KdbService
+from .data.storage import StorageServices
+from .infra.persistence.local_storage import LocalStorageService
+from .infra.persistence.s3_storage import S3StorageService
+from .infra.rag.pg_embeddings import PgEmbeddingsManager
+from .infra.rag.semantic_chunks import SemanticChunks
+from .infra.secrets.aws_secrets_manager import AwsSecretsManager
+from .infra.vertex_model import VertexModels
+from .utils.file_utils import validate_file_name_format
+
+logger = getLogger(__name__)
+
+
+class KdbManager:
+    def __init__(
+        self, embeddings_model, kdb_service: Literal["pg"], kdb_params: Dict[Any, Any]
+    ):
+        self.kdb_service = kdb_service
+        self.kdb_params = kdb_params
+        self.embeddings_model = embeddings_model
+
+    def retrieve_kdb_service(self):
+        return PgEmbeddingsManager(self.embeddings_model, **self.kdb_params)
+
+
+class PersistenceManager:
+    def __init__(
+        self,
+        storage_service: Literal["s3", "local"],
+        source_storage_route,
+        target_storage_route,
+    ):
+        self.storage_service = storage_service
+        self.source_storage_route = source_storage_route
+        self.target_storage_route = target_storage_route
+
+    def retrieve_storage_service(self):
+        if self.storage_service == StorageServices.S3.value:
+            return S3StorageService(
+                origin_bucket_name=self.source_storage_route,
+                target_bucket_name=self.target_storage_route,
+            )
+        elif self.storage_service == StorageServices.LOCAL.value:
+            return LocalStorageService(
+                source_storage_route=self.source_storage_route,
+                target_storage_route=self.target_storage_route,
+            )
+        else:
+            raise ValueError(f"Unsupported storage service: {self.storage_service}")
+
+
+class ChunksManager:
+    def __init__(
+        self,
+        gcp_project_id: str,
+        gcp_project_location: str,
+        gcp_secret_name: str,
+        langsmith_api_key: str,
+        langsmith_project_name: str,
+        storage_service: Literal["s3", "local"],
+        kdb_service: Literal["pg"],
+        kdb_params: Dict[Any, Any],
+        llm_model_id: str = "claude-3-5-haiku@20241022",
+        embeddings_model_id: str = "text-multilingual-embedding-002",
+        target_language: str = "es",
+    ):
+        self.gcp_project_id = gcp_project_id
+        self.gcp_project_location = gcp_project_location
+        self.aws_secrets_manager = AwsSecretsManager()
+        self.gcp_secret_name = gcp_secret_name
+        self.llm_model_id = llm_model_id
+        self.target_language = target_language
+        self.gcp_sa_dict = self._get_gcp_sa_dict(gcp_secret_name)
+        self.storage_service = storage_service
+        self.kdb_params = kdb_params
+        self.kdb_service = kdb_service
+        self.vertex_model = self._get_vertex_model()
+        self.embeddings_model = self.vertex_model.load_embeddings_model(
+            embeddings_model_id
+        )
+        self.langsmith_api_key = langsmith_api_key
+        self.langsmith_project_name = langsmith_project_name
+        self.langsmith_client = Client(api_key=self.langsmith_api_key)
+        self.kdb_manager = KdbManager(self.embeddings_model, "pg", self.kdb_params)
+        self.pg_embeddings_manager = self.kdb_manager.retrieve_kdb_service()
+        self.rag_chunker = SemanticChunks(self.embeddings_model)
+        self.kdb_service = KdbService(
+            self.pg_embeddings_manager,
+        )
+
+    def _get_gcp_sa_dict(self, gcp_secret_name: str):
+        vertex_gcp_sa = self.aws_secrets_manager.get_secret(gcp_secret_name)
+        vertex_gcp_sa_dict = json.loads(vertex_gcp_sa)
+        return vertex_gcp_sa_dict
+
+    def _get_vertex_model(self):
+        vertex_model = VertexModels(
+            self.gcp_project_id,
+            self.gcp_project_location,
+            self.gcp_sa_dict,
+            llm_model_id=self.llm_model_id,
+        )
+        return vertex_model
+
+    def provision_vector_store(self):
+        try:
+            self.kdb_service.configure_kdb()
+            self.kdb_service.create_vector_store_hsnw_index()
+        except Exception as e:
+            logger.error(f"Error configuring vector store: {e}")
+
+    def index_documents_in_vector_store(self, docs: list[Document]):
+        try:
+            self.kdb_service.index_documents_in_vector_store(docs)
+        except Exception as e:
+            logger.error(f"Error indexing documents in vector store: {e}")
+
+    def search_records(self, query):
+        return self.kdb_service.search(query)
+
+    def tracing(func):
+        async def gen_tracing_context(self, *args, **kwargs):
+            with tracing_context(
+                enabled=True,
+                project_name=self.langsmith_project_name,
+                client=self.langsmith_client,
+            ):
+                return await func(self, *args, **kwargs)
+
+        return gen_tracing_context
+
+    @tracing
+    async def gen_context_chunks(
+        self, file_key: str, source_storage_route: str, target_storage_route: str
+    ):
+        try:
+            validate_file_name_format(file_key)
+            persistence_layer = PersistenceManager(
+                self.storage_service, source_storage_route, target_storage_route
+            )
+            persistence_service = persistence_layer.retrieve_storage_service()
+            target_bucket_file_tags = {}
+            if persistence_service.supports_tagging:
+                target_bucket_file_tags = persistence_service.retrieve_file_tags(
+                    file_key, target_storage_route
+                )
+            rag_chunker = SemanticChunks(self.embeddings_model)
+            kdb_manager = KdbManager(self.embeddings_model, "pg", self.kdb_params)
+            kdb_service = kdb_manager.retrieve_kdb_service()
+            context_chunks_in_document_service = ContextChunksInDocumentService(
+                ai_application_service=self.vertex_model,
+                persistence_service=persistence_service,
+                rag_chunker=rag_chunker,
+                embeddings_manager=kdb_service,
+                target_language=self.target_language,
+            )
+            context_chunks = (
+                await context_chunks_in_document_service.get_context_chunks_in_document(
+                    file_key, target_bucket_file_tags
+                )
+            )
+            return context_chunks
+        except Exception as e:
+            print(f"Error getting context chunks in document: {e}")
+            raise e

wizit_context_ingestor/utils/file_utils.py

@@ -1,13 +1,12 @@
 import re
 
 
-def has_invalid_file_name_format(file_name):
+def validate_file_name_format(file_name):
     """Check if file name has special characters or spaces instead of underscores"""
-    # Check for spaces
-    if " " in file_name:
+    # Check for special characters (anything that's not alphanumeric, underscore, dash, dot, slash, or backslash)
+    if re.search(r"[^a-zA-Z0-9_.\-/\\]", file_name) is None:
         return True
-
-    # Check for special characters (anything that's not alphanumeric, underscore, dash, or dot)
-    if re.search(r"[^a-zA-Z0-9_.-]", file_name):
-        return True
-    return False
+    else:
+        raise ValueError(
+            "Invalid file name format, do not provide special characters or spaces (instead use underscores or hyphens)"
+        )

{wizit_context_ingestor-0.3.0b2.dist-info → wizit_context_ingestor-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: wizit-context-ingestor
-Version: 0.3.0b2
+Version: 0.4.0
 Summary: Contextual Rag with Cloud Solutions
 Requires-Dist: anthropic[vertex]>=0.66.0
 Requires-Dist: boto3>=1.40.23
@@ -8,10 +8,13 @@ Requires-Dist: langchain-aws>=0.2.31
 Requires-Dist: langchain-chroma>=0.2.6
 Requires-Dist: langchain-experimental>=0.3.4
 Requires-Dist: langchain-google-vertexai>=2.0.28
+Requires-Dist: langchain-postgres>=0.0.16
 Requires-Dist: langchain-redis>=0.2.3
 Requires-Dist: langgraph>=0.6.8
 Requires-Dist: pillow>=11.3.0
+Requires-Dist: psycopg2-binary>=2.9.11
 Requires-Dist: pymupdf>=1.26.4
+Requires-Dist: sqlalchemy[asyncio]>=2.0.43
 Requires-Python: >=3.12
 Description-Content-Type: text/markdown