Spaces:
Running
Running
| from __future__ import annotations | |
| import os | |
| from datetime import datetime | |
| import requests | |
| class AirflowTriggerError(RuntimeError): | |
| pass | |
| class AirflowService: | |
| def __init__( | |
| self, | |
| base_url: str | None = None, | |
| username: str | None = None, | |
| password: str | None = None, | |
| ): | |
| self.base_url = (base_url or os.getenv("AIRFLOW_API_BASE_URL") or "http://localhost:8080").rstrip("/") | |
| self.username = username or os.getenv("AIRFLOW_API_USERNAME") or "admin" | |
| self.password = password or os.getenv("AIRFLOW_API_PASSWORD") or "admin" | |
| def trigger_pdf_ingestion( | |
| self, | |
| job_id: int, | |
| pdf_path: str, | |
| source_name: str, | |
| use_llm: bool, | |
| llm_model: str, | |
| ollama_base_url: str, | |
| ) -> dict: | |
| dag_id = "cert_study_pdf_ingestion" | |
| url = f"{self.base_url}/api/v1/dags/{dag_id}/dagRuns" | |
| run_id = f"cert_study_job_{job_id}_{datetime.utcnow().strftime('%Y%m%dT%H%M%S')}" | |
| payload = { | |
| "dag_run_id": run_id, | |
| "conf": { | |
| "job_id": job_id, | |
| "pdf_path": pdf_path, | |
| "source_name": source_name, | |
| "use_llm": use_llm, | |
| "llm_model": llm_model, | |
| "ollama_base_url": ollama_base_url, | |
| }, | |
| } | |
| try: | |
| response = requests.post( | |
| url, | |
| json=payload, | |
| auth=(self.username, self.password), | |
| timeout=10, | |
| ) | |
| except requests.RequestException as exc: | |
| raise AirflowTriggerError(f"Airflow API 연결 실패: {exc}") from exc | |
| if response.status_code >= 400: | |
| raise AirflowTriggerError( | |
| f"Airflow DAG trigger 실패 ({response.status_code}): {response.text[:500]}" | |
| ) | |
| return response.json() | |
| def trigger_visual_analysis( | |
| self, | |
| source_name: str | None = None, | |
| limit: int = 20, | |
| model: str | None = None, | |
| ) -> dict: | |
| dag_id = "cert_study_visual_analysis" | |
| url = f"{self.base_url}/api/v1/dags/{dag_id}/dagRuns" | |
| run_id = f"cert_study_visual_{datetime.utcnow().strftime('%Y%m%dT%H%M%S')}" | |
| payload = { | |
| "dag_run_id": run_id, | |
| "conf": { | |
| "source_name": source_name or "", | |
| "limit": int(limit), | |
| "model": model or os.getenv("OLLAMA_VISUAL_MODEL", "qwen3-vl:8b-instruct-q4_K_M"), | |
| }, | |
| } | |
| try: | |
| response = requests.post( | |
| url, | |
| json=payload, | |
| auth=(self.username, self.password), | |
| timeout=10, | |
| ) | |
| except requests.RequestException as exc: | |
| raise AirflowTriggerError(f"Airflow API 연결 실패: {exc}") from exc | |
| if response.status_code >= 400: | |
| raise AirflowTriggerError( | |
| f"Airflow 이미지 분석 DAG trigger 실패 ({response.status_code}): {response.text[:500]}" | |
| ) | |
| return response.json() | |
| def list_dag_runs(self, dag_id: str, limit: int = 5) -> list[dict]: | |
| url = f"{self.base_url}/api/v1/dags/{dag_id}/dagRuns" | |
| try: | |
| response = requests.get( | |
| url, | |
| params={"limit": max(int(limit), 50)}, | |
| auth=(self.username, self.password), | |
| timeout=8, | |
| ) | |
| except requests.RequestException as exc: | |
| raise AirflowTriggerError(f"Airflow API 연결 실패: {exc}") from exc | |
| if response.status_code >= 400: | |
| raise AirflowTriggerError( | |
| f"Airflow DAG 상태 조회 실패 ({response.status_code}): {response.text[:500]}" | |
| ) | |
| runs = response.json().get("dag_runs", []) | |
| return sorted( | |
| runs, | |
| key=lambda run: run.get("logical_date") or run.get("execution_date") or run.get("start_date") or "", | |
| reverse=True, | |
| )[: int(limit)] | |