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

wizit_context_ingestor/__init__.py

@@ -1,3 +1,4 @@
-from .main import ChunksManager, TranscriptionManager
+from .main import TranscriptionManager
+from .main_chunks import ChunksManager
 
 __all__ = ["ChunksManager", "TranscriptionManager"]

wizit_context_ingestor/application/context_chunk_service.py

@@ -1,19 +1,20 @@
 import asyncio
+import logging
+from typing import Any, Dict, List, Optional
+
+from langchain_core.documents import Document
+from langchain_core.messages.human import HumanMessage
 from langchain_core.output_parsers.pydantic import PydanticOutputParser
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.documents import Document
+
 from ..data.prompts import CONTEXT_CHUNKS_IN_DOCUMENT_SYSTEM_PROMPT, ContextChunk
-from langchain_core.messages.human import HumanMessage
 from ..workflows.context_workflow import ContextWorkflow
-from typing import Dict, Any, Optional, List
 from .interfaces import (
     AiApplicationService,
+    EmbeddingsManager,
     PersistenceService,
     RagChunker,
-    EmbeddingsManager,
 )
-import logging
-
 
 logger = logging.getLogger(__name__)
 
@@ -39,7 +40,7 @@ class ContextChunksInDocumentService:
         self.rag_chunker = rag_chunker
         self.embeddings_manager = embeddings_manager
         self.target_language = target_language
-        self.embeddings_manager.init_vector_store()
+        # self.embeddings_manager.init_vector_store()
         self.chat_model = self.ai_application_service.load_chat_model()
         # TODO
         self.context_additional_instructions = ""
@@ -85,74 +86,6 @@ class ContextChunksInDocumentService:
             logger.error(f"Failed to retrieve context chunks in document: {str(e)}")
             raise
 
-    # def _retrieve_context_chunk_in_document(
-    #     self,
-    #     markdown_content: str,
-    #     chunk: Document,
-    #     chunk_metadata: Optional[Dict[str, Any]] = None,
-    # ) -> Document:
-    #     """Retrieve context chunks in document."""
-    #     try:
-    #         chunk_output_parser = PydanticOutputParser(pydantic_object=ContextChunk)
-    #         # Create the prompt template with image
-    #         prompt = ChatPromptTemplate.from_messages(
-    #             [
-    #                 ("system", CONTEXT_CHUNKS_IN_DOCUMENT_SYSTEM_PROMPT),
-    #                 (
-    #                     "human",
-    #                     [
-    #                         {
-    #                             "type": "text",
-    #                             "text": f"Generate context for the following chunk: <chunk>{chunk.page_content}</chunk>,  ensure all content chunks are generated in '{self.target_language}' language",
-    #                         }
-    #                     ],
-    #                 ),
-    #             ]
-    #         ).partial(
-    #             document_content=markdown_content,
-    #             format_instructions=chunk_output_parser.get_format_instructions(),
-    #         )
-    #         model_with_structured_output = self.chat_model.with_structured_output(
-    #             ContextChunk
-    #         )
-    #         # Create the chain
-    #         chain = prompt | model_with_structured_output
-    #         # Process the image
-    #         results = chain.invoke({})
-    #         # chunk.page_content = (
-    #         #     f"Context:{results.context}, Content:{chunk.page_content}"
-    #         # )
-    #         chunk.metadata["context"] = results.context
-    #         if chunk_metadata:
-    #             for key, value in chunk_metadata.items():
-    #                 chunk.metadata[key] = value
-    #         return chunk
-
-    #     except Exception as e:
-    #         logger.error(f"Failed to retrieve context chunks in document: {str(e)}")
-    #         raise
-
-    # def retrieve_context_chunks_in_document(
-    #     self,
-    #     markdown_content: str,
-    #     chunks: List[Document],
-    #     chunks_metadata: Optional[Dict[str, Any]] = None,
-    # ) -> List[Document]:
-    #     """Retrieve context chunks in document."""
-    #     try:
-    #         context_chunks = list(
-    #             map(
-    #                 lambda chunk: self._retrieve_context_chunk_in_document(
-    #                     markdown_content, chunk, chunks_metadata
-    #                 ),
-    #                 chunks,
-    #             )
-    #         )
-    #         return context_chunks
-    #     except Exception as e:
-    #         logger.error(f"Failed to retrieve context chunks in document: {str(e)}")
-    #         raise
-
     async def retrieve_context_chunks_in_document_with_workflow(
         self,
         markdown_content: str,
@@ -205,26 +138,7 @@ class ContextChunksInDocumentService:
                 )
             )
             logger.info(f"Context chunks generated:{len(context_chunks)}")
