96 lines
2.7 KiB
Python
96 lines
2.7 KiB
Python
import json
|
|
import uuid
|
|
import redis
|
|
from app.config import settings
|
|
from app.models.db import SessionLocal, AnalysisTask, Match
|
|
from app.core.analyzer import DocumentAnalyzer
|
|
from app.celery_app import celery_app
|
|
from datetime import datetime
|
|
|
|
_analyzer: DocumentAnalyzer | None = None
|
|
_redis: redis.Redis | None = None
|
|
|
|
|
|
def get_analyzer() -> DocumentAnalyzer:
|
|
global _analyzer
|
|
if _analyzer is None:
|
|
_analyzer = DocumentAnalyzer()
|
|
return _analyzer
|
|
|
|
|
|
def get_redis() -> redis.Redis:
|
|
global _redis
|
|
if _redis is None:
|
|
_redis = redis.from_url(settings.REDIS_URL)
|
|
return _redis
|
|
|
|
|
|
def _publish(task_id: str, stage: str, progress: int):
|
|
try:
|
|
get_redis().publish(
|
|
f"progress:{task_id}",
|
|
json.dumps({"stage": stage, "progress": progress}),
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
@celery_app.task(bind=True, name="tasks.analyze_document")
|
|
def analyze_document(self, text: str, task_id: str):
|
|
db = SessionLocal()
|
|
try:
|
|
analyzer = get_analyzer()
|
|
|
|
_publish(task_id, "fingerprint", 10)
|
|
exact = analyzer._level_1_exact_match(text)
|
|
|
|
_publish(task_id, "fuzzy", 40)
|
|
fuzzy = analyzer._level_2_fuzzy_match(text)
|
|
|
|
_publish(task_id, "semantic", 70)
|
|
semantic = analyzer._level_3_semantic_match(text)
|
|
|
|
matches = exact + fuzzy + semantic
|
|
overall_similarity = analyzer._calculate_overall_similarity(matches)
|
|
|
|
_publish(task_id, "saving", 90)
|
|
|
|
task = db.query(AnalysisTask).filter(AnalysisTask.id == task_id).first()
|
|
if not task:
|
|
return
|
|
|
|
task.status = "completed"
|
|
task.overall_similarity = overall_similarity
|
|
task.completed_at = datetime.utcnow()
|
|
|
|
for match in matches[:settings.TOP_K_RESULTS]:
|
|
raw_id = match.get("source_id")
|
|
try:
|
|
doc_id = uuid.UUID(str(raw_id)) if raw_id is not None else None
|
|
except (ValueError, TypeError, AttributeError):
|
|
doc_id = None
|
|
|
|
db.add(Match(
|
|
task_id=task_id,
|
|
doc_id=doc_id,
|
|
similarity=float(match["similarity"]),
|
|
method=match["method"],
|
|
fragment_start=match.get("fragment_start"),
|
|
fragment_end=match.get("fragment_end"),
|
|
fragment_text=match.get("fragment_text"),
|
|
))
|
|
|
|
db.commit()
|
|
_publish(task_id, "done", 100)
|
|
|
|
except Exception as exc:
|
|
task = db.query(AnalysisTask).filter(AnalysisTask.id == task_id).first()
|
|
if task:
|
|
task.status = "failed"
|
|
db.commit()
|
|
_publish(task_id, "failed", 0)
|
|
raise self.retry(exc=exc, countdown=5, max_retries=2)
|
|
|
|
finally:
|
|
db.close()
|