exAichatbot_agent/exAiChatBot-chatbot2.0-agent/scripts/api_clients.py

"""
api_clients.py
────────────────────────────────────────────
외부 API 클라이언트 모듈: LLM, TEI Embedding, vLLM Reranker
"""

import os
import time
import json
from typing import List, Dict, Any, Optional, Tuple
import httpx
import numpy as np
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type


# ────────────────────────────────────────────
# 환경 변수 설정
# ────────────────────────────────────────────
# 신규 게이트웨이 방식: HTTPS + 호스트명 라우팅 + Bearer 토큰 + OpenAI 호환 포맷
# 예) LLM_BASE_URL=https://llm-ai.ex.co.kr  (클라이언트가 /v1/chat/completions 등을 붙임)
LLM_BASE_URL = os.getenv("LLM_BASE_URL", os.getenv("SGLANG_BASE_URL", "http://localhost:16000"))
LLM_MODEL_NAME = os.getenv("LLM_MODEL_NAME", "default")  # 모델 이름 (예: Qwen/Qwen3.6-27B-FP8)

TEI_EMBED_URL = os.getenv("TEI_EMBED_URL", "http://localhost:16001")
EMBED_MODEL_NAME = os.getenv("EMBED_MODEL_NAME", "Qwen/Qwen3-Embedding-8B")  # OpenAI 임베딩 model 필드
TEI_RERANK_URL = os.getenv("TEI_RERANK_URL", "http://localhost:16002")
TEI_RERANK_MODEL = os.getenv("TEI_RERANK_MODEL", "Qwen/Qwen3-Reranker-8B")  # 리랭커 model 필드
API_TIMEOUT = float(os.getenv("API_TIMEOUT", "60"))

# 공통 인증/TLS: 모든 모델 API가 동일 게이트웨이를 사용하므로 단일 키 사용 (LLM_API_KEY는 하위호환)
MODEL_API_KEY = os.getenv("MODEL_API_KEY", os.getenv("LLM_API_KEY", ""))
LLM_API_KEY = MODEL_API_KEY
# 자체서명 인증서면 검증 비활성화 (curl -k 와 동일). 기본 false
MODEL_VERIFY_SSL = os.getenv("MODEL_VERIFY_SSL", "false").strip().lower() in ("1", "true", "yes")
# openai: /v1/embeddings (게이트웨이) | tei: /embed (로컬 TEI·exdev 호스트 AI)
EMBED_API_STYLE = os.getenv("EMBED_API_STYLE", "openai").strip().lower()
# gateway: /score (queries+documents) | legacy: /rerank (query+documents, exdev reranker_server)
RERANK_API_STYLE = os.getenv("RERANK_API_STYLE", "gateway").strip().lower()


def _normalize_base_url(url: str) -> str:
    """URL 끝의 / 또는 /v1 접미사만 제거 (rstrip('/v1')는 포트 16001 등을 깨뜨림)."""
    url = url.rstrip("/")
    if url.endswith("/v1"):
        return url[:-3]
    return url


def _auth_headers() -> Dict[str, str]:
    """Bearer 인증 헤더 (키가 있을 때만)"""
    return {"Authorization": f"Bearer {MODEL_API_KEY}"} if MODEL_API_KEY else {}


