84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
import json
|
|
from typing import List, Union
|
|
|
|
import chromadb
|
|
from chromadb.utils import embedding_functions
|
|
|
|
# === 경로 설정 (모두 로컬) ===
|
|
EMBEDDING_MODEL_PATH = "./models/ko-sroberta-multitask"
|
|
|
|
# 2. 벡터 DB 설정
|
|
persist_directory = "./chroma_db"
|
|
chroma_client = chromadb.PersistentClient(path=persist_directory)
|
|
|
|
# 임베딩 함수 설정
|
|
embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
|
|
model_name=EMBEDDING_MODEL_PATH, # 로컬 폴더 경로 가능
|
|
device="cpu",
|
|
normalize_embeddings=True
|
|
)
|
|
|
|
# 컬렉션 생성 또는 가져오기
|
|
collection = chroma_client.get_or_create_collection(
|
|
name="orgchart",
|
|
embedding_function=embedding_fn,
|
|
metadata={"hnsw:space": "cosine"}
|
|
)
|
|
|
|
|
|
def init(sessionId: str, data: List[Union[str, dict]]):
|
|
"""
|
|
데이터를 벡터 DB에 초기화(저장)합니다.
|
|
|
|
Args:
|
|
sessionId (str): 세션 ID
|
|
data (List[Union[str, dict]]): 저장할 데이터 리스트 (문자열 또는 딕셔너리)
|
|
"""
|
|
print(f'{sessionId} init start')
|
|
|
|
# 문서 ID 생성
|
|
doc_ids = [f"{sessionId}_{i}" for i in range(len(data))]
|
|
|
|
# 데이터 처리: 문자열이면 그대로, 객체면 JSON 문자열로 변환
|
|
documents = []
|
|
|
|
doc_list = []
|
|
for q in data:
|
|
# 각 필드를 안전하게 추출 (None 방어)
|
|
name = q.get('name') or ""
|
|
dept = q.get('deptNm') or ""
|
|
grade = q.get('gradeNm') or ""
|
|
position = q.get('ptsnNm') or ""
|
|
office_phone = q.get('ofcePhn') or ""
|
|
mobile_phone = q.get('mblPhn') or ""
|
|
chief_name = q.get('chiefNm') or ""
|
|
state_code = q.get('state') or ""
|
|
|
|
# 상태 코드 한글화
|
|
state_map = {'C': '재직', 'T': '퇴사', 'H': '휴직'}
|
|
status = state_map.get(state_code, "정보없음")
|
|
|
|
# [핵심] 검색 엔진이 좋아할만한 서술형 문장 생성
|
|
# 부서명과 이름을 앞부분에 배치하여 가중치 유도
|
|
if name == '':
|
|
doc = (
|
|
f"부서: {dept}. "
|
|
f"해당 {dept}의 부서장은 {chief_name}입니다."
|
|
)
|
|
else:
|
|
doc = (
|
|
f"부서: {dept}. 이름: {name}. {dept} 소속의 {name} {grade}입니다. "
|
|
f"직위는 {position}이며 현재 {status} 중입니다. "
|
|
f"사내 전화번호(사선)는 {office_phone}입니다."
|
|
)
|
|
doc_list.append(doc)
|
|
|
|
# 벡터 DB에 추가
|
|
collection.add(
|
|
documents=doc_list,
|
|
ids=doc_ids,
|
|
metadatas=[{"sessionId": sessionId, 'deptCd': d.get('deptCd') or "", 'gradeCd': d.get('gradeCd') or ""} for d in data]
|
|
)
|
|
|
|
print(f'{sessionId} init end')
|