add way to split xlsx for storage
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user