test.py

# 인코딩 확인 및 설정
import sys
import io
import firebase_admin
from firebase_admin import credentials, firestore

import openai
import faiss
import numpy as np
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document
from langchain.docstore import InMemoryDocstore
from langchain_community.vectorstores import FAISS

# 표준 출력 인코딩 설정 (터미널의 인코딩을 UTF-8로 강제 설정)
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

# Firestore 설정
cred = credentials.Certificate(r'C:\Webcrawling\halmoney20-398d5-firebase-adminsdk-si4rf-8fad21bcda.json')
firebase_admin.initialize_app(cred)
db = firestore.client()

# OpenAI API 키 설정
openai.api_key = 'sk-FORak5TKpSAKo7lZfQku4GQF32CTsQlyZeb-e-_fQpT3BlbkFJQq6k9UsfHirVVkZsLVAMDa370xkiHmZHKN85HTp5oA'

# OpenAI 임베딩 함수
def embed_documents(documents):
    response = openai.Embedding.create(
        input=documents,
        model="text-embedding-ada-002"
    )
    embeddings = [item['embedding'] for item in response['data']]
    return embeddings

# Vector Store 구축 함수
def vector_store(all_answers):
    """
    직무명과 자소서 답변이 포함된 데이터를 벡터화하고 Vector DB를 생성하는 함수
    """
    print("문서를 분할하고 벡터화를 시작합니다...")

    # 텍스트 데이터 분할 (답변만 사용하여 분할)
    recur_text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=10,
        length_function=len
    )

    chunks = []
    metadata = []  # 직무명과 매핑된 메타데이터 저장용
    for job, answer in all_answers:
        chunked_texts = recur_text_splitter.split_text(answer)
        chunks.extend(chunked_texts)
        metadata.extend([{"job": job} for _ in chunked_texts])  # 직무 정보와 매칭

    # 임베딩 생성
    doc_embeddings = embed_documents(chunks)

    # FAISS 인덱스 생성 및 임베딩 추가
    index = faiss.IndexFlatL2(1536)
    index.add(np.array(doc_embeddings).astype(np.float32))

    # 문서 저장소 및 ID 매핑 설정
    documents = [Document(page_content=chunk, metadata=meta) for chunk, meta in zip(chunks, metadata)]
    docstore = InMemoryDocstore({str(i): doc for i, doc in enumerate(documents)})
    index_to_docstore_id = {i: str(i) for i in range(len(documents))}

    # 벡터 스토어 생성
    vector_store = FAISS(index=index, embedding_function=embed_documents,
                         docstore=docstore, index_to_docstore_id=index_to_docstore_id)
    
    print("문서를 벡터화하고 벡터 스토어에 저장을 완료했습니다.")
    return vector_store

# Firebase에서 직무별 답변 가져오기 함수
def get_cv_answers():
    """
    Firebase에서 직무별 자소서 데이터를 가져와서 직무명과 답변을 함께 반환하는 함수
    """
    doc_job = db.collection('coverLetter').stream()
    job_list = [job.id for job in doc_job]  # 직무 리스트
    
    all_answers = []  # 모든 직무의 자소서 답변 저장
    for jobCategory in job_list:
        cv_answers = []  # 직무별로 답변 리스트 초기화
        company_docs = db.collection('coverLetter').document(jobCategory).collection('company').stream()

        # 각 회사 문서에서 답변 가져와 리스트에 저장
        for doc in company_docs:
            data = doc.to_dict()
            if 'a1' in data:
                cv_answers.append((jobCategory, data['a1']))  # 직무명과 답변 저장
            if 'a2' in data:
                cv_answers.append((jobCategory, data['a2']))  # 직무명과 답변 저장

        all_answers.extend(cv_answers)  # 모든 답변을 한 리스트에 저장
    
    return all_answers

# LLM 호출 및 직무 관련 역량 키워드 추출
def LLM_abilities(query, vector_store):
    # 쿼리 벡터화
    query_embedding = embed_documents([query])[0]
    
    # 벡터 스토어에서 유사한 문서 검색
    D, I = vector_store.index.search(np.array([query_embedding]).astype(np.float32), k=5)

    # 검색된 문서 가져오기
    results = []
    for idx in I[0]:
        doc_id = vector_store.index_to_docstore_id[idx]
        document = vector_store.docstore.search(doc_id)
        results.append(document.page_content)
    
    combined_text = " ".join(results)  # 검색된 문서 결합
    
    # GPT-4에 검색된 문서와 쿼리 전달
    prompt = f"'{combined_text}'는 쿼리 답변에 참고할 직무별 합격자소서 데이터입니다. 이를 바탕으로 '{query}' 직무에 필요한 역량 키워드를 추출해줘"

    response = openai.ChatCompletion.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ],
        max_tokens=1000,
        temperature=0.7
    )

    if response and 'choices' in response:
        answer = response['choices'][0]['message']['content'].strip()
        print(f"GPT-4 응답: {answer}")
        return answer
    else:
        print("OpenAI API 응답이 유효하지 않습니다.")
        return None

# 실행 함수
def main():
    # Firebase에서 직무명과 답변을 가져옴
    all_answers = get_cv_answers()
    
    # 벡터 스토어 생성
    vector_db = vector_store(all_answers)
    
    # LLM을 통한 역량 추출 (예시 쿼리)
    query = "회계/세무 직무에 필요한 역량 키워드를 추출해줘. 설명없이"
    LLM_abilities(query, vector_db)

# 프로그램 실행
if __name__ == "__main__":
    main()