# ────────────────────────────────────────────
# LLM Client (OpenAI 호환 Chat Completions)
# ────────────────────────────────────────────
class SGLangClient:
    """
    OpenAI 호환 LLM API 클라이언트
    
    지원 서비스:
    - SGLang (self-hosted)
    - vLLM (self-hosted)
    - Ollama (self-hosted)
    - 내부 LLM (OpenAI 호환 형식)
    - OpenAI API (API 키 필요)
    """
    
    def __init__(
        self, 
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        model_name: Optional[str] = None,
        timeout: float = API_TIMEOUT
    ):
        raw_url = base_url or LLM_BASE_URL
        self.base_url = _normalize_base_url(raw_url)
        self.model_name = model_name or LLM_MODEL_NAME
        self.timeout = timeout
        
        # API 키 설정 (있으면 헤더에 추가)
        api_key = api_key or MODEL_API_KEY
        headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
        if api_key:
            print(f"[LLM Client] API 키 인증 활성화")
        
        self.client = httpx.Client(timeout=timeout, headers=headers, verify=MODEL_VERIFY_SSL)
        print(f"[LLM Client] 초기화: {self.base_url}, model={self.model_name}, "
              f"timeout={timeout}s, verify_ssl={MODEL_VERIFY_SSL}")
    
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=10),
        retry=retry_if_exception_type((httpx.TimeoutException, httpx.ConnectError)),
    )
    def chat_completion(
        self,
        messages: List[Dict[str, str]],
        max_tokens: int = 64,
        temperature: float = 0.1,
        model: Optional[str] = None,
        **kwargs
    ) -> str:
        """
        Chat Completions API 호출 (OpenAI 호환 형식)
        
        Args:
            messages: [{"role": "system"|"user"|"assistant", "content": "..."}]
            max_tokens: 최대 생성 토큰 수
            temperature: 샘플링 온도 (0=deterministic)
            model: 모델 이름 (지정하지 않으면 초기화 시 설정한 모델 사용)
        
        Returns:
            생성된 텍스트
        """
        endpoint = f"{self.base_url}/v1/chat/completions"
        
        payload = {
            "model": model or self.model_name,
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "chat_template_kwargs": {"enable_thinking": False},
            **kwargs
        }
        
        try:
            response = self.client.post(endpoint, json=payload)
            response.raise_for_status()
            data = response.json()
            return data["choices"][0]["message"]["content"]
        except httpx.HTTPStatusError as e:
            print(f"[LLM Client] HTTP 오류: {e.response.status_code} - {e.response.text}")
            raise
        except Exception as e:
            print(f"[LLM Client] 오류: {e}")
            raise

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=10),
        retry=retry_if_exception_type((httpx.TimeoutException, httpx.ConnectError)),
    )
    def chat_completion_message(
        self,
        messages: List[Dict[str, Any]],
        max_tokens: int = 512,
        temperature: float = 0.1,
        model: Optional[str] = None,
        tools: Optional[List[Dict[str, Any]]] = None,
        tool_choice: Optional[Any] = "auto",
        **kwargs
    ) -> Dict[str, Any]:
        """Chat Completions → assistant message dict (content + tool_calls)."""
        endpoint = f"{self.base_url}/v1/chat/completions"
        payload: Dict[str, Any] = {
            "model": model or self.model_name,
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "chat_template_kwargs": {"enable_thinking": False},
            **kwargs,
        }
        if tools:
            payload["tools"] = tools
            payload["tool_choice"] = tool_choice

        response = self.client.post(endpoint, json=payload)
        try:
            response.raise_for_status()
        except httpx.HTTPStatusError as e:
            print(f"[LLM Client] tool-call HTTP 오류: {e.response.status_code} - {e.response.text}")
            raise
        data = response.json()
        return data["choices"][0]["message"]
    
    def __del__(self):
        if hasattr(self, 'client'):
            self.client.close()


