42 lines
1.5 KiB
Python
42 lines
1.5 KiB
Python
from langchain.embeddings.base import Embeddings
|
|
from typing import List, Dict, Any
|
|
import requests
|
|
import os
|
|
|
|
embedding_host = os.getenv("EMBEDDING_HOST", "http://183.111.96.67:31740")
|
|
# embedding_host = os.getenv("EMBEDDING_HOST", "http://vector-module-01jsg2688r20ddyj62jmkafavg-groupuser")
|
|
|
|
class WeaviateCustomEmbeddings(Embeddings):
|
|
"""커스텀 엔드포인트를 사용하는 임베딩 모델"""
|
|
|
|
def __init__(self):
|
|
self.api_url = embedding_host.rstrip("/")
|
|
print(self.api_url)
|
|
|
|
def embed_documents(self, texts: List[str]) -> List[List[float]]:
|
|
"""문서 리스트를 임베딩합니다."""
|
|
embeddings = []
|
|
for text in texts:
|
|
embeddings.append(self.embed_query(text))
|
|
return embeddings
|
|
|
|
def embed_query(self, text: str) -> List[float]:
|
|
"""쿼리 텍스트를 임베딩합니다."""
|
|
query_text = text.question if hasattr(text, 'question') else str(text)
|
|
|
|
response = requests.post(
|
|
f"{self.api_url}/embed",
|
|
json={"inputs": query_text},
|
|
headers={
|
|
"Authorization": "Bearer 01k1vwr8fs0wnwcd3y0pskcbah"
|
|
}
|
|
)
|
|
if response.status_code != 200:
|
|
raise ValueError(f"임베딩 API 호출 실패: {self.api_url}/embed, {query_text}, {response.status_code}, {response.text}")
|
|
|
|
vectors = response.json()
|
|
|
|
# 응답은 [[...]] 구조 → 첫 번째 벡터만 반환
|
|
return vectors[0]
|
|
#return response.json()["vector"]
|