src/embedding_client.py

"""
embedding_client.py — 语义向量召回客户端
支持两种模式：
1. API 模式（DeepSeek / OpenAI / Qwen 兼容嵌入接口）
2. 本地 sentence-transformers 模式（可选，需安装）

优先级：API → sentence-transformers → token cosine fallback
"""
from __future__ import annotations

import json
import os
import re
from pathlib import Path
from typing import Optional

try:
    from dotenv import load_dotenv
except ModuleNotFoundError:  # python-dotenv is optional; env vars still work.
    def load_dotenv(*args, **kwargs):
        return False


load_dotenv()


class EmbeddingClient:
    """统一嵌入接口。未配置时自动回退。"""

    def __init__(self) -> None:
        self.api_key = os.getenv("EMBEDDING_API_KEY") or os.getenv("LLM_API_KEY")
        self.base_url = os.getenv("EMBEDDING_BASE_URL") or os.getenv("LLM_BASE_URL")
        self.model = os.getenv("EMBEDDING_MODEL", "text-embedding-v1")
        self.provider = os.getenv("EMBEDDING_PROVIDER", "").lower()
        self._client = None
        self._local_model = None
        self.mode: Optional[str] = None  # "api" | "local" | None

        # 尝试加载 API 模式
        if self.api_key and self.base_url:
            try:
                from openai import OpenAI
                self._client = OpenAI(api_key=self.api_key, base_url=self.base_url)
                self.mode = "api"
                print(f"[OK] EmbeddingClient: API 模式 ({self.base_url})")
            except Exception as exc:
                print(f"[WARN] EmbeddingClient API 初始化失败：{exc}")

        # API 失败则尝试本地模型  
        if self.mode is None:
            self._try_load_local()

    def _try_load_local(self) -> None:
        try:
            from sentence_transformers import SentenceTransformer
            model_name = os.getenv("LOCAL_EMBEDDING_MODEL", "BAAI/bge-small-zh-v1.5")
            self._local_model = SentenceTransformer(model_name)
            self.mode = "local"
            print(f"[OK] EmbeddingClient: 本地模式 ({model_name})")
        except Exception:
            self.mode = None

    def is_available(self) -> bool:
        return self.mode is not None

    def embed(self, texts: list[str]) -> Optional[list[list[float]]]:
        """批量获取嵌入向量。返回 List[List[float]] 或 None。"""
        if self.mode == "api":
            return self._embed_api(texts)
        if self.mode == "local":
            return self._embed_local(texts)
        return None

    def _embed_api(self, texts: list[str]) -> Optional[list[list[float]]]:
        try:
            resp = self._client.embeddings.create(model=self.model, input=texts)
            return [d.embedding for d in resp.data]
        except Exception as exc:
            print(f"[WARN] Embedding API 调用失败：{exc}")
            return None

    def _embed_local(self, texts: list[str]) -> Optional[list[list[float]]]:
        try:
            embeddings = self._local_model.encode(texts, normalize_embeddings=True)
            return embeddings.tolist()
        except Exception as exc:
            print(f"[WARN] 本地嵌入失败：{exc}")
            return None

    def cosine_similarity(self, vec1: list[float], vec2: list[float]) -> float:
        """余弦相似度。"""
        if not vec1 or not vec2:
            return 0.0
        dot = sum(a * b for a, b in zip(vec1, vec2))
        norm1 = sum(a * a for a in vec1) ** 0.5
        norm2 = sum(b * b for b in vec2) ** 0.5
        if norm1 == 0 or norm2 == 0:
            return 0.0
        return dot / (norm1 * norm2)

    def semantic_similarity(self, resume_text: str, job_text: str) -> Optional[float]:
        """计算简历与岗位的语义相似度（0~1）。"""
        embeddings = self.embed([resume_text, job_text])
        if embeddings is None:
            return None
        return self.cosine_similarity(embeddings[0], embeddings[1])

    def rank_by_semantic(
        self, resume_text: str, jobs: list[dict], top_k: int = 8
    ) -> list[tuple[int, float]]:
        """用语义向量对岗位重新排序。返回 [(index, score), ...] 按分数降序。"""
        if not self.is_available():
            return []
        texts = [f"{j.get('title', '')} {j.get('jd', '')}" for j in jobs]
        all_texts = [resume_text] + texts
        embeddings = self.embed(all_texts)
        if embeddings is None:
            return []
        resume_emb = embeddings[0]
        scores: list[tuple[int, float]] = []
        for i, job_emb in enumerate(embeddings[1:], start=0):
            sim = self.cosine_similarity(resume_emb, job_emb)
            scores.append((i, sim))
        scores.sort(key=lambda x: x[1], reverse=True)
        return scores[:top_k]


def get_embedding_client() -> EmbeddingClient:
    """工厂函数。"""
    return EmbeddingClient()