structverify 0.3.0__py3-none-any.whl

structverify/retrieval/custom_csv_source.py

@@ -0,0 +1,386 @@
+"""structverify.retrieval.custom_csv_source — Custom CSV DataSource.
+
+회사 업로드 CSV를 BaseDataSource 인터페이스로 노출.
+가정 CSV 형태: indicator,year,region,value,unit (행 = 시점별 값).
+
+레퍼런스: retrieval/kosis_source.py (KOSISDataSource).
+search_catalog: 키워드 매칭 (임베딩/DB 없음). fetch_evidence: 다음 단계.
+"""
+from __future__ import annotations
+
+import csv
+import os
+import re
+from typing import Any
+
+from .base import BaseDataSource, CatalogCandidate, EvidenceData
+from .registry import register_datasource
+from structverify.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+# CSV 컬럼 → 표준 필드 기본 매핑 (config.column_mapping 으로 덮어쓰기)
+_DEFAULT_COLUMN_MAPPING = {
+    "indicator": "indicator",
+    "time_period": "year",
+    "region": "region",
+    "value": "value",
+    "unit": "unit",
+    "operator": "operator",   # (선택) 규정/한도 비교 연산자: <= >= < > = (없으면 일치 비교)
+}
+
+
+def _cosine(a: list[float], b: list[float]) -> float:
+    dot = sum(x * y for x, y in zip(a, b))
+    na = sum(x * x for x in a) ** 0.5
+    nb = sum(y * y for y in b) ** 0.5
+    return dot / (na * nb) if na and nb else 0.0
+
+
+@register_datasource("custom_csv")
+class CustomCSVDataSource(BaseDataSource):
+    """회사 업로드 CSV 데이터소스.
+
+    config 예 (config/default.yaml의 data_sources.custom_csv):
+        base_path: "./uploads/custom"
+        catalog_json: "tables.csv"          # CSV 파일명 (base_path 기준)
+        column_mapping: {value: "val", ...} # CSV 컬럼명 커스텀 (선택)
+        csv_path: "/abs/path.csv"           # (테스트 편의) CSV 경로 직접 주입
+    """
+
+    name = "custom_csv"
+
+    def __init__(self, **config: Any):
+        self.config = config
+        self.base_path = config.get("base_path", "")
+        self.catalog_json = config.get("catalog_json", "")
+        # 테스트 편의: csv_path 직접 주입 우선. 없으면 base_path/catalog_json로 해석.
+        self.csv_path = config.get("csv_path") or self._resolve_csv_path()
+        # CSV 컬럼명 매핑 — config 우선, 없으면 기본
+        self.column_mapping = {
+            **_DEFAULT_COLUMN_MAPPING,
+            **(config.get("column_mapping") or {}),
+        }
+        # ── 임베딩 검색 (지표가 방대할 때) ──
+        # embedding config가 있고 distinct 지표 수가 임계 초과면 임베딩 의미검색을 쓴다.
+        # (지표가 적으면 키워드 매칭으로 충분 — 임베딩 비용 회피.)
+        self._embed_cfg = config.get("embedding") or {}
+        self._embed_threshold = int(config.get("embed_threshold", 200))
+        self._embed_client = None
+        self._ind_index: list[dict] | None = None   # [{key, vec}]
+        logger.info(
+            f"[CustomCSVDataSource] 초기화: csv_path={self.csv_path!r}, "
+            f"column_mapping={self.column_mapping}"
+        )
+
+    # ── CSV 헬퍼 ──
+
+    def _resolve_csv_path(self) -> str:
+        """base_path/catalog_json 으로 CSV 경로 해석."""
+        if self.base_path and self.catalog_json:
+            return os.path.join(self.base_path, self.catalog_json)
+        return self.catalog_json or self.base_path or ""
+
+    def _read_rows(self) -> list[dict[str, str]]:
+        """CSV 전체 행을 dict 리스트로. 경로 없거나 못 읽으면 []."""
+        if not self.csv_path or not os.path.exists(self.csv_path):
+            logger.warning(f"[CustomCSVDataSource] CSV 경로 없음: {self.csv_path!r}")
+            return []
+        try:
+            with open(self.csv_path, encoding="utf-8") as f:
+                return list(csv.DictReader(f))
+        except Exception as e:  # noqa: BLE001
+            logger.warning(f"[CustomCSVDataSource] CSV 읽기 실패: {e}")
+            return []
+
+    @staticmethod
+    def _match_score(query: str, indicator: str) -> float:
+        """query↔indicator 매칭 점수. 한국어 띄어쓰기/접두어 변이에 견고.
+
+        완전일치 1.0 / 부분포함 0.8 / 공백제거 부분포함 0.75 /
+        부분-토큰 겹침 비례(≤0.6) / 없으면 0.0.
+
+        핵심: 한국어는 띄어쓰기가 흔들린다("고객 수" vs "고객수")·수식어가 붙는다
+        ("누적 고객 수" vs "총 고객수"). 그래서 토큰 완전일치가 아니라 *부분 문자열*
+        토큰 겹침(고객 ⊂ 고객수)까지 인정한다. (검색 리콜 우선 — 정밀도는 랭커/fetch가.)
+        """
+        q, ind = query.strip().lower(), indicator.strip().lower()
+        if not q or not ind:
+            return 0.0
+        if q == ind:
+            return 1.0
+        if q in ind or ind in q:
+            return 0.8
+        # 공백 제거 후 부분포함 ("고객 수" → "고객수" ⊂ "총고객수")
+        qn, indn = q.replace(" ", ""), ind.replace(" ", "")
+        if qn and indn and (qn in indn or indn in qn):
+            return 0.75
+        # 부분-토큰 겹침: query 토큰이 indicator 토큰에 (부분)포함되면 hit
+        qtok = re.findall(r"\w+", q)
+        itok = re.findall(r"\w+", ind)
+        if qtok and itok:
+            hits = sum(1 for t in qtok if any(t in it or it in t for it in itok))
+            if hits:
+                return min(0.6, 0.6 * hits / len(qtok))
+        return 0.0
+
+    # ── 임베딩 검색 (지표 방대 시) ──
+
+    def _distinct_indicators(self) -> list[tuple[str, str]]:
+        ind_c = self.column_mapping["indicator"]
+        reg_c = self.column_mapping["region"]
+        seen: set = set()
+        out: list[tuple[str, str]] = []
+        for r in self._read_rows():
+            i = (r.get(ind_c) or "").strip()
+            if not i:
+                continue
+            reg = (r.get(reg_c) or "").strip()
+            k = (i, reg)
+            if k not in seen:
+                seen.add(k)
+                out.append(k)
+        return out
+
+    def _embed_enabled(self) -> bool:
+        """임베딩 의미검색 사용 여부.
+
+        config.use_embedding 으로 명시 제어(기본 auto):
+          · "true"/"on"/"embedding" → 항상 임베딩(임베딩 config가 있으면)
+          · "false"/"off"/"keyword" → 항상 키워드
+          · "auto"(기본)            → 지표가 embed_threshold 초과면 임베딩, 아니면 키워드
+        """
+        if not self._embed_cfg:
+            return False
+        mode = str(self.config.get("use_embedding", "auto")).strip().lower()
+        if mode in ("false", "off", "no", "keyword"):
+            return False
+        if mode in ("true", "on", "yes", "embedding"):
+            return True
+        return len(self._distinct_indicators()) > self._embed_threshold
+
+    def _index_cache_path(self, texts: list[str]) -> str | None:
+        """임베딩 인덱스 디스크 캐시 경로. 지표 집합+모델이 바뀌면 키가 달라져 자동 재구축.
+
+        비활성화: embedding.cache=false. 위치: embedding.cache_dir (기본 임시디렉토리).
+        """
+        import hashlib
+        import os
+        import tempfile
+        if self._embed_cfg.get("cache") is False:
+            return None
+        base = self._embed_cfg.get("cache_dir") or os.path.join(
+            tempfile.gettempdir(), "structverify_emb"
+        )
+        model = f"{self._embed_cfg.get('provider', '')}:{self._embed_cfg.get('model') or 'default'}"
+        _ident = getattr(self, "table", None) or getattr(self, "csv_path", None) or self.name
+        key = hashlib.md5(
+            ("|".join(sorted(texts)) + "|" + model + "|" + str(_ident)).encode("utf-8")
+        ).hexdigest()[:16]
+        return os.path.join(base, f"{self.name}_{key}.json")
+
+    async def _ensure_ind_index(self) -> None:
+        if self._ind_index is not None:
+            return
+        import json
+        import os
+        from structverify.utils.embedding_client import EmbeddingClient
+        if self._embed_client is None:
+            self._embed_client = EmbeddingClient(self._embed_cfg)
+        pairs = self._distinct_indicators()
+        texts = [f"{i} ({r})" if r else i for i, r in pairs]
+
+        # ── 디스크 캐시 로드 (첫 사용 후 재실행/새 프로세스에서 즉시 로드) ──
+        cache_path = self._index_cache_path(texts)
+        if cache_path and os.path.exists(cache_path):
+            try:
+                with open(cache_path, encoding="utf-8") as f:
+                    self._ind_index = json.load(f)
+                logger.info(
+                    f"[{self.name}] 지표 임베딩 인덱스 캐시 로드 {len(self._ind_index)}개 "
+                    f"({cache_path})"
+                )
+                return
+            except Exception as e:  # noqa: BLE001 — 캐시 손상 시 재구축
+                logger.warning(f"[{self.name}] 인덱스 캐시 로드 실패 → 재구축: {e}")
+
+        self._ind_index = []
+        # 배치 임베딩 — 지표가 200+개여도 직렬 N회가 아니라 배치 몇 회로(속도·429 방지).
+        _CHUNK = 96   # upstage/openai 배치 입력 상한 여유
+        for k in range(0, len(texts), _CHUNK):
+            chunk_pairs = pairs[k:k + _CHUNK]
+            chunk_texts = texts[k:k + _CHUNK]
+            try:
+                vecs = await self._embed_client.embed_batch(chunk_texts, role="passage")
+            except Exception as e:  # noqa: BLE001
+                logger.warning(f"[{self.name}] 지표 배치 임베딩 실패(chunk {k}): {e}")
+                vecs = []
+            for (ind, reg), text, vec in zip(chunk_pairs, chunk_texts, vecs or []):
+                if vec:
+                    self._ind_index.append({"id": f"{ind}|{reg}", "name": text, "vec": vec})
+        logger.info(
+            f"[{self.name}] 지표 임베딩 인덱스 {len(self._ind_index)}개 "
+            f"(배치 {(len(texts) + _CHUNK - 1) // _CHUNK}회, 첫 구축)"
+        )
+        # ── 디스크 캐시 저장 (다음 실행부터 즉시 로드) ──
+        if cache_path and self._ind_index:
+            try:
+                os.makedirs(os.path.dirname(cache_path), exist_ok=True)
+                with open(cache_path, "w", encoding="utf-8") as f:
+                    json.dump(self._ind_index, f)
+                logger.info(f"[{self.name}] 지표 임베딩 인덱스 캐시 저장 → {cache_path}")
+            except Exception as e:  # noqa: BLE001 — 저장 실패는 무시(다음에 재구축)
+                logger.debug(f"[{self.name}] 인덱스 캐시 저장 실패: {e}")
+
+    async def _search_embedding(self, query: str, top_k: int) -> list[CatalogCandidate]:
+        await self._ensure_ind_index()
+        qv = await self._embed_client.embed(query, role="query")
+        scored = sorted(
+            ((_cosine(qv, it["vec"]), it) for it in (self._ind_index or [])),
+            key=lambda x: x[0], reverse=True,
+        )
+        cands = [{"id": it["id"], "name": it["name"], "score": float(s)} for s, it in scored[:top_k]]
+        logger.info(f"[CustomCSVDataSource] 임베딩 검색(query={query!r}): {len(cands)}개 후보")
+        return cands
+
+    # ── BaseDataSource 인터페이스 ──
+
+    async def search_catalog(
+        self,
+        query: str,
+        category: list[str] | None = None,
+        top_k: int = 10,
+        context: dict[str, Any] | None = None,
+    ) -> list[CatalogCandidate]:
+        # 지표가 방대하면 임베딩 의미검색, 아니면 키워드 매칭.
+        if self._embed_enabled():
+            return await self._search_embedding(query, top_k)
+        """CSV indicator 컬럼을 query로 키워드 매칭 → 후보 반환.
+
+        같은 (indicator, region)은 시점별 여러 행이 있어도 후보 1개로 묶음.
+        반환: [{"id": "<indicator>|<region>", "name": "<indicator> (<region>)",
+                "score": <매칭점수>}] — score 내림차순, 매칭 없으면 [].
+        """
+        ind_col = self.column_mapping["indicator"]
+        reg_col = self.column_mapping["region"]
+
+        # (indicator, region) 단위로 묶어 표 1개 = 후보 1개
+        groups: dict[tuple[str, str], float] = {}
+        for row in self._read_rows():
+            indicator = (row.get(ind_col) or "").strip()
+            if not indicator:
+                continue
+            region = (row.get(reg_col) or "").strip()
+            score = self._match_score(query, indicator)
+            if score <= 0.0:
+                continue
+            groups.setdefault((indicator, region), score)  # 같은 그룹 첫 점수 유지
+
+        candidates: list[CatalogCandidate] = [
+            {
+                "id": f"{ind}|{reg}",
+                "name": f"{ind} ({reg})" if reg else ind,
+                "score": sc,
+            }
+            for (ind, reg), sc in groups.items()
+        ]
+        candidates.sort(key=lambda c: c["score"], reverse=True)
+        logger.info(
+            f"[CustomCSVDataSource] search_catalog(query={query!r}): "
+            f"{len(candidates)}개 후보"
+        )
+        return candidates[:top_k]
+
+    @staticmethod
+    def _year_key(year: Any) -> int:
+        """연도 정렬용 키. 파싱 실패 시 -1 (가장 뒤)."""
+        try:
+            return int(str(year).strip())
+        except (TypeError, ValueError):
+            return -1
+
+    async def fetch_evidence(
+        self,
+        candidate_id: str,
+        params: dict[str, Any] | None = None,
+        workspace: Any = None,   # 인터페이스 호환 (CSV는 캐시 없음 — 받되 무시)
+    ) -> EvidenceData | None:
+        """candidate_id("지표|지역") + params["time_period"]로 CSV 행 매칭 → EvidenceData.
+
+        time_period 매칭: year 컬럼과 문자열 정규화 비교 ("2024" == 2024).
+        time_period 없으면 가장 최근 연도 행. 맞는 행 없으면 None.
+        """
+        params = params or {}
+        indicator, _, region = candidate_id.partition("|")  # "지표|지역" 분해
+        indicator, region = indicator.strip(), region.strip()
+
+        ind_col = self.column_mapping["indicator"]
+        reg_col = self.column_mapping["region"]
+        year_col = self.column_mapping["time_period"]
+        val_col = self.column_mapping["value"]
+        unit_col = self.column_mapping["unit"]
+
+        # 해당 indicator(+region) 행만 추림
+        rows = [
+            r for r in self._read_rows()
+            if (r.get(ind_col) or "").strip() == indicator
+            and (r.get(reg_col) or "").strip() == region
+        ]
+        if not rows:
+            return None
+
+        tp = params.get("time_period")
+        tp_str = str(tp).strip() if tp is not None else ""
+        if tp_str:
+            matched = next(
+                (r for r in rows if str(r.get(year_col, "")).strip() == tp_str),
+                None,
+            )
+            # 이 지표에 해당 연도 행이 없을 때 — 두 경우로 나뉜다:
+            if matched is None:
+                #  (b) 그 연도가 *소스의 다른 지표엔* 존재 → 연도별 데이터가 있다는 뜻.
+                #      이 표(연도 없는 누적 등)로 폴백하면 누적값을 특정연도값인 척
+                #      씌우게 되므로 거부(None) → 에이전트가 '연간 …' 같은 연도별
+                #      지표 후보를 시도하게 유도. (총 주문건수 vs 연간 주문건수 혼동 방지.)
+                tp_in_source = any(
+                    str(r.get(year_col, "")).strip() == tp_str
+                    for r in self._read_rows()
+                )
+                if tp_in_source:
+                    logger.info(
+                        f"[{self.name}] time_period={tp_str!r}이 소스 다른 지표엔 존재 "
+                        f"→ 이 표('{indicator}') 부적합, fetch 거부(연도별 후보 유도)"
+                    )
+                    return None
+                #  (a) 그 연도가 소스 어디에도 없음(예: "2023년 말 기준" as-of 시점,
+                #      정답은 연도 없는 누적 총계) → 최근/유일 행으로 폴백.
+                matched = max(rows, key=lambda r: self._year_key(r.get(year_col)))
+                logger.info(
+                    f"[{self.name}] time_period={tp_str!r} 소스에 없음 → "
+                    f"누적/최근 행 폴백 (year={matched.get(year_col)!r})"
+                )
+        else:
+            matched = max(rows, key=lambda r: self._year_key(r.get(year_col)))
+
+        if matched is None:
+            return None
+
+        try:
+            value = float((matched.get(val_col) or "").strip())
+        except (TypeError, ValueError):
+            logger.warning(
+                f"[CustomCSVDataSource] value 파싱 실패: {matched.get(val_col)!r}"
+            )
+            return None
+
+        op_col = self.column_mapping.get("operator", "operator")
+        return {
+            "value": value,
+            "unit": (matched.get(unit_col) or "").strip(),
+            "time_period": str(matched.get(year_col, "")).strip(),
+            "operator": (matched.get(op_col) or "").strip(),   # 규정 비교 연산자 (없으면 "")
+            "source": self.name,   # custom_csv / custom_db (하위클래스)
+            "indicator": indicator,
+            "region": region,
+            "matched_row": matched,
+        }