Coverage for src/local_deep_research/domain_classifier/classifier.py: 93%

1"""Domain classifier using LLM for categorization."""

3import json

4from typing import Dict, List, Optional

6from loguru import logger

7from sqlalchemy.orm import Session

9from ..config.llm_config import get_llm

10from ..database.models import ResearchResource

11from ..database.session_context import get_user_db_session

12from ..utilities.sql_utils import escape_like as _escape_like

13from ..utilities.json_utils import extract_json, get_llm_response_text

14from .models import DomainClassification

17# Predefined categories for domain classification

18DOMAIN_CATEGORIES = {

19 "Academic & Research": [

20 "University/Education",

21 "Scientific Journal",

22 "Research Institution",

23 "Academic Database",

24 ],

25 "News & Media": [

26 "General News",

27 "Tech News",

28 "Business News",

29 "Entertainment News",

30 "Local/Regional News",

31 ],

32 "Reference & Documentation": [

33 "Encyclopedia",

34 "Technical Documentation",

35 "API Documentation",

36 "Tutorial/Guide",

37 "Dictionary/Glossary",

38 ],

39 "Social & Community": [

40 "Social Network",

41 "Forum/Discussion",

42 "Q&A Platform", # Question-and-answer focused sites (StackOverflow, Quora)

43 "Blog Platform", # Structured publishing platforms (Medium, WordPress.com)

44 "Personal Blog", # Individual author blogs

45 ],

46 "Business & Commerce": [

47 "E-commerce",

48 "Corporate Website",

49 "B2B Platform",

50 "Financial Service",

51 "Marketing/Advertising",

52 ],

53 "Technology": [

54 "Software Development",

55 "Cloud Service",

56 "Open Source Project",

57 "Tech Company",

58 "Developer Tools",

59 ],

60 "Government & Organization": [

61 "Government Agency",

62 "Non-profit",

63 "International Organization",

64 "Think Tank",

65 "Industry Association",

66 ],

67 "Entertainment & Lifestyle": [

68 "Streaming Service",

69 "Gaming",

70 "Sports",

71 "Arts & Culture",

72 "Travel & Tourism",

73 ],

74 "Professional & Industry": [

75 "Healthcare",

76 "Legal",

77 "Real Estate",

78 "Manufacturing",

79 "Energy & Utilities",

80 ],

81 "Other": ["Personal Website", "Miscellaneous", "Unknown"],

82}

85class DomainClassifier:

86 """Classify domains using LLM with predefined categories."""

88 def __init__(self, username: str, settings_snapshot: dict = None):

89 """Initialize the domain classifier.

91 Args:

92 username: Username for database session

93 settings_snapshot: Settings snapshot for LLM configuration

94 """

95 self.username = username

96 self.settings_snapshot = settings_snapshot

97 self.llm = None

99 def _get_llm(self):

100 """Get or initialize LLM instance."""

101 if self.llm is None:

102 self.llm = get_llm(settings_snapshot=self.settings_snapshot)

103 return self.llm

104

105 def close(self) -> None:

106 """Close the LLM client if one was created."""

107 from ..utilities.resource_utils import safe_close

108

109 safe_close(self.llm, "classifier LLM")

110

111 def _get_domain_samples(

112 self, domain: str, session: Session, limit: int = 5

113 ) -> List[Dict]:

114 """Get sample resources from a domain.

115

116 Args:

117 domain: Domain to get samples for

118 session: Database session

119 limit: Maximum number of samples

120

121 Returns:

122 List of resource samples

123 """

124 resources = (

125 session.query(ResearchResource)

126 .filter(

127 ResearchResource.url.like(

128 f"%{_escape_like(domain)}%", escape="\\"

129 )

130 )

131 .limit(limit)

132 .all()

133 )

134

135 samples = []

136 for resource in resources:

137 samples.append(

138 {

139 "title": resource.title or "Untitled",

140 "url": resource.url,

141 "preview": resource.content_preview[:200]

142 if resource.content_preview

143 else None,

144 }

145 )

146

147 return samples

148

149 def _build_classification_prompt(

150 self, domain: str, samples: List[Dict]

151 ) -> str:

152 """Build prompt for LLM classification.

153

154 This method uses actual content samples (titles, previews) from the domain

155 rather than relying solely on domain name patterns, providing more

156 accurate classification based on actual site content.

157

158 Args:

159 domain: Domain to classify

160 samples: Sample resources from the domain (titles, URLs, content previews)

161

162 Returns:

163 Formatted prompt string

164 """

165 # Format categories for prompt

