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

wizit_context_ingestor/__init__.py

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

wizit_context_ingestor/application/context_chunk_service.py

@@ -2,8 +2,15 @@ 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
+from .interfaces import (
+    AiApplicationService,
+    PersistenceService,
+    RagChunker,
+    EmbeddingsManager,
+)
 import logging
 
 
@@ -21,7 +28,7 @@ class ContextChunksInDocumentService:
         persistence_service: PersistenceService,
         rag_chunker: RagChunker,
         embeddings_manager: EmbeddingsManager,
-        target_language: str = 'es'
+        target_language: str = "es",
     ):
         """
         Initialize the ChunkerService.
@@ -33,48 +40,144 @@ class ContextChunksInDocumentService:
         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(self, markdown_content: str, chunk: Document, chunk_metadata: Optional[Dict[str, Any]] = None) -> Document:
+    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:
-            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()
+            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,
+                    }
+                },
             )
-            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
+            # 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]:
+    # 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_chunks = list(map(
-                lambda chunk: self._retrieve_context_chunk_in_document(markdown_content, chunk, chunks_metadata),
-                chunks
-            ))
+            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)}")
@@ -85,18 +188,27 @@ class ContextChunksInDocumentService:
         Get the context chunks in a document.
         """
         try:
-            markdown_content = self.persistence_service.load_markdown_file_content(file_key)
+            markdown_content = self.persistence_service.load_markdown_file_content(
+                file_key
+            )
             langchain_rag_document = Document(
+                id=file_key,
                 page_content=markdown_content,
-                metadata={
-                    "source": file_key
-                }
+                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(markdown_content, chunks, file_tags)
+            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:
@@ -108,7 +220,9 @@ class ContextChunksInDocumentService:
         Delete the context chunks in a document.
         """
         try:
-            self.embeddings_manager.delete_documents_by_source_id(file_key)
+            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/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/application/transcription_service.py

@@ -1,76 +1,153 @@
 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.
+    Service for transcribing documents.
     """
 
     def __init__(
         self,
         ai_application_service: AiApplicationService,
         persistence_service: PersistenceService,
-        target_language: str = 'es'
+        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}"
+    # 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}"
+                                },
                             }
-                        },
-                        {
-                            "type": "text",
-                            "text": f"Transcribe the document, ensure all content transcribed is using '{self.target_language}' language"
-                        }]
+                        ]
                     ),
-                ]).partial(
-                    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
+                ]
+            },
+            {
+                "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]:
         """
@@ -81,15 +158,21 @@ class TranscriptionService:
         document_pages = parse_doc_model_service.parse_document_to_base64()
         parsed_pages = []
         for page in document_pages:
-            page = self.parse_doc_page(page)
+            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]] = {}):
+    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)
+        self.persistence_service.save_parsed_document(
+            file_key, parsed_document, file_tags
+        )

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]