wizit-context-ingestor 0.2.5b2__tar.gz → 0.3.0b1__tar.gz

{wizit_context_ingestor-0.2.5b2 → wizit_context_ingestor-0.3.0b1}/PKG-INFO

@@ -1,13 +1,15 @@
 Metadata-Version: 2.3
 Name: wizit-context-ingestor
-Version: 0.2.5b2
+Version: 0.3.0b1
 Summary: Contextual Rag with Cloud Solutions
 Requires-Dist: anthropic[vertex]>=0.66.0
 Requires-Dist: boto3>=1.40.23
 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-redis>=0.2.3
+Requires-Dist: langgraph>=0.6.8
 Requires-Dist: pillow>=11.3.0
 Requires-Dist: pymupdf>=1.26.4
 Requires-Python: >=3.12
@@ -137,6 +139,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.5b2 → wizit_context_ingestor-0.3.0b1}/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.5b2 → wizit_context_ingestor-0.3.0b1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "wizit_context_ingestor"
-version = "0.2.5-beta-2"
+version = "0.3.0-beta-1"
 description = "Contextual Rag with Cloud Solutions"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -8,9 +8,11 @@ dependencies = [
     "anthropic[vertex]>=0.66.0",
     "boto3>=1.40.23",
     "langchain-aws>=0.2.31",
+    "langchain-chroma>=0.2.6",
     "langchain-experimental>=0.3.4",
     "langchain-google-vertexai>=2.0.28",
     "langchain-redis>=0.2.3",
+    "langgraph>=0.6.8",
     "pillow>=11.3.0",
     "pymupdf>=1.26.4",
 ]

wizit_context_ingestor-0.3.0b1/src/wizit_context_ingestor/__init__.py

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

wizit_context_ingestor-0.3.0b1/src/wizit_context_ingestor/application/context_chunk_service.py

@@ -0,0 +1,228 @@
+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,
+    PersistenceService,
+    RagChunker,
+    EmbeddingsManager,
+)
+import logging
+
+
+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"
+
+    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 = workflow.invoke(
+                {
+                    "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:{result['context']}, Content:{chunk.page_content}"
+            # )
+            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
+
+    # 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
+
+    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 = list(
+                map(
+                    lambda chunk: self._retrieve_context_chunk_in_document_with_workflow(
+                        compiled_context_workflow,
+                        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
+
+    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 = self.retrieve_context_chunks_in_document_with_workflow(
+                markdown_content, chunks, file_tags
+            )
+            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)}")
+            raise e

{wizit_context_ingestor-0.2.5b2 → wizit_context_ingestor-0.3.0b1}/src/wizit_context_ingestor/application/interfaces.py

@@ -92,7 +92,7 @@ class EmbeddingsManager(ABC):
         pass
 
     @abstractmethod
-    def index_documents(self, documents: List[Document]):
+    def index_documents(self, documents: list[Document]):
         """Index documents."""
         pass
 

wizit_context_ingestor-0.3.0b1/src/wizit_context_ingestor/application/transcription_service.py

@@ -0,0 +1,178 @@
+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: int = 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
+            or transcription_accuracy_threshold > 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()
+
+    # 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
+
+    def parse_doc_page_with_workflow(self, document: ParsedDocPage) -> ParsedDocPage:
+        """Transcribe an image to text using an agent.
+        Args:
+            document: The document with the image to transcribe
+        Returns:
+            Processed text
+        """
+        transcription_workflow = TranscriptionWorkflow(
+            self.chat_model, self.transcription_additional_instructions
+        )
+        compiled_transcription_workflow = transcription_workflow.gen_workflow()
+        compiled_transcription_workflow = compiled_transcription_workflow.compile()
+        result = compiled_transcription_workflow.invoke(
+            {
+                "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 result["transcription"]:
+            document.page_text = result["transcription"]
+        else:
+            raise ValueError("No transcription found")
+        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
+
+    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.3.0b1/src/wizit_context_ingestor/data/kdb.py

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