Spaces:
Sleeping
Sleeping
File size: 3,816 Bytes
3a7eb07 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """
Entity Extractor Agent
Mining-specific Named Entity Recognition
"""
import logging
from typing import Any, Dict, List, Optional
from app.agents.base import BaseAgent
logger = logging.getLogger(__name__)
class EntityExtractorAgent(BaseAgent):
"""
Named Entity Recognition Agent for Mining Documents.
Extracts mining-specific entities:
- Equipment names and models
- Chemical compounds and gases
- Mine locations and sections
- Personnel and roles
- Dates and schedules
- Regulatory references
"""
def __init__(self):
# Cerebras: 1M tokens/day free, 2600+ TPS, 60K TPM — better for high-volume extraction than Groq
super().__init__(model_name="gpt-oss-120b", provider="cerebras")
@property
def system_prompt(self) -> str:
return """You are a named entity extraction agent specialized in mining documents.
Extract the following entity types:
1. EQUIPMENT
- Mining machinery (excavators, haul trucks, drills)
- Brand names and models (Caterpillar D11, Komatsu PC8000)
- Equipment IDs and serial numbers
- Tools and instruments
2. CHEMICALS
- Gases (methane, CO, H2S, oxygen)
- Minerals and ores
- Explosives and blasting agents
- Dust types (coal dust, silica)
- Hazardous substances
3. LOCATIONS
- Mine names
- Sections and portals
- Underground levels
- Surface areas
- Geographic coordinates
4. PERSONNEL
- Names (anonymize if needed)
- Roles (Safety Officer, Foreman, Engineer)
- Departments and teams
- Certifications
5. DATES
- Specific dates
- Deadlines
- Scheduled events
- Time periods
6. REGULATIONS
- MSHA regulations (30 CFR citations)
- OSHA standards
- EPA requirements
- State regulations
- Company policies
Be precise and avoid duplicates. Extract exactly as written in the document.
"""
async def analyze(
self, text: str, context: Optional[Dict] = None
) -> Dict[str, Any]:
"""
Extract named entities from document.
Returns dict with keys: equipment, chemicals, locations, personnel,
dates, regulations (all List[str]), entity_count (int)
"""
prompt = f"""Extract all mining-specific named entities from this document.
Document content ({len(text)} chars total, showing up to 15000):
{self._prepare_text(text)}
Respond with a JSON object:
{{
"equipment": ["<equipment name or model>"],
"chemicals": ["<chemical compound, gas, or mineral>"],
"locations": ["<mine name, section, or area>"],
"personnel": ["<role or name>"],
"dates": ["<date or time period>"],
"regulations": ["<e.g. 30 CFR 75.400, OSHA 1910.134>"]
}}
Notes:
- List each unique entity only once
- Use exact text from document
- For personnel, prefer roles over names for privacy
- Include regulation citations in standard format
"""
result = await self._generate_json(prompt)
entities = {
"equipment": self._deduplicate(result.get("equipment", [])),
"chemicals": self._deduplicate(result.get("chemicals", [])),
"locations": self._deduplicate(result.get("locations", [])),
"personnel": self._deduplicate(result.get("personnel", [])),
"dates": self._deduplicate(result.get("dates", [])),
"regulations": self._deduplicate(result.get("regulations", [])),
}
entities["entity_count"] = sum(len(v) for v in entities.values())
return entities
def _deduplicate(self, items: List[str]) -> List[str]:
"""Remove duplicates while preserving order"""
seen = set()
result = []
for item in items:
if item and item.lower() not in seen:
seen.add(item.lower())
result.append(item)
return result
|