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 ..security import url_authority_without_userinfo

13from ..utilities.sql_utils import escape_like as _escape_like

14from ..utilities.json_utils import extract_json, get_llm_response_text

15from .models import DomainClassification

16from ..utilities.llm_utils import invoke_llm_sync

19# Predefined categories for domain classification

20DOMAIN_CATEGORIES = {

21 "Academic & Research": [

22 "University/Education",

23 "Scientific Journal",

24 "Research Institution",

25 "Academic Database",

26 ],

27 "News & Media": [

28 "General News",

29 "Tech News",

30 "Business News",

31 "Entertainment News",

32 "Local/Regional News",

33 ],

34 "Reference & Documentation": [

35 "Encyclopedia",

36 "Technical Documentation",

37 "API Documentation",

38 "Tutorial/Guide",

39 "Dictionary/Glossary",

40 ],

41 "Social & Community": [

42 "Social Network",

43 "Forum/Discussion",

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

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

46 "Personal Blog", # Individual author blogs

47 ],

48 "Business & Commerce": [

49 "E-commerce",

50 "Corporate Website",

51 "B2B Platform",

52 "Financial Service",

53 "Marketing/Advertising",

54 ],

55 "Technology": [

56 "Software Development",

57 "Cloud Service",

58 "Open Source Project",

59 "Tech Company",

60 "Developer Tools",

61 ],

62 "Government & Organization": [

63 "Government Agency",

64 "Non-profit",

65 "International Organization",

66 "Think Tank",

67 "Industry Association",

68 ],

69 "Entertainment & Lifestyle": [

70 "Streaming Service",

71 "Gaming",

72 "Sports",

73 "Arts & Culture",

74 "Travel & Tourism",

75 ],

76 "Professional & Industry": [

77 "Healthcare",

78 "Legal",

79 "Real Estate",

80 "Manufacturing",

81 "Energy & Utilities",

82 ],

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

84}

87class DomainClassifier:

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

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

91 """Initialize the domain classifier.

93 Args:

94 username: Username for database session

95 settings_snapshot: Settings snapshot for LLM configuration

96 """

97 self.username = username

98 self.settings_snapshot = settings_snapshot

99 self.llm = None

100

101 def _get_llm(self):

102 """Get or initialize LLM instance."""

103 if self.llm is None:

104 # Thread the user so the LLM PEP enforces local-only policy: the

105 # route builds settings_snapshot with a bare SettingsManager (no

106 # `_username`), and classification prompts carry the user's own

107 # resource titles, which must not fail open to a cloud LLM when

108 # the user's default engine is a private retriever.

109 self.llm = get_llm(

110 settings_snapshot=self.settings_snapshot,

111 username=self.username,

112 )

113 return self.llm

114

115 def close(self) -> None:

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

117 from ..utilities.resource_utils import safe_close

118

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

120

121 def _get_domain_samples(

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

123 ) -> List[Dict]:

124 """Get sample resources from a domain.

125

126 Args:

127 domain: Domain to get samples for

128 session: Database session

129 limit: Maximum number of samples

130

131 Returns:

132 List of resource samples

133 """

134 resources = (

135 session.query(ResearchResource)

136 .filter(

137 ResearchResource.url.like(

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

139 )

140 )

141 .limit(limit)

142 .all()

143 )

144

145 samples = []

146 for resource in resources:

147 samples.append(

148 {

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

150 "url": resource.url,

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

152 if resource.content_preview

153 else None,

154 }

155 )

156

157 return samples

158

159 def _build_classification_prompt(

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

161 ) -> str:

162 """Build prompt for LLM classification.

163

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

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

166 accurate classification based on actual site content.

167

168 Args:

169 domain: Domain to classify

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

171

172 Returns:

173 Formatted prompt string

174 """

175 # Format categories for prompt

176 categories_text = []

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

178 subcats_text = ", ".join(subcats)

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

180

181 # Format samples

182 samples_text = []

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

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

185 if sample.get("preview"):

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

187

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

189

190Domain: {domain}

191

192Sample content from this domain:

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

194

195Available Categories:

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

197

198Respond with a JSON object containing:

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

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

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

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

203

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

205

206JSON Response:"""

207

208 def classify_domain(

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

210 ) -> Optional[DomainClassification]:

211 """Classify a single domain using LLM.

212

213 Args:

214 domain: Domain to classify

215 force_update: If True, reclassify even if already exists

216

217 Returns:

218 DomainClassification object or None if failed

