add way to split xlsx for storage

This commit is contained in:
camilleAND
2026-01-14 16:37:53 +01:00
parent b6449addb4
commit 0bdee3025b
3 changed files with 286 additions and 3 deletions
@@ -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:
+28 -1
View File
@@ -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
+4
View File
@@ -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",