166 categories_text = []

167 for main_cat, subcats in DOMAIN_CATEGORIES.items():

168 subcats_text = ", ".join(subcats)

169 categories_text.append(f"{main_cat}: {subcats_text}")

170

171 # Format samples

172 samples_text = []

173 for i, sample in enumerate(samples[:5], 1):

174 samples_text.append(f"{i}. Title: {sample['title']}")

175 if sample.get("preview"):

176 samples_text.append(f" Preview: {sample['preview'][:100]}...")

177

178 return f"""Classify the following domain into one of the predefined categories.

179

180Domain: {domain}

181

182Sample content from this domain:

183{chr(10).join(samples_text) if samples_text else "No samples available"}

184

185Available Categories:

186{chr(10).join(categories_text)}

187

188Respond with a JSON object containing:

189- "category": The main category (e.g., "News & Media")

190- "subcategory": The specific subcategory (e.g., "Tech News")

191- "confidence": A confidence score between 0 and 1

192- "reasoning": A brief explanation (max 100 words) of why this classification was chosen

193

194Focus on accuracy. If uncertain, use "Other" category with "Unknown" subcategory.

195

196JSON Response:"""

197

198 def classify_domain(

199 self, domain: str, force_update: bool = False

200 ) -> Optional[DomainClassification]:

201 """Classify a single domain using LLM.

202

203 Args:

204 domain: Domain to classify

205 force_update: If True, reclassify even if already exists

206

207 Returns:

208 DomainClassification object or None if failed

209 """

210 try:

211 with get_user_db_session(self.username) as session:

212 # Check if already classified

213 existing = (

214 session.query(DomainClassification)

215 .filter_by(domain=domain)

216 .first()

217 )

218

219 if existing and not force_update:

220 logger.info(

221 f"Domain {domain} already classified as {existing.category}"

222 )

223 return existing

224

225 # Get sample resources

226 samples = self._get_domain_samples(domain, session)

227

228 # Build prompt and get classification

229 prompt = self._build_classification_prompt(domain, samples)

230 llm = self._get_llm()

231

232 response = llm.invoke(prompt)

233 response_text = get_llm_response_text(response)

234

235 result = extract_json(response_text, expected_type=dict)

236 if result is None:

237 raise ValueError("Could not parse JSON from LLM response") # noqa: TRY301 — inside db session; except logs and returns None

238

239 # Create or update classification

240 if existing: 240 ↛ 241line 240 didn't jump to line 241 because the condition on line 240 was never true

241 existing.category = result.get("category", "Other")

242 existing.subcategory = result.get("subcategory", "Unknown")

243 existing.confidence = float(result.get("confidence", 0.5))

244 existing.reasoning = result.get("reasoning", "")

245 existing.sample_titles = json.dumps(

246 [s["title"] for s in samples]

247 )

248 existing.sample_count = len(samples)

249 classification = existing

250 else:

251 classification = DomainClassification(

252 domain=domain,

253 category=result.get("category", "Other"),

254 subcategory=result.get("subcategory", "Unknown"),

255 confidence=float(result.get("confidence", 0.5)),

256 reasoning=result.get("reasoning", ""),

257 sample_titles=json.dumps([s["title"] for s in samples]),

258 sample_count=len(samples),

259 )

260 session.add(classification)

261

262 session.commit()

263 logger.info(

264 f"Classified {domain} as {classification.category}/{classification.subcategory} with confidence {classification.confidence}"

265 )

266 return classification

267

268 except Exception:

269 logger.exception(f"Error classifying domain {domain}")

270 return None

271

272 def classify_all_domains(

273 self, force_update: bool = False, progress_callback=None

274 ) -> Dict:

275 """Classify all unique domains in the database.

276

277 Args:

278 force_update: If True, reclassify all domains

279 progress_callback: Optional callback function for progress updates

280

281 Returns:

282 Dictionary with classification results

283 """

284 results = {

285 "total": 0,

286 "classified": 0,

287 "failed": 0,

288 "skipped": 0,

289 "domains": [],

290 }

291

292 try:

293 with get_user_db_session(self.username) as session:

294 # Get all unique domains

295 from urllib.parse import urlparse

296

297 resources = session.query(ResearchResource.url).distinct().all()

298 domains = set()

299

300 for (url,) in resources:

301 if url:

302 try:

303 parsed = urlparse(url)

304 domain = parsed.netloc.lower()

305 if domain.startswith("www."):

306 domain = domain[4:]

307 if domain:

308 domains.add(domain)

309 except (ValueError, AttributeError):