-            # upsert validation
-            try:
-                print(f"deleting chunks: {file_key}")
-                self.delete_document_context_chunks(file_key)
-            except Exception as e:
-                logger.error(f"could not delete by source: {e}")
-            self.embeddings_manager.index_documents(context_chunks)
             return context_chunks
         except Exception as e:
-            logger.error("Error get_context_chunks_in_document")
-            raise e
-
-    def delete_document_context_chunks(self, file_key: str):
-        """
-        Delete the context chunks in a document.
-        """
-        try:
-            self.embeddings_manager.delete_documents_by_metadata_key(
-                self.metadata_source, file_key
-            )
-        except Exception as e:
-            logger.error(f"Error delete_document_context_chunks: {str(e)}")
+            logger.error(f"Error: {str(e)}")
             raise e

wizit_context_ingestor/application/interfaces.py

@@ -1,13 +1,19 @@
 """
 Application interfaces defining application layer contracts.
 """
+
 from abc import ABC, abstractmethod
-from ..domain.models import ParsedDocPage, ParsedDoc
-from typing import List, Union, Optional
-from langchain_core.documents import Document
+from typing import List, Optional, Union
+
+from langchain.indexes import IndexingResult, SQLRecordManager
 from langchain_aws import ChatBedrockConverse
+from langchain_core.documents import Document
 from langchain_google_vertexai import ChatVertexAI
 from langchain_google_vertexai.model_garden import ChatAnthropicVertex
+from langchain_postgres import PGVectorStore
+
+from ..domain.models import ParsedDoc, ParsedDocPage
+
 
 class TranscriptionService(ABC):
     """Interface for transcription services."""
@@ -17,6 +23,7 @@ class TranscriptionService(ABC):
         """Parse a document page."""
         pass
 
+
 class AiApplicationService(ABC):
     """Interface for AI application services."""
 
@@ -26,7 +33,9 @@ class AiApplicationService(ABC):
     #     pass
 
     @abstractmethod
-    def load_chat_model(self, **kwargs) -> Union[ChatVertexAI, ChatAnthropicVertex, ChatBedrockConverse]:
+    def load_chat_model(
+        self, **kwargs
+    ) -> Union[ChatVertexAI, ChatAnthropicVertex, ChatBedrockConverse]:
         """Load a chat model."""
         pass
 
@@ -40,7 +49,9 @@ class PersistenceService(ABC):
     """Interface for persistence services."""
 
     @abstractmethod
