Source code for agency.tools.database
import re
from typing import Type, Dict, List
from pydantic import BaseModel, Field
from langchain.tools import BaseTool
from django.forms import model_to_dict
from django.db.models import Q
from django.apps import apps
from django.db import models as django_models
from agency.utils import app_log
[docs]
class DatabaseSchemaGenerator:
"""Automatically generates database schema for the chatbot."""
[docs]
@staticmethod
def get_field_info(field) -> Dict:
"""Extract useful information from a model field.
Args:
field: Django model field
Returns:
Dictionary with field information
"""
info = {
'name': field.name,
'type': field.get_internal_type(),
'is_relation': field.is_relation,
'null': field.null,
'blank': field.blank,
}
if hasattr(field, 'max_length') and field.max_length:
info['max_length'] = field.max_length
if field.is_relation:
info['related_model'] = field.related_model.__name__
info['relation_type'] = type(field).__name__
if hasattr(field, 'choices') and field.choices:
info['choices'] = [choice[0] for choice in field.choices]
return info
[docs]
@staticmethod
def get_model_info(model) -> Dict:
"""Extract complete information from a model.
Args:
model: Django model class
Returns:
Dictionary with model information
"""
fields = []
searchable_fields = []
date_fields = []
for field in model._meta.fields:
field_info = DatabaseSchemaGenerator.get_field_info(field)
fields.append(field_info)
if field_info['type'] in ['CharField', 'TextField']:
searchable_fields.append(field.name)
if field_info['type'] in ['DateTimeField', 'DateField']:
date_fields.append(field.name)
m2m_fields = []
for field in model._meta.many_to_many:
m2m_info = DatabaseSchemaGenerator.get_field_info(field)
m2m_fields.append(m2m_info)
return {
'name': model.__name__,
'app_label': model._meta.app_label,
'table_name': model._meta.db_table,
'verbose_name': str(model._meta.verbose_name),
'verbose_name_plural': str(model._meta.verbose_name_plural),
'fields': fields,
'm2m_fields': m2m_fields,
'searchable_fields': searchable_fields,
'date_fields': date_fields,
'has_user_relation': any(f['name'] == 'user' and f['is_relation'] for f in fields),
'has_created_at': any(f['name'] == 'created_at' for f in fields),
'doc_string': model.__doc__.strip() if model.__doc__ else None,
}
[docs]
@staticmethod
def generate_schema(app_names: List[str] = None) -> Dict:
"""Generate complete schema for specified apps.
Args:
app_names: App labels to include. Default: installed domain apps
(``agency``, ``chat``, ``boe``, ``bopv``) — never the legacy
``system`` label.
Returns:
Dictionary with schema information
"""
if app_names is None:
preferred = ('agency', 'chat', 'boe', 'bopv')
app_names = [label for label in preferred if apps.is_installed(label)]
if not app_names:
app_names = ['agency']
schema = {
'models': {},
'apps': {},
'summary': {
'total_models': 0,
'models_with_user_relation': 0,
'searchable_models': 0,
}
}
for app_name in app_names:
try:
app_config = apps.get_app_config(app_name)
app_models = []
for model in app_config.get_models():
model_info = DatabaseSchemaGenerator.get_model_info(model)
schema['models'][model.__name__] = model_info
app_models.append(model.__name__)
schema['summary']['total_models'] += 1
if model_info['has_user_relation']:
schema['summary']['models_with_user_relation'] += 1
if model_info['searchable_fields']:
schema['summary']['searchable_models'] += 1
schema['apps'][app_name] = {
'models': app_models,
'verbose_name': app_config.verbose_name,
}
except LookupError as e:
app_log.debug(f"Skipping unknown app {app_name}: {e}")
except Exception as e:
app_log.warning(f"Could not process app {app_name}: {e}")
return schema
[docs]
@staticmethod
def generate_tool_description(schema: Dict) -> str:
"""Generate tool description based on schema.
Args:
schema: Database schema dictionary
Returns:
Formatted tool description string
"""
description_parts = [
"Use this tool to retrieve objects from the local models database.",
"",
"Available models:"
]
for app_name, app_info in schema['apps'].items():
if app_info['models']:
description_parts.append(f"\n{app_info['verbose_name']} ({app_name}):")
for model_name in app_info['models']:
model_info = schema['models'][model_name]
model_desc = f" \"{model_name}\": {model_info['verbose_name']}"
extra_info = []
if model_info['has_user_relation']:
extra_info.append("user-related")
if model_info['searchable_fields']:
extra_info.append(f"searchable by: {', '.join(model_info['searchable_fields'])}")
if model_info['date_fields']:
extra_info.append(f"date fields: {', '.join(model_info['date_fields'])}")
if extra_info:
model_desc += f" ({', '.join(extra_info)})"
if model_info['doc_string']:
model_desc += f"\n {model_info['doc_string']}"
description_parts.append(model_desc)
description_parts.extend([
"",
"Parameters:",
"model: Choose one of the available models listed above",
"user: User ID (optional, only for models that have user relationships)",
"query: Search term to filter results (optional)",
"q_filter: Django Q object filter as text (e.g., \"Q(tools=True)\", \"Q(name__icontains='test')\")",
"limit: Maximum number of results to return (default: 10)",
])
if 'Message' in schema['models']:
description_parts.extend([
"",
"Message History Examples:",
"Get recent messages: model=\"Message\", user=1, limit=10",
"Search by content: model=\"Message\", query=\"search term\"",
"Filter by date: model=\"Message\", q_filter=\"Q(created_at__gte='2024-01-01')\"",
])
description_parts.extend([
"",
"Q Filter Examples:",
"Simple filter: q_filter=\"Q(active=True)\"",
"Text search: q_filter=\"Q(name__icontains='search_term')\"",
"Date range: q_filter=\"Q(created_at__gte='2024-01-01') & Q(created_at__lte='2024-12-31')\"",
"Complex filter: q_filter=\"Q(category='important') | Q(priority='high')\"",
"",
"Returns a JSON list of objects from the specified model.",
"Report any errors that occur during the query."
])
return "\n".join(description_parts)
[docs]
class DatabaseReadTool(BaseTool):
"""Tool to read data from the local database with auto-generated schema."""
name: str = "database_read"
description: str = "" # Will be auto-generated
db_schema: Dict = Field(default_factory=dict, description="Database schema information")
[docs]
def __init__(self, app_names: List[str] = None, **kwargs):
"""Initialize the tool with auto-generated schema.
Args:
app_names: List of app names to include in schema
**kwargs: Additional arguments for BaseTool
"""
super().__init__(**kwargs)
self.db_schema = DatabaseSchemaGenerator.generate_schema(app_names)
self.description = DatabaseSchemaGenerator.generate_tool_description(self.db_schema)
app_log.info(f"DatabaseReadTool initialized with {self.db_schema['summary']['total_models']} models")
[docs]
class Arguments(BaseModel):
"""Input parameters for DatabaseReadTool.
Args:
model: Database model to query
user: User ID if you know it (optional)
query: Search query for content filtering (optional)
q_filter: Django Q object filter as text (optional)
limit: Maximum number of objects to return (default: 10)
"""
model: str = Field(description="Database model to query")
user: int|None = Field(description="User ID if you know it", default=None)
query: str = Field(description="Search query for content filtering, only if needed", default="")
q_filter: str = Field(description="Django Q object filter as text", default="")
limit: int = Field(default=10, description="Maximum number of objects to return (default 10)")
args_schema: Type[BaseModel] = Arguments
return_direct: bool = False
[docs]
def get_model_class(self, model_name: str):
"""Get model class by name.
Args:
model_name: Name of the model to retrieve
Returns:
Django model class
Raises:
ValueError: If model not found or cannot be loaded
"""
if model_name not in self.db_schema['models']:
available_models = list(self.db_schema['models'].keys())
raise ValueError(f"Model '{model_name}' not found. Available models: {available_models}")
model_info = self.db_schema['models'][model_name]
app_name = model_info['app_label']
try:
app_config = apps.get_app_config(app_name)
return app_config.get_model(model_name)
except Exception as e:
raise ValueError(f"Could not load model {model_name} from app {app_name}: {e}")
def _safe_eval_q_filter(self, q_filter_str: str) -> Q:
"""Safely evaluate a Q filter string.
Args:
q_filter_str: String representation of Django Q object
Returns:
Django Q object
Raises:
ValueError: If filter contains invalid characters or dangerous patterns
"""
if not q_filter_str.strip():
return Q()
allowed_chars = set('Q()&|~,=<>!:\'" ._0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ')
if not all(c in allowed_chars for c in q_filter_str):
raise ValueError("Invalid characters in Q filter")
dangerous_patterns = [
r'__import__', r'eval\(', r'exec\(', r'open\(', r'file\(',
r'input\(', r'raw_input\(', r'compile\(', r'globals\(',
r'locals\(', r'vars\(', r'dir\(', r'help\(', r'breakpoint\(',
]
for pattern in dangerous_patterns:
if re.search(pattern, q_filter_str, re.IGNORECASE):
raise ValueError(f"Potentially dangerous pattern detected: {pattern}")
try:
safe_dict = {'Q': Q}
result = eval(q_filter_str, {"__builtins__": {}}, safe_dict)
if not isinstance(result, Q):
raise ValueError("Expression must evaluate to a Q object")
return result
except Exception as e:
raise ValueError(f"Invalid Q filter expression: {str(e)}")
def _run(self,
model: str,
user: int|None = None,
query: str = "",
q_filter: str = "",
limit: int = 10,
) -> list:
"""Execute search on the specified model.
Args:
model: Model name to query
user: Optional user ID for filtering
query: Optional search term for text fields
q_filter: Optional Django Q filter string
limit: Maximum number of results (default: 10)
Returns:
List of model objects as dictionaries, or error dict
"""
try:
model_class = self.get_model_class(model)
model_info = self.db_schema['models'][model]
except ValueError as e:
return {"error": str(e)}
objects = model_class.objects.all()
if model_info['has_user_relation'] and user:
objects = objects.filter(user=user)
if q_filter.strip():
try:
q_obj = self._safe_eval_q_filter(q_filter)
objects = objects.filter(q_obj)
except Exception as e:
return {"error": f"Error applying Q filter: {str(e)}"}
if query.strip() and model_info['searchable_fields']:
text_q = Q()
for field_name in model_info['searchable_fields']:
text_q |= Q(**{f"{field_name}__icontains": query})
objects = objects.filter(text_q)
if model_info['has_created_at']:
objects = objects.order_by('-created_at')
try:
results = [model_to_dict(o) for o in objects[:limit]]
app_log.info(f"DatabaseReadTool found {len(results)} objects for {model}")
return results
except Exception as e:
return {"error": f"Error serializing results: {str(e)}"}
[docs]
def get_schema_info(self) -> Dict:
"""Return schema information for debugging.
Returns:
Dictionary with schema information
"""
return self.db_schema