219 """

220 try:

221 with get_user_db_session(self.username) as session:

222 # Check if already classified

223 existing = (

224 session.query(DomainClassification)

225 .filter_by(domain=domain)

226 .first()

227 )

228

229 if existing and not force_update:

230 logger.info(

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

232 )

233 return existing

234

235 # Get sample resources

236 samples = self._get_domain_samples(domain, session)

237

238 # Build prompt and get classification

239 prompt = self._build_classification_prompt(domain, samples)

240 llm = self._get_llm()

241

242 response = invoke_llm_sync(llm, prompt)

243 response_text = get_llm_response_text(response)

244

245 result = extract_json(response_text, expected_type=dict)

246 if result is None:

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

248

249 # Create or update classification

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

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

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

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

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

255 existing.sample_titles = json.dumps(

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

257 )

258 existing.sample_count = len(samples)

259 classification = existing

260 else:

261 classification = DomainClassification(

262 domain=domain,

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

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

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

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

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

268 sample_count=len(samples),

269 )

270 session.add(classification)

271

272 session.commit()

273 logger.info(

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

275 )

276 return classification

277

278 except Exception as e:

279 # Type name only: this frame holds the sample resource URLs,

280 # which a traceback's locals or the exception text can echo.

281 logger.warning(

282 f"Error classifying domain {domain} ({type(e).__name__})"

283 )

284 return None

285

286 def classify_all_domains(

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

288 ) -> Dict:

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

290

291 Args:

292 force_update: If True, reclassify all domains

293 progress_callback: Optional callback function for progress updates

294

295 Returns:

296 Dictionary with classification results

297 """

298 results = {

299 "total": 0,

300 "classified": 0,

301 "failed": 0,

302 "skipped": 0,

303 "domains": [],

304 }

305

306 try:

307 with get_user_db_session(self.username) as session:

308 # Get all unique domains

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

310 domains = set()

311

312 for (url,) in resources:

313 if url:

314 try:

315 # Authority minus userinfo (empty when that

316 # is ambiguous): a stored URL's user:pass@

317 # would otherwise land in the logs, the

318 # progress payload and the persisted

319 # DomainClassification key.

320 domain = url_authority_without_userinfo(url).lower()

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

322 domain = domain[4:]

323 if domain:

324 domains.add(domain)

325 except (ValueError, AttributeError):

326 logger.debug("Skipping malformed URL")

327 continue

328

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

330 logger.info(

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

332 )

333

334 # Classify each domain ONE BY ONE

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

336 logger.info(

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

338 )

339

340 if progress_callback:

341 progress_callback(

342 {

343 "current": i,

344 "total": results["total"],

345 "domain": domain,

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

347 }

348 )

349

350 try:

351 # Check if already classified

352 if not force_update:

353 existing = (

354 session.query(DomainClassification)

355 .filter_by(domain=domain)

356 .first()

357 )

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

359 results["skipped"] += 1

360 results["domains"].append(

361 {

362 "domain": domain,

363 "status": "skipped",

364 "category": existing.category,

365 "subcategory": existing.subcategory,

366 }

367 )

368 logger.info(

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

370 )

371 continue

372

373 # Classify this single domain

374 classification = self.classify_domain(

375 domain, force_update

376 )

377

378 if classification:

379 results["classified"] += 1

380 results["domains"].append(

381 {

382 "domain": domain,

383 "status": "classified",

384 "category": classification.category,

385 "subcategory": classification.subcategory,

386 "confidence": classification.confidence,

387 }

388 )

389 logger.info(

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

391 )

392 else:

393 results["failed"] += 1

394 results["domains"].append(

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

396 )

397 logger.warning(

398 f"Failed to classify domain {domain}"

399 )

400

401 except Exception as e:

402 # Type name only: this frame holds the raw

403 # resource URLs, which a traceback can echo.

404 logger.warning(

405 f"Error classifying domain {domain} "

406 f"({type(e).__name__})"

407 )

408 results["failed"] += 1

409 results["domains"].append(

410 {

411 "domain": domain,

412 "status": "failed",

413 "error": "Classification failed",

414 }

415 )

416

417 logger.info(

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

419 )

420 return results

421

422 except Exception as e:

423 # Type name only: this frame holds the raw resource URLs.

424 logger.warning(

425 f"Error in classify_all_domains ({type(e).__name__})"

426 )

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

428 return results

429

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

431 """Get existing classification for a domain.

432

433 Args:

434 domain: Domain to look up

435

436 Returns:

437 DomainClassification object or None if not found

438 """

439 try:

440 with get_user_db_session(self.username) as session:

441 return (

442 session.query(DomainClassification)

443 .filter_by(domain=domain)

444 .first()

445 )

446 except Exception:

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

448 return None

449

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

451 """Get all domain classifications.

452

453 Returns:

454 List of all DomainClassification objects

455 """

456 try:

457 with get_user_db_session(self.username) as session:

458 return (

459 session.query(DomainClassification)

460 .order_by(

461 DomainClassification.category,

462 DomainClassification.domain,

463 )

464 .all()

465 )

466 except Exception:

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

468 return []

469

470 def get_categories_summary(self) -> Dict:

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

472

473 Returns:

474 Dictionary with category counts and domains

475 """

476 try:

477 with get_user_db_session(self.username) as session:

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

479

480 summary = {}

481 for classification in classifications:

482 cat = classification.category

483 if cat not in summary:

484 summary[cat] = {

485 "count": 0,

486 "domains": [],

487 "subcategories": {},

488 }

489

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

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

492 {

493 "domain": classification.domain,

494 "subcategory": classification.subcategory,

495 "confidence": classification.confidence,

496 }

497 )

498

499 subcat = classification.subcategory

500 if subcat:

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

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

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

504

505 return summary

506

507 except Exception:

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

509 return {}