-    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."""
         pass
 
@@ -70,38 +81,56 @@ class EmbeddingsManager(ABC):
     @abstractmethod
     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"
     ):
         """Configure the vector store."""
         pass
 
+    # @abstractmethod
+    # async def init_vector_store(
+    #     self,
+    #     table_name: str = "tenant_embeddings",
+    #     content_column: str = "document",
+    #     metadata_json_column: str = "cmetadata",
+    #     id_column: str = "id",
+    # ):
+    #     """Initialize the vector store."""
+    #     pass
+
     @abstractmethod
-    def init_vector_store(
+    def retrieve_vector_store(
         self,
-        table_name: str = "langchain_pg_embedding",
-        content_column: str = "document",
-        metadata_json_column: str = "cmetadata",
-        id_column: str = "id",
-    ):
-        """Initialize the vector store."""
+    ) -> tuple[PGVectorStore, SQLRecordManager]:
+        """Retrieve the vector store."""
         pass
 
     @abstractmethod
-    def index_documents(self, documents: list[Document]):
+    def index_documents(
+        self,
+        docs: list[Document],
+    ) -> IndexingResult:
         """Index documents."""
         pass
 
     @abstractmethod
-    def get_documents_keys_by_source_id(self, source_id: str):
-        """Get documents keys by source ID."""
+    def search_records(
+        self,
+        query: str,
+    ) -> list[Document]:
+        """Search documents."""
         pass
 
     @abstractmethod
-    def delete_documents_by_source_id(self, source_id: str):
-        """Delete documents by source ID."""
+    def create_index(
+        self,
+    ):
         pass
+
+    # @abstractmethod
+    # def get_documents_keys_by_source_id(self, source_id: str):
+    #     """Get documents keys by source ID."""
+    #     pass
+
+    # @abstractmethod
+    # def delete_documents_by_source_id(self, source_id: str):
+    #     """Delete documents by source ID."""
+    #     pass

wizit_context_ingestor/application/kdb_service.py

@@ -0,0 +1,59 @@
+import logging
+
+from langchain.indexes import SQLRecordManager
+from langchain_core.documents import Document
+from langchain_postgres import PGVectorStore
+
+from .interfaces import (
+    EmbeddingsManager,
+    RagChunker,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class KdbService:
+    """
+    Service for chunking documents.
+    """
+
+    def __init__(
+        self,
+        embeddings_manager: EmbeddingsManager,
+    ):
+        """
+        Initialize the ChunkerService.
+        """
+        self.embeddings_manager = embeddings_manager
+        self._vector_store = None
+        self._records_manager = None
+
+    def configure_kdb(self):
+        try:
+            self.embeddings_manager.configure_vector_store()
+        except Exception as e:
+            raise Exception(f"Error configuring KDB: {e}")
+
+    def create_vector_store_hsnw_index(self):
+        try:
+            self.embeddings_manager.create_index()
+        except Exception as e:
+            logger.error(f"Error creating vector store index: {e}")
+            raise Exception(f"Error creating vector store index: {e}")
+
+    def search(self, query: str) -> list[Document]:
+        try:
+            records = []
+            records = self.embeddings_manager.search_records(query)
+            print(records)
+            return records
+        except Exception as e:
+            logger.error(f"Error indexing documents: {e}")
+            raise Exception(f"Error indexing documents: {e}")
+
+    def index_documents_in_vector_store(self, documents: list[Document]) -> None:
+        try:
+            self.embeddings_manager.index_documents(documents)
+        except Exception as e:
+            logger.error(f"Error indexing documents: {e}")
+            raise Exception(f"Error indexing documents: {e}")

wizit_context_ingestor/data/kdb.py

@@ -2,9 +2,12 @@ from enum import Enum
 from typing import Literal
 
 
-class KdbServices(Enum):
+class KdbServices(str, Enum):
     REDIS = "redis"
     CHROMA = "chroma"
+    PG = "pg"
 
 
-kdb_services = Literal[KdbServices.REDIS.value, KdbServices.CHROMA.value]
+kdb_services = Literal[
+    KdbServices.REDIS.value, KdbServices.CHROMA.value, KdbServices.PG.value
+]

wizit_context_ingestor/data/storage.py

@@ -2,7 +2,7 @@ from enum import Enum
 from typing import Literal
 
 
-class StorageServices(Enum):
+class StorageServices(str, Enum):
     S3 = "s3"
     LOCAL = "local"
 

wizit_context_ingestor/infra/rag/chroma_embeddings.py

@@ -46,26 +46,29 @@ class ChromaEmbeddingsManager(EmbeddingsManager):
             logger.error(f"Failed to initialize ChromaEmbeddingsManager: {str(e)}")
             raise
 
-    def configure_vector_store(
+    async def configure_vector_store(
         self,
         table_name: str = "",
         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",
     ):
         """Configure the vector store."""
         pass
 
-    def init_vector_store(
+    async def init_vector_store(
         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]):
+    async def index_documents(self, documents: list[Document]):
         """
         Add documents to the vector store with their embeddings.
 
