utils/get_embedding.py

import json
import os
import sys

import pysqlite3
__import__('pysqlite3')
sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')

import chromadb
from chromadb.utils import embedding_functions

import openai

# OpenAI API 키 설정
openai.api_key = os.getenv("OPENAI_API_KEY")

# ChromaDB 클라이언트 초기화 (영구 저장소 사용)
client = chromadb.PersistentClient(path="./data/chroma_db")

# OpenAI 임베딩 함수 설정
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
                api_key=openai.api_key,
                model_name="text-embedding-3-small"
            )

# 컬렉션 이름 설정
COLLECTION_NAME = "busan_data_navigation"

def get_chroma():
    """
    로컬에 저장된 ChromaDB 컬렉션을 로드합니다.
    """
    try:
        collection = client.get_collection(name=COLLECTION_NAME, embedding_function=openai_ef)
        print(f"컬렉션 '{COLLECTION_NAME}'을 성공적으로 로드했습니다.")
        print(f"현재 컬렉션에는 {collection.count()} 개의 문서가 있습니다.")
        return collection
    except ValueError as e:
        print(f"오류: 컬렉션 '{COLLECTION_NAME}'을 로드하는 데 실패했습니다.")
        print(f"오류 메시지: {str(e)}")
        print("컬렉션이 존재하지 않거나 손상되었을 수 있습니다.")
        return None

# 메인 실행 부분
if __name__ == "__main__":
    # ChromaDB 컬렉션 로드
    collection = get_chroma()
    
    if collection:
        # 쿼리 실행 (테스트용)
        results = collection.query(
            query_texts=["부산광역시 강서구 어업 면허에 대해 알려주세요"],
            n_results=10
        )
        with open('./data/id_to_metadata.json', 'r') as f:
            id_to_metadata= json.load(f)
        titles = [id_to_metadata[_id]['title'].strip() for _id in results['ids'][0]]

        print("쿼리 결과:\n", '\n'.join(titles))
    else:
        print("컬렉션을 로드할 수 없어 쿼리를 실행할 수 없습니다.")