wizit-context-ingestor 0.2.5b3__tar.gz → 0.4.1__tar.gz

{wizit_context_ingestor-0.2.5b3 → wizit_context_ingestor-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: wizit-context-ingestor
-Version: 0.2.5b3
+Version: 0.4.1
 Summary: Contextual Rag with Cloud Solutions
 Requires-Dist: anthropic[vertex]>=0.66.0
 Requires-Dist: boto3>=1.40.23
@@ -8,9 +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
 
@@ -138,6 +142,13 @@ Finally
     poetry publish -r tbbcmegaingestor
 ```
 
+# USAGE
+
+## For transcriptions
+
+----- TODO ---
+You can provide number of retries and a transcription quality threshold
+
 ## License
 
 This project is licensed under the Apache License - see the LICENSE file for details.

{wizit_context_ingestor-0.2.5b3 → wizit_context_ingestor-0.4.1}/README.md

@@ -122,6 +122,13 @@ Finally
     poetry publish -r tbbcmegaingestor
 ```
 
+# USAGE
+
+## For transcriptions
+
+----- TODO ---
+You can provide number of retries and a transcription quality threshold
+
 ## License
 
 This project is licensed under the Apache License - see the LICENSE file for details.

{wizit_context_ingestor-0.2.5b3 → wizit_context_ingestor-0.4.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "wizit_context_ingestor"
-version = "0.2.5-beta-3"
+version = "0.4.1"
 description = "Contextual Rag with Cloud Solutions"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -11,9 +11,13 @@ dependencies = [
     "langchain-chroma>=0.2.6",
     "langchain-experimental>=0.3.4",
     "langchain-google-vertexai>=2.0.28",
+    "langchain-postgres>=0.0.16",
     "langchain-redis>=0.2.3",
+    "langgraph>=0.6.8",
     "pillow>=11.3.0",
+    "psycopg2-binary>=2.9.11",
     "pymupdf>=1.26.4",
+    "sqlalchemy[asyncio]>=2.0.43",
 ]
 
 [dependency-groups]

wizit_context_ingestor-0.4.1/src/wizit_context_ingestor/__init__.py

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

wizit_context_ingestor-0.4.1/src/wizit_context_ingestor/application/context_chunk_service.py

@@ -0,0 +1,145 @@
+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 ..data.prompts import CONTEXT_CHUNKS_IN_DOCUMENT_SYSTEM_PROMPT, ContextChunk
+from ..workflows.context_workflow import ContextWorkflow
+from .interfaces import (
+    AiApplicationService,
+    EmbeddingsManager,
+    PersistenceService,
+    RagChunker,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ContextChunksInDocumentService:
+    """
+    Service for chunking documents.
+    """
+
+    def __init__(
+        self,
+        ai_application_service: AiApplicationService,
+        persistence_service: PersistenceService,
+        rag_chunker: RagChunker,
+        embeddings_manager: EmbeddingsManager,
+        target_language: str = "es",
+    ):
+        """
+        Initialize the ChunkerService.
+        """
+        self.ai_application_service = ai_application_service
+        self.persistence_service = persistence_service
+        self.rag_chunker = rag_chunker
+        self.embeddings_manager = embeddings_manager
+        self.target_language = target_language
+        # self.embeddings_manager.init_vector_store()
+        self.chat_model = self.ai_application_service.load_chat_model()
+        # TODO
+        self.context_additional_instructions = ""
+        self.metadata_source = "source"
+
+    async def _retrieve_context_chunk_in_document_with_workflow(
+        self,
+        workflow,
+        markdown_content: str,
+        chunk: Document,
+        chunk_metadata: Optional[Dict[str, Any]] = None,
+    ) -> Document:
+        """Retrieve context chunks in document."""
+        try:
+            result = await workflow.ainvoke(
+                {
+                    "messages": [
+                        HumanMessage(
+                            content=[
+                                {
+                                    "type": "text",
+                                    "text": f"Retrieve a complete context for the following chunk: <chunk>{chunk.page_content}</chunk>,  ensure all content chunks are generated with the same document's language.",
+                                },
+                            ]
+                        )
+                    ],
+                    "document_content": markdown_content,
+                },
+                {
+                    "configurable": {
+                        "transcription_accuracy_threshold": 0.95,
+                        "max_transcription_retries": 2,
+                    }
+                },
+            )
+            chunk.page_content = f"<context>\n{result['context']}\n</context>\n <content>\n{chunk.page_content}\n</content>"
+            # INFO: prevent context in metadata because it's already included in the chunk content, also generates issues when text is long
+            # chunk.metadata["context"] = result["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
+
+    async def retrieve_context_chunks_in_document_with_workflow(
+        self,
+        markdown_content: str,
+        chunks: List[Document],
+        chunks_metadata: Optional[Dict[str, Any]] = None,
+    ) -> List[Document]:
+        """Retrieve context chunks in document."""
+        try:
+            context_workflow = ContextWorkflow(
+                self.chat_model, self.context_additional_instructions
+            )
+            compiled_context_workflow = context_workflow.gen_workflow()
+            compiled_context_workflow = compiled_context_workflow.compile()
+            context_chunks_workflow_invocations = list(
+                map(
+                    lambda chunk: self._retrieve_context_chunk_in_document_with_workflow(
+                        compiled_context_workflow,
+                        markdown_content,
+                        chunk,
+                        chunks_metadata,
+                    ),
+                    chunks,
+                )
+            )
+            context_chunks = await asyncio.gather(*context_chunks_workflow_invocations)
+            return context_chunks
+        except Exception as e:
+            logger.error(f"Failed to retrieve context chunks in document: {str(e)}")
+            raise
+
+    async def get_context_chunks_in_document(self, file_key: str, file_tags: dict = {}):
+        """
+        Get the context chunks in a document.
+        """
+        try:
+            markdown_content = self.persistence_service.load_markdown_file_content(
+                file_key
+            )
+            langchain_rag_document = Document(
+                id=file_key,
+                page_content=markdown_content,
+                metadata={self.metadata_source: file_key},
+            )
+            logger.info(f"Document loaded:{file_key}")
+            chunks = self.rag_chunker.gen_chunks_for_document(langchain_rag_document)
+            logger.info(f"Chunks generated:{len(chunks)}")
+            context_chunks = (
+                await self.retrieve_context_chunks_in_document_with_workflow(
+                    markdown_content, chunks, file_tags
+                )
+            )
+            logger.info(f"Context chunks generated:{len(context_chunks)}")
+            return context_chunks
+        except Exception as e:
+            logger.error(f"Error: {str(e)}")
+            raise e

{wizit_context_ingestor-0.2.5b3 → wizit_context_ingestor-0.4.1}/src/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-0.4.1/src/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-0.4.1/src/wizit_context_ingestor/application/transcription_service.py

@@ -0,0 +1,209 @@
+import asyncio
+from typing import Tuple, List, Dict, Optional
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers.pydantic import PydanticOutputParser
+from langchain_core.messages import HumanMessage
+from logging import getLogger
+from ..data.prompts import IMAGE_TRANSCRIPTION_SYSTEM_PROMPT, Transcription
+from ..domain.models import ParsedDoc, ParsedDocPage
+from ..domain.services import ParseDocModelService
+from .interfaces import AiApplicationService, PersistenceService
+from ..workflows.transcription_workflow import TranscriptionWorkflow
+
+logger = getLogger(__name__)
+
+
+class TranscriptionService:
+    """
+    Service for transcribing documents.
+    """
+
+    def __init__(
+        self,
+        ai_application_service: AiApplicationService,
+        persistence_service: PersistenceService,
+        target_language: str = "es",
+        transcription_additional_instructions: str = "",
+        transcription_accuracy_threshold: float = 0.90,
+        max_transcription_retries: int = 2,
+    ):
+        self.ai_application_service = ai_application_service
+        self.persistence_service = persistence_service
+        self.target_language = target_language
+        if (
+            transcription_accuracy_threshold < 0.0
+            or transcription_accuracy_threshold > 0.95
+        ):
+            raise ValueError(
+                "transcription_accuracy_threshold must be between 0 and 95"
+            )
+        if max_transcription_retries < 1 or max_transcription_retries > 3:
+            raise ValueError(
+                "max_transcription_retries must be between 1 and 3 to prevent token exhaustion"
+            )
+        self.transcription_accuracy_threshold = transcription_accuracy_threshold
+        self.max_transcription_retries = max_transcription_retries
+        self.transcription_additional_instructions = (
+            transcription_additional_instructions
+        )
+        self.chat_model = self.ai_application_service.load_chat_model()
+        self.transcription_workflow = TranscriptionWorkflow(
+            self.chat_model, self.transcription_additional_instructions
+        )
+        self.compiled_transcription_workflow = (
+            self.transcription_workflow.gen_workflow()
+        )
+        self.compiled_transcription_workflow = (
+            self.compiled_transcription_workflow.compile()
+        )
+
+    # def parse_doc_page(self, document: ParsedDocPage) -> ParsedDocPage:
+    #     """Transcribe an image to text.
+    #     Args:
+    #         document: The document with the image to transcribe
+    #     Returns:
+    #         Processed text
+    #     """
+    #     try:
+    #         # Create the prompt template with image
+    #         transcription_output_parser = PydanticOutputParser(
+    #             pydantic_object=Transcription
+    #         )
+    #         prompt = ChatPromptTemplate.from_messages(
+    #             [
+    #                 ("system", IMAGE_TRANSCRIPTION_SYSTEM_PROMPT),
+    #                 (
+    #                     "human",
+    #                     [
+    #                         {
+    #                             "type": "image",
+    #                             "image_url": {
+    #                                 "url": f"data:image/png;base64,{document.page_base64}"
+    #                             },
+    #                         },
+    #                         {
+    #                             "type": "text",
+    #                             "text": "Transcribe the document, ensure all content transcribed accurately",
+    #                         },
+    #                     ],
+    #                 ),
+    #             ]
+    #         ).partial(
+    #             transcription_additional_instructions=self.transcription_additional_instructions,
+    #             format_instructions=transcription_output_parser.get_format_instructions(),
+    #         )
+    #         model_with_structured_output = self.chat_model.with_structured_output(
+    #             Transcription
+    #         )
+    #         # Create the chain
+    #         chain = prompt | model_with_structured_output
+    #         # Process the image
+    #         chain = chain.with_retry(
+    #             stop_after_attempt=3, exponential_jitter_params={"initial": 60}
+    #         )
+    #         result = chain.invoke({})
+    #         if result.transcription:
+    #             document.page_text = result.transcription
+    #         else:
+    #             raise ValueError("No transcription found")
+    #         return document
+    #     except Exception as e:
+    #         logger.error(f"Failed to parse document page: {str(e)}")
+    #         raise
+
+    async def parse_doc_page_with_workflow(
+        self, document: ParsedDocPage, retries: int = 0
+    ) -> ParsedDocPage:
+        """Transcribe an image to text using an agent.
+        Args:
+            document: The document with the image to transcribe
+        Returns:
+            Processed text
+        """
+        if retries > 1:
+            logger.info("Max retries exceeded")
+            return document
+        result = await self.compiled_transcription_workflow.ainvoke(
+            {
+                "messages": [
+                    HumanMessage(
+                        content=[
+                            {
+                                "type": "text",
+                                "text": "Transcribe the document, ensure all content transcribed accurately. transcription must be in the same language of source document.",
+                            },
+                        ]
+                    ),
+                    HumanMessage(
+                        content=[
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": f"data:image/png;base64,{document.page_base64}"
+                                },
+                            }
+                        ]
+                    ),
+                ]
+            },
+            {
+                "configurable": {
+                    "transcription_accuracy_threshold": self.transcription_accuracy_threshold,
+                    "max_transcription_retries": self.max_transcription_retries,
+                }
+            },
+        )
+        if "transcription" in result:
+            document.page_text = result["transcription"]
+        else:
+            return await self.parse_doc_page_with_workflow(
+                document, retries=retries + 1
+            )
+        return document
+
+    # def process_document(self, file_key: str) -> Tuple[List[ParsedDocPage], ParsedDoc]:
+    #     """
+    #     Process a document by parsing it and returning the parsed content.
+    #     """
+    #     raw_file_path = self.persistence_service.retrieve_raw_file(file_key)
+    #     parse_doc_model_service = ParseDocModelService(raw_file_path)
+    #     document_pages = parse_doc_model_service.parse_document_to_base64()
+    #     parsed_pages = []
+    #     for page in document_pages:
+    #         page = self.parse_doc_page_with_workflow(page)
+    #         parsed_pages.append(page)
+    #     logger.info(f"Parsed {len(parsed_pages)} pages")
+    #     parsed_document = parse_doc_model_service.create_md_content(parsed_pages)
+    #     return parsed_pages, parsed_document
+
+    async def process_document(
+        self, file_key: str
+    ) -> Tuple[List[ParsedDocPage], ParsedDoc]:
+        """
+        Process a document by parsing it and returning the parsed content.
+        """
+        raw_file_path = self.persistence_service.retrieve_raw_file(file_key)
+        parse_doc_model_service = ParseDocModelService(raw_file_path)
+        document_pages = parse_doc_model_service.parse_document_to_base64()
+        parse_pages_workflow_tasks = []
+        parsed_pages = []
+        for page in document_pages:
+            parse_pages_workflow_tasks.append(self.parse_doc_page_with_workflow(page))
+        # here
+        parsed_pages = await asyncio.gather(*parse_pages_workflow_tasks)
+        logger.info(f"Parsed {len(parsed_pages)} pages")
+        parsed_document = parse_doc_model_service.create_md_content(parsed_pages)
+        return parsed_pages, parsed_document
+
+    def save_parsed_document(
+        self,
+        file_key: str,
+        parsed_document: ParsedDoc,
+        file_tags: Optional[Dict[str, str]] = {},
+    ):
+        """
+        Save the parsed document to a file.
+        """
+        self.persistence_service.save_parsed_document(
+            file_key, parsed_document, file_tags
+        )

wizit_context_ingestor-0.4.1/src/wizit_context_ingestor/data/kdb.py

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