@@ -85,7 +88,7 @@ class ChromaEmbeddingsManager(EmbeddingsManager):
         """
         try:
             logger.info(f"Indexing {len(documents)} documents in vector store")
-            self.chroma.add_documents(documents)
+            await self.chroma.aadd_documents(documents)
         except Exception as e:
             logger.error(f"Error indexing documents: {str(e)}")
             raise
@@ -110,12 +113,14 @@ 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):
+    async 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})
+            await self.chroma.adelete(where={metadata_key: metadata_value})
         except Exception as error:
             logger.error(
                 f"Error deleting documents by filter: {str(filter)}, error: {error} "

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

@@ -3,6 +3,7 @@
 # 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
@@ -58,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/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
 
 
@@ -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)"
                 )

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:
-        return True
-
     # 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):
+    if re.search(r"[^a-zA-Z0-9_.\-/\\]", file_name) is None:
         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.0b8.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.0b8
+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
 

{wizit_context_ingestor-0.3.0b8.dist-info → wizit_context_ingestor-0.4.0.dist-info}/RECORD

@@ -1,13 +1,14 @@
 wizit_context_ingestor/.DS_Store,sha256=c7hZ0C8v2hxprMlCgmvxXDl92phew3iSATJzE1yYTBs,6148
-wizit_context_ingestor/__init__.py,sha256=TSTm5qSpNNCz9ilKYkXRUxupvmWG2AHfv7RBWFw8T4c,107
+wizit_context_ingestor/__init__.py,sha256=aFOBTwoH2ZH4KvCzzdfQ0PH1yBV3aLdXXye-VtuXLRo,131
 wizit_context_ingestor/application/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-wizit_context_ingestor/application/context_chunk_service.py,sha256=LYRKBsY30IT2LczkgkYdPx7W3yycRy-0m7t3KKgq6Nw,9046
-wizit_context_ingestor/application/interfaces.py,sha256=W0qonE3t-S-zwAoKtDYc4oyW_GOILKVmrdy8LnC8MVI,3193
+wizit_context_ingestor/application/context_chunk_service.py,sha256=-efSbpLJS5968Qh6Ho_c6et4-g_L0gPgR-Z9URO_cS0,5504
+wizit_context_ingestor/application/interfaces.py,sha256=oG-a8JbGw-YT8Xtt0SDJ72FZwDNogyBzP93k_tBIC6k,3528
+wizit_context_ingestor/application/kdb_service.py,sha256=CYQ-2aLvUZHCgCwmuMeZsANhLbm54gtfpiKsg1a81qI,1757
 wizit_context_ingestor/application/transcription_service.py,sha256=FlUcMGyAotAO8MmT5UMlPMbgIWVQLg7YO6rJx9ANn7A,8567
 wizit_context_ingestor/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-wizit_context_ingestor/data/kdb.py,sha256=GCkXQmnk2JCXV_VJ-h0k55AOIX8qohzBJN2v-9D1dlU,194
+wizit_context_ingestor/data/kdb.py,sha256=lJ4F0ltVhtqC80QqHsm49a2FnSEI_fep61_Vfd65x34,241
 wizit_context_ingestor/data/prompts.py,sha256=bzgLdjINtXGQVTy4ZZktdcNItbtDQpM7maAQ2UBGdnY,14187
-wizit_context_ingestor/data/storage.py,sha256=aanXY1AV696cShHtDDhlJDhKPouZ1dq2lo_57yhTd20,198
+wizit_context_ingestor/data/storage.py,sha256=e29IH47MfN0qAuzrcZ-VlHw1sk7dT5QqYk6wO273nCk,203
 wizit_context_ingestor/domain/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wizit_context_ingestor/domain/models.py,sha256=DV83PArMyh-VoUqnVF_ohcgStsk549ixdYw98B8o2GI,381
 wizit_context_ingestor/domain/services.py,sha256=dg8UvYSjYsOMphrciZyGvuRriM8Qf08SstvO979XrFc,3344
@@ -16,20 +17,21 @@ wizit_context_ingestor/infra/aws_model.py,sha256=glIaewSdv6PDBXoCe6QgCUIzLCjtM7K
 wizit_context_ingestor/infra/persistence/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wizit_context_ingestor/infra/persistence/local_storage.py,sha256=GtPUvtn8XlgcqwjWmSm2998sgyYlwkF22HoB40ri7c0,2029
 wizit_context_ingestor/infra/persistence/s3_storage.py,sha256=bzlQteLPPGS_Gbh39RkxyoK8G-CEOQewMNPuzPule9k,4906
-wizit_context_ingestor/infra/rag/chroma_embeddings.py,sha256=fV6Ays8Vu4rzwp7kJiFx5HwepGeepk95Kzh_68Qjtkc,4298
-wizit_context_ingestor/infra/rag/pg_embeddings.py,sha256=D7onh27SvqYahYAsLy6DeyklxGyBFYH2DwV42fVCalQ,8157
+wizit_context_ingestor/infra/rag/chroma_embeddings.py,sha256=8VSXw38TX8nfGyJ5SXtbyqDVt3XI6gtEoie0g_v5Ax8,4517
+wizit_context_ingestor/infra/rag/pg_embeddings.py,sha256=f4fq75uot4JPaz8j7fmKxGWGhaLs5WcCtfLWBLvMA_M,7754
 wizit_context_ingestor/infra/rag/redis_embeddings.py,sha256=pCP_I1RLeIUTYMSHkZT6AjIOyHA9A47wyffrZBjiG0s,5107
-wizit_context_ingestor/infra/rag/semantic_chunks.py,sha256=miPG7gGYVIyQzTQMtQwDPknI9JlZroUP0m9E7ahfGfY,2449
+wizit_context_ingestor/infra/rag/semantic_chunks.py,sha256=Peylky5rsWERtu5psmet4SL_VYpnK5hsEIyE0Ua1wGE,2463
 wizit_context_ingestor/infra/secrets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wizit_context_ingestor/infra/secrets/aws_secrets_manager.py,sha256=vukil5sO9tQPTM74wUbyQqR8Z-z0ElyjeF2ns7rbVbQ,1249
 wizit_context_ingestor/infra/vertex_model.py,sha256=6L2C4qH7PSVjdOSzIEZlFtUwu1pgQVXtQBIU5isn644,7582
-wizit_context_ingestor/main.py,sha256=vP9311d9TcXtoxmOE1-4jAzKjrNU5ZaxiJiG2d9IV1w,11687
+wizit_context_ingestor/main.py,sha256=2gJ7B39aS2U742cJQWjVGHbQ-NslYFdGQXzljGd83WA,11685
+wizit_context_ingestor/main_chunks.py,sha256=07zT76VB_RzCXV3toJgTQBPCmiVud860vVzVkMLwDuA,6784
 wizit_context_ingestor/services/.DS_Store,sha256=1lFlJ5EFymdzGAUAaI30vcaaLHt3F1LwpG7xILf9jsM,6148
 wizit_context_ingestor/services/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wizit_context_ingestor/services/chunks.py,sha256=tQQsdsOscZWzqVY5WxVxr3ii62FOJ3nMARaJJz6CvjQ,2011
 wizit_context_ingestor/services/parse_doc.py,sha256=3CyZoGbiUfxbs0SXUWXjQevtusSzTBgvUVeNNSdxJLE,4491
 wizit_context_ingestor/services/pg_embeddings_manager.py,sha256=n1HOmu_Z_Z71H-rVAyJS3FdPKbBckm5W8_XethY8nuM,4998
-wizit_context_ingestor/utils/file_utils.py,sha256=Wuua14LivrfL8oBP-j3ZtSqc7uq4rrQzT-T-p7pxEpM,414
+wizit_context_ingestor/utils/file_utils.py,sha256=BmNN71exw75IKWvdDvyvo4QlGpl5W-nTCgimDnmD4Eo,501
 wizit_context_ingestor/workflows/context_nodes.py,sha256=3qlFcxPUmehx04mQHpmouneKq--To8rwSDHCRFyWICo,3168
 wizit_context_ingestor/workflows/context_state.py,sha256=4MTIUjK-F2pWvIldovWZhMAqqCOpViKbvitJzETkSkY,324
 wizit_context_ingestor/workflows/context_tools.py,sha256=E9VTL3AC0MwSIuc1e-juZK7XCxnZfFv0-KpHfR2CNH4,2764
@@ -39,6 +41,6 @@ wizit_context_ingestor/workflows/transcription_schemas.py,sha256=CQCl7LXD5voxhJO
 wizit_context_ingestor/workflows/transcription_state.py,sha256=2Z_t2aZFEH_nAjdEO6RFBEmi_fwvr9cV0aLS1eIxiCQ,590
 wizit_context_ingestor/workflows/transcription_tools.py,sha256=FtIfWFITn8_Rr5SEobCeR55aJGZoHRMgF2UxRT5vJ-E,1373
 wizit_context_ingestor/workflows/transcription_workflow.py,sha256=77cLsYGdv01Py2GaKYpACuifPeSxH7tkVodvLv97sdg,1621
-wizit_context_ingestor-0.3.0b8.dist-info/WHEEL,sha256=eh7sammvW2TypMMMGKgsM83HyA_3qQ5Lgg3ynoecH3M,79
-wizit_context_ingestor-0.3.0b8.dist-info/METADATA,sha256=7pmXei8lCU0BAPl2m7T9k7nlZFAEOAANgndG3OPbIHY,3768
-wizit_context_ingestor-0.3.0b8.dist-info/RECORD,,
+wizit_context_ingestor-0.4.0.dist-info/WHEEL,sha256=eh7sammvW2TypMMMGKgsM83HyA_3qQ5Lgg3ynoecH3M,79
+wizit_context_ingestor-0.4.0.dist-info/METADATA,sha256=hZxwWZSUb5hQXEoWZ5ulaI0fmPZYSyDmgOWHxizp6oQ,3890
+wizit_context_ingestor-0.4.0.dist-info/RECORD,,