first

org_transformer_offline.py

@@ -0,0 +1,57 @@
+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 = []
+    for item in data:
+        if isinstance(item, str):
+            documents.append(item)
+        else:
+            documents.append(json.dumps(item, ensure_ascii=False))
+
+    # 벡터 DB에 추가
+    collection.add(
+        documents=documents,
+        ids=doc_ids,
+        metadatas=[{"sessionId": sessionId} for _ in doc_ids]
+    )
+
+    print(f'{sessionId} init end')