From 0bdee3025b35b93b20ca83bdeee33583b744e778 Mon Sep 17 00:00:00 2001 From: camilleAND Date: Wed, 14 Jan 2026 16:37:53 +0100 Subject: [PATCH] add way to split xlsx for storage --- .../albert_rag_backend.py | 256 +++++++++++++++++- src/backend/chat/clients/pydantic_ai.py | 29 +- src/backend/pyproject.toml | 4 + 3 files changed, 286 insertions(+), 3 deletions(-) diff --git a/src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py b/src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py index 8af560b..d5aa549 100644 --- a/src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py +++ b/src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py @@ -18,6 +18,169 @@ from chat.agent_rag.document_rag_backends.base_rag_backend import BaseRagBackend logger = logging.getLogger(__name__) +# Albert API token limit for document vectorization +# We use a conservative chunk size to stay well under the limit +ALBERT_MAX_TOKENS = 8192 +ALBERT_CHUNK_SIZE_TOKENS = 5000 # More conservative chunk size with larger safety margin +# Approximate tokens: ~3 characters per token (more conservative estimate for Markdown/Excel) +# Markdown and Excel content often have more tokens per character due to formatting +ALBERT_CHUNK_SIZE_CHARS = ALBERT_CHUNK_SIZE_TOKENS * 3 + + +def _estimate_tokens(content: str) -> int: + """ + Estimate the number of tokens in a text string. + + Uses a conservative approximation: ~3 characters per token. + This is more conservative than 4 chars/token to account for: + - Markdown formatting (headers, lists, tables) + - Excel content with special characters + - Whitespace and punctuation + + Args: + content (str): The text content to estimate. + + Returns: + int: Estimated number of tokens. + """ + return len(content) // 3 + + +def _chunk_content(content: str, max_chars: int = ALBERT_CHUNK_SIZE_CHARS) -> List[str]: + """ + Split content into chunks that fit within Albert's token limit. + + Attempts to split at paragraph boundaries (double newlines) when possible, + otherwise splits at line boundaries, and finally at character boundaries. + Validates that each chunk is under the token limit after splitting. + + Args: + content (str): The content to chunk. + max_chars (int): Maximum characters per chunk (default: ALBERT_CHUNK_SIZE_CHARS). + + Returns: + list[str]: List of content chunks, each under the token limit. + """ + # First check if content fits in one chunk + estimated_tokens = _estimate_tokens(content) + if estimated_tokens <= ALBERT_CHUNK_SIZE_TOKENS: + return [content] + + chunks = [] + remaining = content + + while len(remaining) > 0: + # Check if remaining content fits in one chunk + remaining_tokens = _estimate_tokens(remaining) + if remaining_tokens <= ALBERT_CHUNK_SIZE_TOKENS: + if remaining.strip(): + chunks.append(remaining.strip()) + break + + # Need to split - find the best split point + # Start with max_chars but may need to reduce if token estimate is too high + search_limit = max_chars + + # Try to find a split point that keeps us under token limit + # Reduce search limit if needed to ensure token limit is respected + while search_limit > 100: # Minimum chunk size + # Try to split at paragraph boundary (double newline) + split_pos = remaining.rfind("\n\n", 0, search_limit) + if split_pos == -1: + # Try to split at single newline + split_pos = remaining.rfind("\n", 0, search_limit) + if split_pos == -1: + # Force split at character boundary + split_pos = search_limit + + # Validate that this chunk is under token limit + chunk_candidate = remaining[:split_pos].strip() + if chunk_candidate: + chunk_tokens = _estimate_tokens(chunk_candidate) + if chunk_tokens <= ALBERT_CHUNK_SIZE_TOKENS: + chunks.append(chunk_candidate) + remaining = remaining[split_pos:].lstrip() + break + + # Chunk too large, reduce search limit and try again + search_limit = int(search_limit * 0.8) # Reduce by 20% + else: + # Fallback: force split at a safe size + # This should rarely happen, but ensures we don't get stuck + safe_size = min(max_chars, len(remaining)) + chunk = remaining[:safe_size].strip() + if chunk: + chunks.append(chunk) + remaining = remaining[safe_size:].lstrip() + + # Validate all chunks are under limit and split further if needed + validated_chunks = [] + for chunk_item in chunks: + chunk_tokens = _estimate_tokens(chunk_item) + if chunk_tokens > ALBERT_MAX_TOKENS: + logger.warning( + "Chunk still exceeds token limit (%d tokens, max: %d), forcing split further", + chunk_tokens, + ALBERT_MAX_TOKENS, + ) + # Force split this chunk further using a more conservative size + # Use a size that ensures we stay well under the token limit + # Target: ~5000 tokens max per chunk (conservative) + max_safe_chars = ALBERT_CHUNK_SIZE_TOKENS * 3 # 6000 * 3 = 18000 chars for ~5000 tokens + remaining_chunk = chunk_item + while len(remaining_chunk) > 0: + remaining_tokens = _estimate_tokens(remaining_chunk) + if remaining_tokens <= ALBERT_CHUNK_SIZE_TOKENS: + if remaining_chunk.strip(): + validated_chunks.append(remaining_chunk.strip()) + break + + # Find a safe split point + split_pos = min(max_safe_chars, len(remaining_chunk)) + # Try to split at a line boundary if possible + line_split = remaining_chunk.rfind("\n", 0, split_pos) + if line_split > max_safe_chars * 0.5: # Only use if it's not too small + split_pos = line_split + + sub_chunk = remaining_chunk[:split_pos].strip() + if sub_chunk: + sub_tokens = _estimate_tokens(sub_chunk) + # Double-check this sub-chunk is safe + if sub_tokens > ALBERT_MAX_TOKENS: + # Still too large, use even smaller size + logger.warning( + "Sub-chunk still too large (%d tokens), using smaller split", + sub_tokens, + ) + split_pos = ALBERT_CHUNK_SIZE_TOKENS * 2 # 12000 chars for ~3000 tokens + sub_chunk = remaining_chunk[:split_pos].strip() + validated_chunks.append(sub_chunk) + remaining_chunk = remaining_chunk[split_pos:].lstrip() + else: + validated_chunks.append(chunk_item) + + # Final validation - ensure NO chunk exceeds the limit + final_chunks = [] + for chunk in validated_chunks: + chunk_tokens = _estimate_tokens(chunk) + if chunk_tokens > ALBERT_MAX_TOKENS: + logger.error( + "CRITICAL: Chunk still exceeds limit after all splitting attempts: %d tokens", + chunk_tokens, + ) + # Emergency split: use very conservative size + emergency_size = ALBERT_CHUNK_SIZE_TOKENS * 2 # 12000 chars + remaining = chunk + while len(remaining) > 0: + emergency_chunk = remaining[:emergency_size].strip() + if emergency_chunk: + final_chunks.append(emergency_chunk) + remaining = remaining[emergency_size:].lstrip() + else: + final_chunks.append(chunk) + + return final_chunks + class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attributes """ @@ -170,7 +333,42 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att """ Store the document content in the Albert collection. This method should handle the logic to send the document content to the Albert API. + + If the document is too large (exceeds Albert's token limit), it will be automatically + split into multiple chunks and stored as separate documents. + Args: + name (str): The name of the document. + content (str): The content of the document in Markdown format. + """ + # Check if content needs to be chunked + estimated_tokens = _estimate_tokens(content) + + if estimated_tokens > ALBERT_MAX_TOKENS: + logger.info( + "Document '%s' is too large (%d estimated tokens, limit: %d). " + "Splitting into chunks.", + name, + estimated_tokens, + ALBERT_MAX_TOKENS, + ) + chunks = _chunk_content(content) + logger.info("Split document '%s' into %d chunks", name, len(chunks)) + + # Store each chunk as a separate document + for i, chunk in enumerate(chunks, start=1): + chunk_name = f"{name}_part_{i}" if len(chunks) > 1 else name + self._store_single_document(chunk_name, chunk) + else: + # Document fits within limit, store as-is + self._store_single_document(name, content) + + def _store_single_document(self, name: str, content: str) -> None: + """ + Store a single document chunk in the Albert collection. + + Internal method that performs the actual API call to store one document. + Args: name (str): The name of the document. content (str): The content of the document in Markdown format. @@ -185,14 +383,68 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att }, timeout=settings.ALBERT_API_TIMEOUT, ) - logger.debug(response.json()) + logger.debug("Stored document '%s': %s", name, response.json()) response.raise_for_status() async def astore_document(self, name: str, content: str) -> None: """ Store the document content in the Albert collection. This method should handle the logic to send the document content to the Albert API. + + If the document is too large (exceeds Albert's token limit), it will be automatically + split into multiple chunks and stored as separate documents. + Args: + name (str): The name of the document. + content (str): The content of the document in Markdown format. + """ + # Check if content needs to be chunked + estimated_tokens = _estimate_tokens(content) + + if estimated_tokens > ALBERT_MAX_TOKENS: + logger.info( + "Document '%s' is too large (%d estimated tokens, limit: %d). " + "Splitting into chunks.", + name, + estimated_tokens, + ALBERT_MAX_TOKENS, + ) + chunks = _chunk_content(content) + logger.info("Split document '%s' into %d chunks", name, len(chunks)) + + # Validate chunks before storing + for i, chunk in enumerate(chunks, start=1): + chunk_tokens = _estimate_tokens(chunk) + logger.debug( + "Chunk %d/%d: %d chars, ~%d tokens", + i, + len(chunks), + len(chunk), + chunk_tokens, + ) + if chunk_tokens > ALBERT_MAX_TOKENS: + logger.error( + "Chunk %d/%d still exceeds token limit: %d tokens (max: %d)", + i, + len(chunks), + chunk_tokens, + ALBERT_MAX_TOKENS, + ) + + # Store each chunk as a separate document + for i, chunk in enumerate(chunks, start=1): + chunk_name = f"{name}_part_{i}" if len(chunks) > 1 else name + await self._astore_single_document(chunk_name, chunk) + else: + # Document fits within limit, store as-is + await self._astore_single_document(name, content) + + async def _astore_single_document(self, name: str, content: str) -> None: + """ + Store a single document chunk in the Albert collection. + + Internal method that performs the actual API call to store one document. + Args: name (str): The name of the document. content (str): The content of the document in Markdown format. @@ -210,7 +462,7 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att }, timeout=settings.ALBERT_API_TIMEOUT, ) - logger.debug(response.json()) + logger.debug("Stored document '%s': %s", name, response.json()) response.raise_for_status() def search(self, query, results_count: int = 4) -> RAGWebResults: diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index f39453d..9ba149b 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -72,6 +72,7 @@ from chat.clients.pydantic_ui_message_converter import ( ui_message_to_user_content, ) from chat.mcp_servers import get_mcp_servers +from chat.tools.data_analysis import add_data_analysis_tool from chat.tools.document_generic_search_rag import add_document_rag_search_tool_from_setting from chat.tools.document_search_rag import add_document_rag_search_tool from chat.tools.document_summarize import document_summarize @@ -151,6 +152,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes deps_type=ContextDeps, ) add_document_rag_search_tool_from_setting(self.conversation_agent, self.user) + add_data_analysis_tool(self.conversation_agent) @property def _stop_cache_key(self): @@ -289,7 +291,24 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes content=document.data, ) - if not document.media_type.startswith("text/"): + # Don't convert tabular files (CSV, Excel) to Markdown - keep originals for data_analysis tool + # Tabular files are already text-based or can be used directly + is_tabular_file = ( + document.media_type in [ + "text/csv", + "application/csv", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + "application/vnd.ms-excel", + "application/excel", + ] + or any( + document.identifier.lower().endswith(ext) + for ext in [".csv", ".xlsx", ".xls", ".xlsm", ".xlsb"] + ) + ) + + # Only convert non-text files that are not tabular files + if not document.media_type.startswith("text/") and not is_tabular_file: md_attachment = await models.ChatConversationAttachment.objects.acreate( conversation=self.conversation, uploaded_by=self.user, @@ -487,6 +506,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes .aexists() ) + document_urls = [] if not conversation_has_documents and not has_not_pdf_docs: # No documents to process @@ -521,6 +541,13 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes async def summarize(ctx: RunContext, *args, **kwargs) -> ToolReturn: """Wrap the document_summarize tool to provide context and add the tool.""" return await document_summarize(ctx, *args, **kwargs) + + if not conversation_has_documents and not has_not_pdf_docs: + # No documents to process + pass + elif has_not_pdf_docs: + # Already handled above with RAG tool + pass else: conversation_documents = [ cd diff --git a/src/backend/pyproject.toml b/src/backend/pyproject.toml index 57cf007..568ff9d 100644 --- a/src/backend/pyproject.toml +++ b/src/backend/pyproject.toml @@ -53,6 +53,10 @@ dependencies = [ "markitdown==0.0.2", "mozilla-django-oidc==4.0.1", "nested-multipart-parser==1.6.0", + "matplotlib==3.9.2", + "numpy==2.1.3", + "openpyxl==3.1.5", + "pandas==2.2.3", "posthog==7.0.0", "pydantic==2.12.4", "pydantic-ai-slim[openai,mistral,mcp,evals,logfire]==1.17.0",