""" 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}", "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'}")