310 logger.debug("Skipping malformed URL")

311 continue

312

313 results["total"] = len(domains)

314 logger.info(

315 f"Found {results['total']} unique domains to process"

316 )

317

318 # Classify each domain ONE BY ONE

319 for i, domain in enumerate(sorted(domains), 1):

320 logger.info(

321 f"Processing domain {i}/{results['total']}: {domain}"

322 )

323

324 if progress_callback:

325 progress_callback(

326 {

327 "current": i,

328 "total": results["total"],

329 "domain": domain,

330 "percentage": (i / results["total"]) * 100,

331 }

332 )

333

334 try:

335 # Check if already classified

336 if not force_update:

337 existing = (

338 session.query(DomainClassification)

339 .filter_by(domain=domain)

340 .first()

341 )

342 if existing: 342 ↛ 358line 342 didn't jump to line 358 because the condition on line 342 was always true

343 results["skipped"] += 1

344 results["domains"].append(

345 {

346 "domain": domain,

347 "status": "skipped",

348 "category": existing.category,

349 "subcategory": existing.subcategory,

350 }

351 )

352 logger.info(

353 f"Domain {domain} already classified, skipping"

354 )

355 continue

356

357 # Classify this single domain

358 classification = self.classify_domain(

359 domain, force_update

360 )

361

362 if classification:

363 results["classified"] += 1

364 results["domains"].append(

365 {

366 "domain": domain,

367 "status": "classified",

368 "category": classification.category,

369 "subcategory": classification.subcategory,

370 "confidence": classification.confidence,

371 }

372 )

373 logger.info(

374 f"Successfully classified {domain} as {classification.category}"

375 )

376 else:

377 results["failed"] += 1

378 results["domains"].append(

379 {"domain": domain, "status": "failed"}

380 )

381 logger.warning(

382 f"Failed to classify domain {domain}"

383 )

384

385 except Exception:

386 logger.exception(f"Error classifying domain {domain}")

387 results["failed"] += 1

388 results["domains"].append(

389 {

390 "domain": domain,

391 "status": "failed",

392 "error": "Classification failed",

393 }

394 )

395

396 logger.info(

397 f"Classification complete: {results['classified']} classified, {results['skipped']} skipped, {results['failed']} failed"

398 )

399 return results

400

401 except Exception:

402 logger.exception("Error in classify_all_domains")

403 results["error"] = "Classification failed"

404 return results

405

406 def get_classification(self, domain: str) -> Optional[DomainClassification]:

407 """Get existing classification for a domain.

408

409 Args:

410 domain: Domain to look up

411

412 Returns:

413 DomainClassification object or None if not found

414 """

415 try:

416 with get_user_db_session(self.username) as session:

417 return (

418 session.query(DomainClassification)

419 .filter_by(domain=domain)

420 .first()

421 )

422 except Exception:

423 logger.exception(f"Error getting classification for {domain}")

424 return None

425

426 def get_all_classifications(self) -> List[DomainClassification]:

427 """Get all domain classifications.

428

429 Returns:

430 List of all DomainClassification objects

431 """

432 try:

433 with get_user_db_session(self.username) as session:

434 return (

435 session.query(DomainClassification)

436 .order_by(

437 DomainClassification.category,

438 DomainClassification.domain,

439 )

440 .all()

441 )

442 except Exception:

443 logger.exception("Error getting all classifications")

444 return []

445

446 def get_categories_summary(self) -> Dict:

447 """Get summary of domain classifications by category.

448

449 Returns:

450 Dictionary with category counts and domains

451 """

452 try:

453 with get_user_db_session(self.username) as session:

454 classifications = session.query(DomainClassification).all()

455

456 summary = {}

457 for classification in classifications:

458 cat = classification.category

459 if cat not in summary:

460 summary[cat] = {

461 "count": 0,

462 "domains": [],

463 "subcategories": {},

464 }

465

466 summary[cat]["count"] += 1

467 summary[cat]["domains"].append(

468 {

469 "domain": classification.domain,

470 "subcategory": classification.subcategory,

471 "confidence": classification.confidence,

472 }

473 )

474

475 subcat = classification.subcategory

476 if subcat:

477 if subcat not in summary[cat]["subcategories"]: 477 ↛ 479line 477 didn't jump to line 479 because the condition on line 477 was always true

478 summary[cat]["subcategories"][subcat] = 0

479 summary[cat]["subcategories"][subcat] += 1

480

481 return summary

482

483 except Exception:

484 logger.exception("Error getting categories summary")

485 return {}