Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 0bdee3025b | |||
| b6449addb4 |
@@ -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
|
||||
|
||||
@@ -0,0 +1,873 @@
|
||||
"""Data analysis tool for tabular files (CSV, Excel)."""
|
||||
|
||||
import base64
|
||||
import functools
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from io import BytesIO
|
||||
from typing import Any, Dict
|
||||
|
||||
import matplotlib
|
||||
import numpy as np
|
||||
|
||||
matplotlib.use("Agg") # Use non-interactive backend
|
||||
import matplotlib.pyplot as plt
|
||||
import pandas as pd
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.files.storage import default_storage
|
||||
from django.db.models import Q
|
||||
from asgiref.sync import sync_to_async
|
||||
from pydantic_ai import Agent, RunContext
|
||||
from pydantic_ai.exceptions import ModelRetry
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
|
||||
from core.file_upload.enums import AttachmentStatus
|
||||
from core.file_upload.utils import generate_retrieve_policy
|
||||
|
||||
from chat.agents.base import BaseAgent, prepare_custom_model
|
||||
from chat.models import ChatConversationAttachment
|
||||
from chat.tools.exceptions import ModelCannotRetry
|
||||
from chat.tools.utils import last_model_retry_soft_fail
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# MIME types for tabular files
|
||||
TABULAR_MIME_TYPES = [
|
||||
"text/csv",
|
||||
"application/csv",
|
||||
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
"application/vnd.ms-excel",
|
||||
"application/excel",
|
||||
]
|
||||
|
||||
|
||||
@sync_to_async
|
||||
def read_tabular_file(attachment) -> bytes:
|
||||
"""Read tabular file content asynchronously."""
|
||||
with default_storage.open(attachment.key, "rb") as f:
|
||||
return f.read()
|
||||
|
||||
|
||||
def detect_csv_separator(file_data: bytes) -> str:
|
||||
"""
|
||||
Detect the CSV separator by analyzing the first few lines.
|
||||
|
||||
Returns the most likely separator: ',', ';', or '\t'
|
||||
"""
|
||||
# Read first 10KB to analyze
|
||||
sample = file_data[:10240].decode("utf-8", errors="ignore")
|
||||
lines = sample.split("\n")[:10] # First 10 lines
|
||||
|
||||
if not lines:
|
||||
return "," # Default to comma
|
||||
|
||||
# Count occurrences of each separator in the first few lines
|
||||
comma_count = sum(line.count(",") for line in lines)
|
||||
semicolon_count = sum(line.count(";") for line in lines)
|
||||
tab_count = sum(line.count("\t") for line in lines)
|
||||
|
||||
# Return the separator with the highest count
|
||||
if tab_count > comma_count and tab_count > semicolon_count:
|
||||
return "\t"
|
||||
elif semicolon_count > comma_count:
|
||||
return ";"
|
||||
else:
|
||||
return "," # Default to comma
|
||||
|
||||
|
||||
def _convert_to_serializable(obj: Any) -> Any:
|
||||
"""
|
||||
Convert pandas/numpy types to Python native types for JSON serialization.
|
||||
|
||||
Handles:
|
||||
- pandas DataFrame/Series
|
||||
- numpy scalars (int64, float64, etc.)
|
||||
- numpy arrays
|
||||
- pandas Timestamp
|
||||
- Other non-serializable types
|
||||
|
||||
Args:
|
||||
obj: The object to convert.
|
||||
|
||||
Returns:
|
||||
A JSON-serializable version of the object.
|
||||
"""
|
||||
|
||||
# Handle pandas DataFrame
|
||||
if isinstance(obj, pd.DataFrame):
|
||||
# Limit number of rows to avoid huge responses
|
||||
if len(obj) > 1000:
|
||||
obj = obj.head(1000)
|
||||
logger.warning("Result truncated to 1000 rows")
|
||||
return obj.to_dict(orient="records")
|
||||
|
||||
# Handle pandas Series
|
||||
if isinstance(obj, pd.Series):
|
||||
# Convert Series to dict, handling index
|
||||
result_dict = obj.to_dict()
|
||||
# Convert any numpy/pandas types in the values
|
||||
return {str(k): _convert_to_serializable(v) for k, v in result_dict.items()}
|
||||
|
||||
# Handle numpy scalars
|
||||
if isinstance(obj, (np.integer, np.floating)):
|
||||
return obj.item() # Convert to Python native int/float
|
||||
|
||||
# Handle numpy arrays
|
||||
if isinstance(obj, np.ndarray):
|
||||
return obj.tolist()
|
||||
|
||||
# Handle pandas Timestamp
|
||||
if isinstance(obj, pd.Timestamp):
|
||||
return obj.isoformat()
|
||||
|
||||
# Handle lists and tuples - recursively convert elements
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [_convert_to_serializable(item) for item in obj]
|
||||
|
||||
# Handle dicts - recursively convert values
|
||||
if isinstance(obj, dict):
|
||||
return {str(k): _convert_to_serializable(v) for k, v in obj.items()}
|
||||
|
||||
# Handle None, bool, int, float, str - these are already serializable
|
||||
if obj is None or isinstance(obj, (bool, int, float, str)):
|
||||
return obj
|
||||
|
||||
# Fallback: try to convert to string
|
||||
try:
|
||||
return str(obj)
|
||||
except Exception:
|
||||
logger.warning("Could not serialize object of type %s, returning None", type(obj))
|
||||
return None
|
||||
|
||||
|
||||
def _is_valid_excel_file(file_data: bytes, file_name: str) -> bool:
|
||||
"""
|
||||
Check if the file data appears to be a valid Excel file.
|
||||
|
||||
XLSX files are ZIP archives, so they should start with ZIP signature (PK\x03\x04).
|
||||
XLS files have a different signature.
|
||||
"""
|
||||
if not file_data:
|
||||
return False
|
||||
|
||||
file_lower = file_name.lower()
|
||||
|
||||
# Check for XLSX (ZIP-based) signature
|
||||
if file_lower.endswith((".xlsx", ".xlsm", ".xlsb")):
|
||||
# XLSX files are ZIP archives, should start with PK\x03\x04
|
||||
return file_data[:4] == b"PK\x03\x04"
|
||||
|
||||
# Check for XLS (OLE2) signature
|
||||
if file_lower.endswith(".xls"):
|
||||
# XLS files are OLE2 compound documents, should start with specific signature
|
||||
# Common signatures: 0xD0CF11E0 (OLE2) or 0x504B0304 (sometimes saved as ZIP)
|
||||
return (
|
||||
file_data[:4] == b"\xd0\xcf\x11\xe0" # OLE2 signature
|
||||
or file_data[:4] == b"PK\x03\x04" # Sometimes XLS are actually ZIP
|
||||
)
|
||||
|
||||
return False
|
||||
|
||||
|
||||
@sync_to_async
|
||||
def load_dataframe(file_data: bytes, content_type: str, file_name: str) -> Dict[str, pd.DataFrame]:
|
||||
"""
|
||||
Load tabular file into pandas DataFrames.
|
||||
|
||||
Returns a dictionary mapping sheet/table names to DataFrames.
|
||||
For CSV files, uses 'default' as the key.
|
||||
For Excel files, uses sheet names as keys.
|
||||
"""
|
||||
try:
|
||||
# Handle CSV files - also accept text/plain if file extension is .csv
|
||||
if content_type in ["text/csv", "application/csv"] or (
|
||||
content_type == "text/plain" and file_name.lower().endswith(".csv")
|
||||
):
|
||||
# Detect the separator
|
||||
separator = detect_csv_separator(file_data)
|
||||
logger.debug("Detected CSV separator: %r", separator)
|
||||
|
||||
# Read CSV with detected separator
|
||||
df = pd.read_csv(
|
||||
BytesIO(file_data),
|
||||
sep=separator,
|
||||
on_bad_lines="skip", # Skip problematic lines
|
||||
engine="python", # More flexible parser
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
if df.empty:
|
||||
raise ValueError("CSV file appears to be empty or could not be parsed")
|
||||
|
||||
return {"default": df}
|
||||
elif content_type in [
|
||||
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
"application/vnd.ms-excel",
|
||||
"application/excel",
|
||||
] or file_name.lower().endswith((".xlsx", ".xls", ".xlsm", ".xlsb")):
|
||||
# Validate Excel file format before attempting to read
|
||||
if not _is_valid_excel_file(file_data, file_name):
|
||||
logger.warning(
|
||||
"File '%s' does not appear to be a valid Excel file. "
|
||||
"File size: %d bytes, First bytes: %r",
|
||||
file_name,
|
||||
len(file_data),
|
||||
file_data[:20] if len(file_data) >= 20 else file_data,
|
||||
)
|
||||
raise ValueError(
|
||||
f"File '{file_name}' does not appear to be a valid Excel file. "
|
||||
"It may be corrupted or in an unsupported format."
|
||||
)
|
||||
|
||||
file_lower = file_name.lower()
|
||||
dataframes = {}
|
||||
|
||||
# Try different engines based on file extension
|
||||
if file_lower.endswith(".xls"):
|
||||
# Old Excel format - try xlrd engine
|
||||
try:
|
||||
logger.debug("Attempting to read .xls file with xlrd engine")
|
||||
excel_file = pd.ExcelFile(BytesIO(file_data), engine="xlrd")
|
||||
dataframes = {
|
||||
sheet_name: excel_file.parse(sheet_name)
|
||||
for sheet_name in excel_file.sheet_names
|
||||
}
|
||||
except Exception as xlrd_error:
|
||||
logger.warning("Failed to read .xls with xlrd: %s", xlrd_error)
|
||||
# Fallback: try openpyxl (sometimes .xls files are actually .xlsx)
|
||||
try:
|
||||
logger.debug("Trying openpyxl as fallback for .xls file")
|
||||
excel_file = pd.ExcelFile(BytesIO(file_data), engine="openpyxl")
|
||||
dataframes = {
|
||||
sheet_name: excel_file.parse(sheet_name)
|
||||
for sheet_name in excel_file.sheet_names
|
||||
}
|
||||
except Exception as openpyxl_error:
|
||||
logger.error("Failed to read .xls with both engines: %s", openpyxl_error)
|
||||
raise ValueError(
|
||||
f"Failed to read Excel file '{file_name}': "
|
||||
f"xlrd error: {xlrd_error}, openpyxl error: {openpyxl_error}"
|
||||
) from openpyxl_error
|
||||
else:
|
||||
# XLSX format - try openpyxl first
|
||||
try:
|
||||
logger.debug("Attempting to read Excel file with openpyxl engine")
|
||||
excel_file = pd.ExcelFile(BytesIO(file_data), engine="openpyxl")
|
||||
dataframes = {
|
||||
sheet_name: excel_file.parse(sheet_name)
|
||||
for sheet_name in excel_file.sheet_names
|
||||
}
|
||||
except Exception as openpyxl_error:
|
||||
logger.warning("Failed to read with openpyxl: %s", openpyxl_error)
|
||||
# Try calamine engine if available (faster and more robust)
|
||||
try:
|
||||
logger.debug("Trying calamine engine as fallback")
|
||||
excel_file = pd.ExcelFile(BytesIO(file_data), engine="calamine")
|
||||
dataframes = {
|
||||
sheet_name: excel_file.parse(sheet_name)
|
||||
for sheet_name in excel_file.sheet_names
|
||||
}
|
||||
except ImportError:
|
||||
logger.debug("calamine engine not available")
|
||||
raise ValueError(
|
||||
f"Failed to read Excel file '{file_name}' with openpyxl: {openpyxl_error}. "
|
||||
"The file may be corrupted or in an unsupported format."
|
||||
) from openpyxl_error
|
||||
except Exception as calamine_error:
|
||||
logger.error("Failed to read with calamine: %s", calamine_error)
|
||||
raise ValueError(
|
||||
f"Failed to read Excel file '{file_name}': "
|
||||
f"openpyxl error: {openpyxl_error}, calamine error: {calamine_error}"
|
||||
) from calamine_error
|
||||
|
||||
if not dataframes:
|
||||
raise ValueError(f"Excel file '{file_name}' contains no readable sheets")
|
||||
|
||||
logger.info(
|
||||
"Successfully loaded Excel file '%s' with %d sheet(s): %s",
|
||||
file_name,
|
||||
len(dataframes),
|
||||
list(dataframes.keys()),
|
||||
)
|
||||
return dataframes
|
||||
else:
|
||||
raise ValueError(f"Unsupported content type: {content_type}")
|
||||
except Exception as e:
|
||||
logger.error("Error loading tabular file: %s", e, exc_info=True)
|
||||
raise ModelCannotRetry(f"Failed to load file: {str(e)}") from e
|
||||
|
||||
|
||||
def generate_metadata(dataframes: Dict[str, pd.DataFrame], file_name: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Generate metadata about the tabular file.
|
||||
|
||||
Returns:
|
||||
Dictionary containing:
|
||||
- sheets: List of sheet/table names
|
||||
- schemas: Dictionary mapping sheet names to their schemas
|
||||
- row_counts: Dictionary mapping sheet names to row counts
|
||||
- column_info: Dictionary mapping sheet names to column information
|
||||
"""
|
||||
metadata = {
|
||||
"file_name": file_name,
|
||||
"sheets": list(dataframes.keys()),
|
||||
"schemas": {},
|
||||
"row_counts": {},
|
||||
"column_info": {},
|
||||
}
|
||||
|
||||
for sheet_name, df in dataframes.items():
|
||||
# Schema: column names and types
|
||||
metadata["schemas"][sheet_name] = {
|
||||
col: str(dtype) for col, dtype in df.dtypes.items()
|
||||
}
|
||||
|
||||
# Row count
|
||||
metadata["row_counts"][sheet_name] = len(df)
|
||||
|
||||
# Column info: name, type, sample values, null counts
|
||||
metadata["column_info"][sheet_name] = {}
|
||||
for col in df.columns:
|
||||
col_info = {
|
||||
"type": str(df[col].dtype),
|
||||
"null_count": int(df[col].isna().sum()),
|
||||
"unique_count": int(df[col].nunique()),
|
||||
}
|
||||
# Add sample values (non-null)
|
||||
sample_values = df[col].dropna().head(5).tolist()
|
||||
if sample_values:
|
||||
col_info["sample_values"] = [str(v) for v in sample_values]
|
||||
metadata["column_info"][sheet_name][col] = col_info
|
||||
|
||||
return metadata
|
||||
|
||||
|
||||
async def generate_query(
|
||||
user_query: str, metadata: Dict[str, Any], query_agent: BaseAgent, ctx: RunContext
|
||||
) -> str:
|
||||
"""
|
||||
Use an LLM agent to generate a pandas query from user query and file metadata.
|
||||
"""
|
||||
metadata_str = json.dumps(metadata, indent=2)
|
||||
|
||||
prompt = f"""You are a data analysis assistant. Given a user query and file metadata, generate a Python pandas query to answer the question.
|
||||
|
||||
File metadata:
|
||||
{metadata_str}
|
||||
|
||||
User query: {user_query}
|
||||
|
||||
Generate a Python code snippet that:
|
||||
1. Uses pandas operations (filter, groupby, aggregate, etc.)
|
||||
2. Works with the dataframes loaded in memory (available as 'dataframes' dict)
|
||||
3. Assigns the final result to a variable named 'result'
|
||||
4. Handles the specific sheet/table if multiple sheets exist
|
||||
5. ALWAYS handles NaN/NA values in boolean conditions using .notna() or .fillna() before filtering
|
||||
6. If the user asks for a plot/graph/chart, create it using matplotlib and save to 'plot_image' variable as base64
|
||||
|
||||
IMPORTANT RULES:
|
||||
- The code MUST assign the final result to a variable named 'result'
|
||||
- When filtering with conditions, ALWAYS check for NaN first: df[df['col'].notna() & (df['col'] > value)]
|
||||
- Use .dropna() if you need to remove rows with missing values
|
||||
- Use .fillna() if you need to replace missing values
|
||||
- If plotting: use plt (already imported), create the plot, convert to base64:
|
||||
```python
|
||||
plt.figure(figsize=(10, 6))
|
||||
# ... your plot code ...
|
||||
buf = BytesIO()
|
||||
plt.savefig(buf, format='png')
|
||||
buf.seek(0)
|
||||
plot_image = base64.b64encode(buf.getvalue()).decode('utf-8')
|
||||
plt.close()
|
||||
```
|
||||
NOTE: Do NOT use import statements - plt, base64, BytesIO are already available.
|
||||
|
||||
Return ONLY the Python code, without markdown formatting or explanations. The code should be executable and use variables:
|
||||
- 'dataframes': dict mapping sheet names to DataFrames
|
||||
- Sheet names available: {', '.join(metadata['sheets'])}
|
||||
|
||||
Example format (without plot):
|
||||
df = dataframes['default']
|
||||
df = df[df['column'].notna()] # Remove NaN values first
|
||||
result = df[df['column'] > 100].groupby('category').sum()
|
||||
|
||||
Example format (with plot):
|
||||
df = dataframes['default']
|
||||
plt.figure(figsize=(10, 6))
|
||||
plt.plot(df.index, df['close'])
|
||||
plt.xlabel('Index')
|
||||
plt.ylabel('Close')
|
||||
plt.title('Close vs Index')
|
||||
buf = BytesIO()
|
||||
plt.savefig(buf, format='png')
|
||||
buf.seek(0)
|
||||
plot_image = base64.b64encode(buf.getvalue()).decode('utf-8')
|
||||
plt.close()
|
||||
result = "Plot generated successfully. The plot image has been saved and is available in the tool response."
|
||||
|
||||
IMPORTANT:
|
||||
- Do NOT use import statements in the code. All necessary modules (pd, plt, np, base64, BytesIO) are already available. Do NOT use anything else than these modules.
|
||||
- When returning the result text, mention that a plot was generated and will be available in the response, but do NOT include the URL in the text - the system will handle displaying it.
|
||||
|
||||
Generate the query code:"""
|
||||
|
||||
try:
|
||||
response = await query_agent.run(prompt, usage=ctx.usage)
|
||||
query_code = response.output.strip()
|
||||
|
||||
# Extract code from markdown code blocks if present
|
||||
if "```python" in query_code:
|
||||
query_code = query_code.split("```python")[1].split("```")[0].strip()
|
||||
elif "```" in query_code:
|
||||
query_code = query_code.split("```")[1].split("```")[0].strip()
|
||||
|
||||
return query_code
|
||||
except Exception as e:
|
||||
logger.error("Error generating query: %s", e, exc_info=True)
|
||||
raise ModelRetry("Failed to generate query. Please try rephrasing your question.") from e
|
||||
|
||||
|
||||
@sync_to_async
|
||||
def execute_query(query_code: str, dataframes: Dict[str, pd.DataFrame]) -> Any:
|
||||
"""
|
||||
Execute the generated pandas query safely.
|
||||
|
||||
Note: Uses exec() in a restricted environment. The query code is generated
|
||||
by an LLM based on file metadata, so it should be relatively safe, but
|
||||
we restrict the available builtins and globals.
|
||||
"""
|
||||
try:
|
||||
# Pre-process dataframes to handle common issues
|
||||
processed_dataframes = {}
|
||||
for name, df in dataframes.items():
|
||||
# Make a copy to avoid modifying original
|
||||
df_processed = df.copy()
|
||||
# Replace common NaN representations
|
||||
df_processed = df_processed.replace(["", " ", "nan", "NaN", "None", "null"], pd.NA)
|
||||
processed_dataframes[name] = df_processed
|
||||
|
||||
# Create a safe execution environment
|
||||
safe_globals = {
|
||||
"pd": pd,
|
||||
"plt": plt,
|
||||
"np": np,
|
||||
"base64": base64,
|
||||
"BytesIO": BytesIO,
|
||||
"dataframes": processed_dataframes,
|
||||
"__builtins__": {
|
||||
"len": len,
|
||||
"str": str,
|
||||
"int": int,
|
||||
"float": float,
|
||||
"bool": bool,
|
||||
"list": list,
|
||||
"dict": dict,
|
||||
"set": set,
|
||||
"tuple": tuple,
|
||||
"range": range,
|
||||
"sum": sum,
|
||||
"max": max,
|
||||
"min": min,
|
||||
"abs": abs,
|
||||
"round": round,
|
||||
},
|
||||
}
|
||||
|
||||
# Clean up query code - remove any import statements that might cause issues
|
||||
# Split by lines and filter out import statements
|
||||
lines = query_code.split("\n")
|
||||
cleaned_lines = [
|
||||
line
|
||||
for line in lines
|
||||
if not line.strip().startswith("import ") and not line.strip().startswith("from ")
|
||||
]
|
||||
query_code = "\n".join(cleaned_lines)
|
||||
|
||||
# Execute the query in a restricted namespace
|
||||
local_vars = {}
|
||||
exec(query_code, safe_globals, local_vars) # noqa: S102
|
||||
|
||||
# Get the result - check if 'result' variable exists, otherwise try 'df'
|
||||
if "result" in local_vars:
|
||||
result = local_vars["result"]
|
||||
elif "df" in local_vars:
|
||||
result = local_vars["df"]
|
||||
else:
|
||||
# If no explicit result variable, get the last expression
|
||||
# This is a fallback - ideally the LLM should assign to 'result'
|
||||
raise ValueError("Query must assign result to 'result' variable")
|
||||
|
||||
# Check if a plot was generated
|
||||
plot_image = None
|
||||
if "plot_image" in local_vars:
|
||||
plot_image = local_vars["plot_image"]
|
||||
logger.info("Plot image generated")
|
||||
|
||||
# Convert result to a serializable format
|
||||
result = _convert_to_serializable(result)
|
||||
|
||||
return {"result": result, "plot_image": plot_image}
|
||||
except Exception as e:
|
||||
logger.error("Error executing query: %s", e, exc_info=True)
|
||||
# Provide more helpful error message
|
||||
error_msg = str(e)
|
||||
if "NaN" in error_msg or "NA" in error_msg:
|
||||
error_msg = (
|
||||
f"{error_msg}. "
|
||||
"The query may need to handle missing values (NaN/NA) using .notna() or .dropna() before filtering."
|
||||
)
|
||||
raise ModelCannotRetry(f"Failed to execute query: {error_msg}") from e
|
||||
|
||||
|
||||
@last_model_retry_soft_fail
|
||||
async def data_analysis(ctx: RunContext, query: str) -> ToolReturn:
|
||||
"""
|
||||
Analyze tabular data files (CSV, Excel) based on user query.
|
||||
Can also generate plots/graphs/charts.
|
||||
|
||||
This tool:
|
||||
1. Loads the tabular file(s) from attachments
|
||||
2. Generates metadata about the file structure
|
||||
3. Uses an LLM to generate a pandas query based on user query
|
||||
4. Executes the query and returns results
|
||||
|
||||
Args:
|
||||
ctx (RunContext): The run context containing the conversation.
|
||||
query (str): The user's data analysis question.
|
||||
|
||||
Returns:
|
||||
ToolReturn: Contains the analysis results and metadata.
|
||||
"""
|
||||
try:
|
||||
# Find tabular files in attachments
|
||||
# First, get all attachments for debugging
|
||||
all_attachments = await sync_to_async(list)(
|
||||
ctx.deps.conversation.attachments.all()
|
||||
)
|
||||
logger.info(
|
||||
"All attachments in conversation: %s",
|
||||
[
|
||||
{
|
||||
"file_name": a.file_name,
|
||||
"content_type": a.content_type,
|
||||
"upload_state": a.upload_state,
|
||||
"conversion_from": a.conversion_from,
|
||||
}
|
||||
for a in all_attachments
|
||||
],
|
||||
)
|
||||
|
||||
# Find tabular files - exclude converted files (they have conversion_from set)
|
||||
# First try by content_type
|
||||
tabular_attachments_by_type = await sync_to_async(list)(
|
||||
ctx.deps.conversation.attachments.filter(
|
||||
content_type__in=TABULAR_MIME_TYPES,
|
||||
upload_state=AttachmentStatus.READY,
|
||||
)
|
||||
.filter(
|
||||
Q(conversion_from__isnull=True) | Q(conversion_from="")
|
||||
)
|
||||
)
|
||||
|
||||
# If no files found by content_type, try by file extension as fallback
|
||||
# (some systems detect CSV as text/plain instead of text/csv)
|
||||
if not tabular_attachments_by_type:
|
||||
csv_extensions = [".csv", ".xlsx", ".xls"]
|
||||
all_ready_attachments = await sync_to_async(list)(
|
||||
ctx.deps.conversation.attachments.filter(
|
||||
upload_state=AttachmentStatus.READY,
|
||||
)
|
||||
.filter(
|
||||
Q(conversion_from__isnull=True) | Q(conversion_from="")
|
||||
)
|
||||
)
|
||||
tabular_attachments = [
|
||||
att
|
||||
for att in all_ready_attachments
|
||||
if any(att.file_name.lower().endswith(ext) for ext in csv_extensions)
|
||||
# Exclude Markdown files (converted files have .md extension or content_type text/markdown)
|
||||
and not att.file_name.lower().endswith(".md")
|
||||
and att.content_type != "text/markdown"
|
||||
]
|
||||
if tabular_attachments:
|
||||
logger.info(
|
||||
"Found %d tabular file(s) by extension fallback (content_type was not recognized): %s",
|
||||
len(tabular_attachments),
|
||||
[f"{a.file_name} ({a.content_type})" for a in tabular_attachments],
|
||||
)
|
||||
else:
|
||||
tabular_attachments = tabular_attachments_by_type
|
||||
|
||||
# If still no files found, check if there are converted files that might have originals
|
||||
# This handles the case where an Excel file was converted to Markdown for RAG
|
||||
if not tabular_attachments:
|
||||
# Look for converted files with tabular extensions
|
||||
csv_extensions = [".csv", ".xlsx", ".xls"]
|
||||
converted_attachments = await sync_to_async(list)(
|
||||
ctx.deps.conversation.attachments.filter(
|
||||
upload_state=AttachmentStatus.READY,
|
||||
)
|
||||
.exclude(
|
||||
Q(conversion_from__isnull=True) | Q(conversion_from="")
|
||||
)
|
||||
)
|
||||
|
||||
# For each converted file, try to find the original
|
||||
for converted_att in converted_attachments:
|
||||
if any(converted_att.file_name.lower().endswith(ext) for ext in csv_extensions):
|
||||
# Try to find the original file using conversion_from key
|
||||
original_key = converted_att.conversion_from
|
||||
if original_key:
|
||||
original_attachment = await sync_to_async(
|
||||
ctx.deps.conversation.attachments.filter(
|
||||
key=original_key,
|
||||
upload_state=AttachmentStatus.READY,
|
||||
).first
|
||||
)()
|
||||
if original_attachment:
|
||||
logger.info(
|
||||
"Found original file '%s' for converted file '%s'",
|
||||
original_attachment.file_name,
|
||||
converted_att.file_name,
|
||||
)
|
||||
tabular_attachments.append(original_attachment)
|
||||
break
|
||||
|
||||
logger.info(
|
||||
"Found %d tabular attachment(s): %s",
|
||||
len(tabular_attachments),
|
||||
[f"{a.file_name} ({a.content_type})" for a in tabular_attachments],
|
||||
)
|
||||
|
||||
if not tabular_attachments:
|
||||
raise ModelCannotRetry(
|
||||
"No tabular files (CSV or Excel) found in the conversation. "
|
||||
"Please upload a CSV or Excel file first. "
|
||||
"Note: If you uploaded an Excel file that was converted to Markdown for RAG, "
|
||||
"the original file must still be available."
|
||||
)
|
||||
|
||||
# Use the first tabular file
|
||||
attachment = tabular_attachments[0]
|
||||
logger.info("Analyzing file: %s (type: %s)", attachment.file_name, attachment.content_type)
|
||||
|
||||
# Load file data
|
||||
file_data = await read_tabular_file(attachment)
|
||||
|
||||
# Validate that this is actually a valid Excel/CSV file (not a converted Markdown file)
|
||||
# Check if it's an Excel file that should have ZIP signature
|
||||
if attachment.file_name.lower().endswith((".xlsx", ".xls", ".xlsm", ".xlsb")):
|
||||
if not _is_valid_excel_file(file_data, attachment.file_name):
|
||||
logger.warning(
|
||||
"File '%s' does not appear to be a valid Excel file. "
|
||||
"It may be a converted Markdown file. Searching for original...",
|
||||
attachment.file_name,
|
||||
)
|
||||
# Try to find the original file
|
||||
# Look for an attachment with the same name but without conversion_from
|
||||
original_attachment = await sync_to_async(
|
||||
ctx.deps.conversation.attachments.filter(
|
||||
file_name=attachment.file_name,
|
||||
upload_state=AttachmentStatus.READY,
|
||||
)
|
||||
.filter(
|
||||
Q(conversion_from__isnull=True) | Q(conversion_from="")
|
||||
)
|
||||
.exclude(pk=attachment.pk)
|
||||
.first
|
||||
)()
|
||||
|
||||
if original_attachment:
|
||||
logger.info(
|
||||
"Found original file '%s' (key: %s), using it instead",
|
||||
original_attachment.file_name,
|
||||
original_attachment.key,
|
||||
)
|
||||
attachment = original_attachment
|
||||
file_data = await read_tabular_file(attachment)
|
||||
elif hasattr(attachment, 'conversion_from') and attachment.conversion_from:
|
||||
# Try to find by key if this file has a conversion_from
|
||||
original_attachment = await sync_to_async(
|
||||
ctx.deps.conversation.attachments.filter(
|
||||
key=attachment.conversion_from,
|
||||
upload_state=AttachmentStatus.READY,
|
||||
).first
|
||||
)()
|
||||
if original_attachment:
|
||||
logger.info(
|
||||
"Found original file via conversion_from: '%s'",
|
||||
original_attachment.file_name,
|
||||
)
|
||||
attachment = original_attachment
|
||||
file_data = await read_tabular_file(attachment)
|
||||
else:
|
||||
raise ModelCannotRetry(
|
||||
f"File '{attachment.file_name}' appears to be a converted Markdown file, "
|
||||
"not the original Excel file. The original file is not available. "
|
||||
"Please re-upload the original Excel file."
|
||||
)
|
||||
else:
|
||||
raise ModelCannotRetry(
|
||||
f"File '{attachment.file_name}' does not appear to be a valid Excel file. "
|
||||
"It may be corrupted or in an unsupported format."
|
||||
)
|
||||
|
||||
# Load into pandas DataFrames
|
||||
dataframes = await load_dataframe(file_data, attachment.content_type, attachment.file_name)
|
||||
|
||||
# Generate metadata
|
||||
metadata = generate_metadata(dataframes, attachment.file_name)
|
||||
logger.debug("File metadata: %s", json.dumps(metadata, indent=2))
|
||||
|
||||
# Generate query using LLM
|
||||
# NOTE:
|
||||
# We intentionally create a "bare" Agent instance here instead of using BaseAgent
|
||||
# with tools enabled. Using BaseAgent would attach all configured tools (including
|
||||
# this data_analysis tool itself), which can cause the model to try to call tools
|
||||
# while we're already inside a tool execution, leading to nested tool calls and
|
||||
# failures like "Failed to generate query. Please try rephrasing your question.".
|
||||
#
|
||||
# Here we reuse the same model configuration as BaseAgent but WITHOUT any tools,
|
||||
# so this internal call is purely text-to-text.
|
||||
llm_config = settings.LLM_CONFIGURATIONS[settings.LLM_DEFAULT_MODEL_HRID]
|
||||
if llm_config.is_custom:
|
||||
model_instance = prepare_custom_model(llm_config)
|
||||
else:
|
||||
# Rely on pydantic-ai's built-in model registry / name inference
|
||||
model_instance = llm_config.model_name
|
||||
|
||||
# Use the same keyword as when using BaseAgent, which forwards to Agent.
|
||||
# On the current pydantic_ai version, the correct kwarg is `output_type`,
|
||||
# not `result_type` (passing `result_type` raises a UserError).
|
||||
query_agent = Agent(model=model_instance, output_type=str)
|
||||
query_code = await generate_query(query, metadata, query_agent, ctx)
|
||||
logger.debug("Generated query: %s", query_code)
|
||||
|
||||
# Execute query
|
||||
try:
|
||||
execution_result = await execute_query(query_code, dataframes)
|
||||
result = execution_result.get("result")
|
||||
plot_image_base64 = execution_result.get("plot_image")
|
||||
except Exception as e:
|
||||
logger.error("Query execution failed: %s", e, exc_info=True)
|
||||
raise ModelRetry(
|
||||
f"Failed to execute the generated query: {str(e)}. "
|
||||
"Please try rephrasing your question."
|
||||
) from e
|
||||
|
||||
# Format result for return
|
||||
return_value = {
|
||||
"query": query,
|
||||
"query_code": query_code,
|
||||
"result": result,
|
||||
"metadata": metadata,
|
||||
}
|
||||
|
||||
# Save plot image to storage if generated
|
||||
plot_url = None
|
||||
plot_attachment = None
|
||||
if plot_image_base64:
|
||||
try:
|
||||
# Decode base64 image
|
||||
plot_image_data = base64.b64decode(plot_image_base64)
|
||||
|
||||
# Generate a unique filename for the plot
|
||||
plot_filename = f"plot_{uuid.uuid4().hex[:8]}.png"
|
||||
plot_key = f"{ctx.deps.conversation.pk}/plots/{plot_filename}"
|
||||
|
||||
# Save to storage
|
||||
await sync_to_async(default_storage.save)(
|
||||
plot_key, BytesIO(plot_image_data)
|
||||
)
|
||||
|
||||
# Create a permanent attachment record in the database
|
||||
plot_attachment = await sync_to_async(ChatConversationAttachment.objects.create)(
|
||||
conversation=ctx.deps.conversation,
|
||||
uploaded_by=ctx.deps.user,
|
||||
key=plot_key,
|
||||
file_name=plot_filename,
|
||||
content_type="image/png",
|
||||
upload_state=AttachmentStatus.READY,
|
||||
size=len(plot_image_data),
|
||||
)
|
||||
|
||||
# Generate presigned URL for immediate access (valid for 1 hour)
|
||||
plot_url = await sync_to_async(generate_retrieve_policy)(plot_key)
|
||||
logger.info(
|
||||
"Plot image saved to storage and database: %s (presigned URL: %s)",
|
||||
plot_key,
|
||||
plot_url[:50] + "..."
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("Failed to save plot image: %s", e, exc_info=True)
|
||||
# Continue without plot URL if save fails
|
||||
|
||||
if plot_url:
|
||||
# Include both local and presigned URLs
|
||||
return_value["plot_url"] = plot_url # Presigned URL for direct access
|
||||
return_value["plot_local_url"] = f"/media-key/{plot_key}" # Local URL for reference
|
||||
# Include attachment ID for reference
|
||||
if plot_attachment:
|
||||
return_value["plot_attachment_id"] = str(plot_attachment.pk)
|
||||
|
||||
return ToolReturn(
|
||||
return_value=return_value,
|
||||
metadata={"file_name": attachment.file_name, "content_type": attachment.content_type},
|
||||
)
|
||||
|
||||
except (ModelCannotRetry, ModelRetry):
|
||||
# Re-raise these as-is
|
||||
raise
|
||||
except Exception as exc:
|
||||
# Unexpected error - stop and inform user
|
||||
logger.exception("Unexpected error in data_analysis: %s", exc)
|
||||
raise ModelCannotRetry(
|
||||
f"An unexpected error occurred during data analysis: {type(exc).__name__}. "
|
||||
"You must explain this to the user and not try to answer based on your knowledge."
|
||||
) from exc
|
||||
|
||||
|
||||
def add_data_analysis_tool(agent: Agent) -> None:
|
||||
"""Add the data analysis tool to an existing agent."""
|
||||
|
||||
@agent.tool(retries=2)
|
||||
@functools.wraps(data_analysis)
|
||||
async def data_analysis_tool(ctx: RunContext, query: str) -> ToolReturn:
|
||||
"""
|
||||
Analyze tabular data files (CSV, Excel) based on user query.
|
||||
|
||||
This tool loads tabular files, generates metadata about their structure,
|
||||
uses an LLM to generate a pandas query based on the user's question,
|
||||
executes the query, and returns the results.
|
||||
|
||||
Use this tool when the user asks questions about data in CSV or Excel files,
|
||||
such as:
|
||||
- "What is the average sales by region?"
|
||||
- "Show me the top 10 products by revenue"
|
||||
- "How many records are in this file?"
|
||||
- "Filter data where column X is greater than Y"
|
||||
|
||||
Args:
|
||||
ctx (RunContext): The run context containing the conversation.
|
||||
query (str): The user's data analysis question.
|
||||
"""
|
||||
# Import here to avoid circular import
|
||||
from chat.tools.data_analysis import data_analysis as _data_analysis
|
||||
|
||||
return await _data_analysis(ctx, query)
|
||||
|
||||
@agent.instructions
|
||||
def data_analysis_instructions() -> str:
|
||||
"""Dynamic system prompt function to add data analysis instructions."""
|
||||
return (
|
||||
"When the user asks questions about data in CSV or Excel files, "
|
||||
"use the data_analysis tool to analyze the data and answer their question. "
|
||||
"The tool will handle loading the file, generating queries, and executing them. "
|
||||
"When a plot is generated, the tool returns a 'plot_url' in the result. "
|
||||
"Use this presigned URL directly in markdown image syntax: . "
|
||||
"Do NOT use local URLs like /media-key/... - always use the presigned URL from plot_url. "
|
||||
"Present the results clearly to the user."
|
||||
)
|
||||
|
||||
@@ -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