# ────────────────────────────────────────────
# TEI Embedding Client
# ────────────────────────────────────────────
class TEIEmbeddingClient:
    """Text Embeddings Inference (TEI) Embedding API 클라이언트"""
    
    def __init__(self, base_url: str = TEI_EMBED_URL, model: str = EMBED_MODEL_NAME, timeout: float = API_TIMEOUT):
        self.base_url = _normalize_base_url(base_url)
        self.model = model
        self.timeout = timeout
        self.client = httpx.Client(timeout=timeout, headers=_auth_headers(), verify=MODEL_VERIFY_SSL)
        self.api_style = EMBED_API_STYLE
        print(f"[EmbeddingClient] 초기화: {self.base_url}, model={self.model}, "
              f"style={self.api_style}, timeout={timeout}s, verify_ssl={MODEL_VERIFY_SSL}")
    
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=10),
        retry=retry_if_exception_type((httpx.TimeoutException, httpx.ConnectError)),
    )
    def embed(self, texts: List[str], normalize: bool = True, is_query: bool = False) -> List[List[float]]:
        """
        텍스트 임베딩 생성 (OpenAI 호환 /v1/embeddings, Qwen3-Embedding-8B)
        
        Args:
            texts: 임베딩할 텍스트 리스트
            normalize: L2 정규화 여부 (코사인/Qdrant 호환). 서버 정규화에 의존하지 않고 클라이언트에서 수행
            is_query: True면 검색 질의(Query), False면 문서(Document)
        
        Returns:
            임베딩 벡터 리스트 [[dim], [dim], ...]
        
        Note:
            Qwen3-Embedding은 Query/Document를 구분합니다:
            - Query: "Instruct: ...\\nQuery: [질문]" 프리픽스
            - Document: 원문 그대로
        """
        # Qwen3-Embedding: Query에만 Instruct 문구 추가
        if is_query:
            processed_texts = [
                f"Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery: {text}"
                for text in texts
            ]
        else:
            processed_texts = list(texts)

        if self.api_style == "tei":
            endpoint = f"{self.base_url}/embed"
            payload = {"inputs": processed_texts, "normalize": normalize}
        else:
            endpoint = f"{self.base_url}/v1/embeddings"
            payload = {"model": self.model, "input": processed_texts}

        try:
            response = self.client.post(endpoint, json=payload)
            response.raise_for_status()
            data = response.json()

            if self.api_style == "tei":
                vectors = data if isinstance(data, list) else data.get("embeddings", data)
            elif isinstance(data, dict) and "data" in data:
                items = sorted(data["data"], key=lambda x: x.get("index", 0))
                vectors = [item["embedding"] for item in items]
            elif isinstance(data, list):
                vectors = data
            elif isinstance(data, dict) and "embeddings" in data:
                vectors = data["embeddings"]
            else:
                raise ValueError(f"예상치 못한 임베딩 응답 형식: {str(data)[:200]}")
            
            if normalize:
                vectors = self._l2_normalize(vectors)
            return vectors
        except httpx.HTTPStatusError as e:
            print(f"[EmbeddingClient] HTTP 오류: {e.response.status_code} - {e.response.text}")
            raise
        except Exception as e:
            print(f"[EmbeddingClient] 오류: {e}")
            raise

    @staticmethod
    def _l2_normalize(vectors: List[List[float]]) -> List[List[float]]:
        """L2 정규화 (코사인 유사도/Qdrant 호환). 0벡터는 그대로 둠."""
        arr = np.asarray(vectors, dtype="float32")
        norms = np.linalg.norm(arr, axis=1, keepdims=True)
        norms[norms == 0] = 1.0
        return (arr / norms).tolist()
    
    def __del__(self):
        if hasattr(self, 'client'):
            self.client.close()


