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 # Thread the user so the LLM PEP enforces local-only policy: the

103 # route builds settings_snapshot with a bare SettingsManager (no

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

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

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

107 self.llm = get_llm(

108 settings_snapshot=self.settings_snapshot,

109 username=self.username,

110 )

111 return self.llm

112

113 def close(self) -> None:

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

115 from ..utilities.resource_utils import safe_close

116

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

118

119 def _get_domain_samples(

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

121 ) -> List[Dict]:

122 """Get sample resources from a domain.

123

124 Args:

125 domain: Domain to get samples for

126 session: Database session

127 limit: Maximum number of samples

128

129 Returns:

130 List of resource samples

131 """

132 resources = (

133 session.query(ResearchResource)

134 .filter(

135 ResearchResource.url.like(

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

137 )

138 )

139 .limit(limit)

140 .all()

141 )

142

143 samples = []

144 for resource in resources:

145 samples.append(

146 {

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

148 "url": resource.url,

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

150 if resource.content_preview

151 else None,

152 }

153 )

154

155 return samples

156

157 def _build_classification_prompt(

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

159 ) -> str:

160 """Build prompt for LLM classification.

161

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

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

164 accurate classification based on actual site content.

165

166 Args:

167 domain: Domain to classify

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

169

170 Returns:

171 Formatted prompt string

172 """

173 # Format categories for prompt

174 categories_text = []

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

176 subcats_text = ", ".join(subcats)

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

178

179 # Format samples

180 samples_text = []

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

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

183 if sample.get("preview"):

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

185

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

187

188Domain: {domain}

189

190Sample content from this domain:

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

192

193Available Categories:

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

195

196Respond with a JSON object containing:

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

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

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

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

201

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

203

204JSON Response:"""

205

206 def classify_domain(

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

208 ) -> Optional[DomainClassification]:

209 """Classify a single domain using LLM.

210

211 Args:

212 domain: Domain to classify

213 force_update: If True, reclassify even if already exists

214

215 Returns:

216 DomainClassification object or None if failed

217 """

218 try:

219 with get_user_db_session(self.username) as session:

220 # Check if already classified

221 existing = (

222 session.query(DomainClassification)

223 .filter_by(domain=domain)

224 .first()

225 )

226

227 if existing and not force_update:

228 logger.info(

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

230 )

231 return existing

232

233 # Get sample resources

234 samples = self._get_domain_samples(domain, session)

235

236 # Build prompt and get classification

237 prompt = self._build_classification_prompt(domain, samples)

238 llm = self._get_llm()

239

240 response = llm.invoke(prompt)

241 response_text = get_llm_response_text(response)

242

243 result = extract_json(response_text, expected_type=dict)

244 if result is None:

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

246

247 # Create or update classification

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

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

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

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

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

253 existing.sample_titles = json.dumps(

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

255 )

256 existing.sample_count = len(samples)

257 classification = existing

258 else:

259 classification = DomainClassification(

260 domain=domain,

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

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

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

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

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

266 sample_count=len(samples),

267 )

268 session.add(classification)

269

270 session.commit()

271 logger.info(

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

273 )

274 return classification

275

276 except Exception:

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

278 return None

279

280 def classify_all_domains(

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

282 ) -> Dict:

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

284

285 Args:

286 force_update: If True, reclassify all domains

287 progress_callback: Optional callback function for progress updates

288

289 Returns:

290 Dictionary with classification results

291 """

292 results = {

293 "total": 0,

294 "classified": 0,

295 "failed": 0,

296 "skipped": 0,

297 "domains": [],

298 }

299

300 try:

301 with get_user_db_session(self.username) as session:

302 # Get all unique domains

303 from urllib.parse import urlparse

304

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

306 domains = set()

307

308 for (url,) in resources:

309 if url:

310 try:

311 parsed = urlparse(url)

312 domain = parsed.netloc.lower()

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

314 domain = domain[4:]

315 if domain:

316 domains.add(domain)

317 except (ValueError, AttributeError):

318 logger.debug("Skipping malformed URL")

319 continue

320

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

322 logger.info(

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

324 )

325

326 # Classify each domain ONE BY ONE

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

328 logger.info(

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

330 )

331

332 if progress_callback:

333 progress_callback(

334 {

335 "current": i,

336 "total": results["total"],

337 "domain": domain,

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

339 }

340 )

341

342 try:

343 # Check if already classified

344 if not force_update:

345 existing = (

346 session.query(DomainClassification)

347 .filter_by(domain=domain)

348 .first()

349 )

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

351 results["skipped"] += 1

352 results["domains"].append(

353 {

354 "domain": domain,

355 "status": "skipped",

356 "category": existing.category,

357 "subcategory": existing.subcategory,

358 }

359 )

360 logger.info(

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

362 )

363 continue

364

365 # Classify this single domain

366 classification = self.classify_domain(

367 domain, force_update

368 )

369

370 if classification:

371 results["classified"] += 1

372 results["domains"].append(

373 {

374 "domain": domain,

375 "status": "classified",

376 "category": classification.category,

377 "subcategory": classification.subcategory,

378 "confidence": classification.confidence,

379 }

380 )

381 logger.info(

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

383 )

384 else:

385 results["failed"] += 1

386 results["domains"].append(

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

388 )

389 logger.warning(

390 f"Failed to classify domain {domain}"

391 )

392

393 except Exception:

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

395 results["failed"] += 1

396 results["domains"].append(

397 {

398 "domain": domain,

399 "status": "failed",

400 "error": "Classification failed",

401 }

402 )

403

404 logger.info(

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

406 )

407 return results

408

409 except Exception:

410 logger.exception("Error in classify_all_domains")

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

412 return results

413

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

415 """Get existing classification for a domain.

416

417 Args:

418 domain: Domain to look up

419

420 Returns:

421 DomainClassification object or None if not found

422 """

423 try:

424 with get_user_db_session(self.username) as session:

425 return (

426 session.query(DomainClassification)

427 .filter_by(domain=domain)

428 .first()

429 )

430 except Exception:

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

432 return None

433

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

435 """Get all domain classifications.

436

437 Returns:

438 List of all DomainClassification objects

439 """

440 try:

441 with get_user_db_session(self.username) as session:

442 return (

443 session.query(DomainClassification)

444 .order_by(

445 DomainClassification.category,

446 DomainClassification.domain,

447 )

448 .all()

449 )

450 except Exception:

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

452 return []

453

454 def get_categories_summary(self) -> Dict:

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

456

457 Returns:

458 Dictionary with category counts and domains

459 """

460 try:

461 with get_user_db_session(self.username) as session:

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

463

464 summary = {}

465 for classification in classifications:

466 cat = classification.category

467 if cat not in summary:

468 summary[cat] = {

469 "count": 0,

470 "domains": [],

471 "subcategories": {},

472 }

473

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

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

476 {

477 "domain": classification.domain,

478 "subcategory": classification.subcategory,

479 "confidence": classification.confidence,

480 }

481 )

482

483 subcat = classification.subcategory

484 if subcat:

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

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

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

488

489 return summary

490

491 except Exception:

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

493 return {}