# ────────────────────────────────────────────
# vLLM Reranker Client
# ────────────────────────────────────────────
class TEIRerankerClient:
    """
    vLLM 기반 Reranker API 클라이언트 (Qwen3-Reranker-8B)
    
    Note: 클래스명은 호환성을 위해 'TEIRerankerClient'로 유지하지만,
          실제로는 vLLM 기반 커스텀 reranker_server.py를 사용합니다.
    """
    
    def __init__(self, base_url: str = TEI_RERANK_URL, model: str = TEI_RERANK_MODEL, timeout: float = API_TIMEOUT):
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.timeout = timeout
        self.client = httpx.Client(timeout=timeout, headers=_auth_headers(), verify=MODEL_VERIFY_SSL)
        self.api_style = RERANK_API_STYLE
        print(f"[Reranker] 초기화: {self.base_url}, model={self.model}, "
              f"style={self.api_style}, timeout={timeout}s, verify_ssl={MODEL_VERIFY_SSL}")
    
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=10),
        retry=retry_if_exception_type((httpx.TimeoutException, httpx.ConnectError)),
    )
    def rerank(
        self,
        query: str,
        documents: List[str],
        top_k: Optional[int] = None,
        return_documents: bool = False
    ) -> List[Dict[str, Any]]:
        """
        문서 재랭킹 (vLLM 기반 Qwen3-Reranker-8B)
        
        Args:
            query: 검색 질의
            documents: 재랭킹할 문서 리스트
            top_k: 상위 k개만 반환 (None이면 전체)
            return_documents: 문서 텍스트도 함께 반환할지 여부
        
        Returns:
            [{"index": int, "score": float, "text": str (옵션)}, ...]
            score가 높은 순서대로 정렬됨
        
        Note:
            vLLM 기반 커스텀 reranker_server.py 사용
            - 내부적으로 "Yes/No" 확률을 계산
            - 요청 형식: {"query": "질문", "documents": ["문서1", "문서2", ...]}
            - 응답 형식: [{"index": int, "score": float}, ...] (score 내림차순 정렬)
        """
        if self.api_style == "legacy":
            endpoint = f"{self.base_url}/rerank"
            payload = {"query": query, "documents": documents}
        else:
            endpoint = f"{self.base_url}/score"
            payload = {
                "model": self.model,
                "queries": f"<Query>: {query}",
                "documents": documents,
            }
        
        try:
            print(f"[Reranker] 요청: endpoint={endpoint}, model={self.model}, query={query[:50]}..., documents={len(documents)}개")
            response = self.client.post(endpoint, json=payload)
            print(f"[Reranker] 응답 상태: {response.status_code}")
            response.raise_for_status()
            data = response.json()
            results = self._parse_scores(data, len(documents))
            print(f"[Reranker] 파싱 결과: {len(results)}개")
            if results:
                print(f"[Reranker] 첫 번째 결과: {results[0]}")
            return results
        except httpx.HTTPStatusError as e:
            print(f"[Reranker] HTTP 오류: {e.response.status_code} - {e.response.text}")
            raise
        except Exception as e:
            print(f"[Reranker] 오류: {e}")
            raise

    @staticmethod
    def _parse_scores(data: Any, num_docs: int) -> List[Dict[str, Any]]:
        """
        다양한 reranker 응답 형식을 [{"index": int, "score": float}, ...]로 정규화.
        지원: {"results": [...]}, {"data": [...]}, {"scores": [...]}, bare list 등
              항목은 {"index","score"} dict 또는 점수(float) 둘 다 허용.
        """
        if isinstance(data, dict):
            for key in ("results", "data", "scores", "output"):
                if key in data:
                    data = data[key]
                    break
        results: List[Dict[str, Any]] = []
        if isinstance(data, list):
            for i, item in enumerate(data):
                if isinstance(item, dict):
                    idx = item.get("index", item.get("idx", i))
                    score = item.get("score", item.get("relevance_score", item.get("logit")))
                    if score is None:
                        continue
                    results.append({"index": int(idx), "score": float(score)})
                elif isinstance(item, (int, float)):
                    results.append({"index": i, "score": float(item)})
        if not results:
            print(f"[Reranker] ⚠️ 점수 파싱 실패 (응답 형식 확인 필요): {str(data)[:200]}")
        return results
    
    def predict(self, pairs: List[Tuple[str, str]], batch_size: int = 16, show_progress_bar: bool = False) -> List[float]:
        """
        CrossEncoder 호환 인터페이스
        
        Args:
            pairs: [(query, doc), (query, doc), ...] 쌍 리스트
            batch_size: 배치 크기 (내부적으로 처리, 호환성 유지용)
            show_progress_bar: 진행률 표시 (미사용, 호환성 유지용)
        
        Returns:
            점수 리스트 [score1, score2, ...]
        """
        if not pairs:
            return []
        
        # 모든 쌍이 같은 query를 가진다고 가정 (일반적인 재랭킹 시나리오)
        query = pairs[0][0]
        documents = [doc for _, doc in pairs]
        
        results = self.rerank(query, documents, return_documents=False)
        
        # index 순서대로 정렬하여 점수 반환
        sorted_results = sorted(results, key=lambda x: x["index"])
        return [r["score"] for r in sorted_results]
    
    def __del__(self):
        if hasattr(self, 'client'):
            self.client.close()


# ────────────────────────────────────────────
# 헬스체크 유틸리티
# ────────────────────────────────────────────
def check_api_health() -> Dict[str, bool]:
    """
    모든 외부 API의 헬스체크 (게이트웨이 기준).
    게이트웨이에는 /health가 없을 수 있으므로 '응답이 오면(상태코드 무관) 도달 가능'으로 판단한다.
    """
    results = {}
    headers = _auth_headers()

    def _reachable(url: str) -> bool:
        try:
            # 어떤 HTTP 응답이든 오면 도달 가능으로 간주 (4xx 포함)
            httpx.get(url, headers=headers, timeout=5, verify=MODEL_VERIFY_SSL)
            return True
        except Exception:
            return False

    llm_base = _normalize_base_url(LLM_BASE_URL)
    results["llm"] = _reachable(f"{llm_base}/v1/models")
    embed_base = TEI_EMBED_URL.rstrip("/")
    results["tei_embed"] = _reachable(
        f"{embed_base}/health" if EMBED_API_STYLE == "tei" else embed_base
    )
    results["tei_rerank"] = _reachable(TEI_RERANK_URL.rstrip("/"))

    return results


if __name__ == "__main__":
    # 간단한 테스트
    print("=== API 클라이언트 테스트 ===")
    print(f"LLM API: {LLM_BASE_URL}")
    print(f"LLM Model: {LLM_MODEL_NAME}")
    print(f"TEI Embed: {TEI_EMBED_URL}")
    print(f"TEI Rerank: {TEI_RERANK_URL}")
    print()
    
    health = check_api_health()
    print("헬스체크 결과:")
    for service, status in health.items():
        print(f"  {service}: {'✅ OK' if status else '❌ FAIL'}")