영리의 테크블로그

Vector DB (Chroma DB), OPENAI 을 활용한 일기 시스템 구현 본문

dev/AI

Vector DB (Chroma DB), OPENAI 을 활용한 일기 시스템 구현

영리0 2024. 10. 10. 03:42

OpenAI의 GPT-와 Chroma Vector DB를 사용해 일기 데이터를 저장하고 검색하는 시스템 구현

사용자가 작성한 일기에서 메타데이터(날짜, 장소, 활동, 감정, 동행인, 날씨)를 자동으로 추출한 뒤 Vector DB에 저장하는 방식으로 작동함.

 

작동 예시


  • 일기 입력 시 OpenAI의 GPT-3.5 모델을 사용해 일기에서 중요한 메타데이터를 자동으로 추출함.
  • 추출된 메타데이터와 일기 텍스트는 Vector DB에 임베딩되어 저장되었음.
  • 사용자가 입력한 질의에 맞는 일기 데이터를 검색하고, 검색된 일기의 메타데이터와 본문을 출력함.

메타데이터 추출 및 임베딩 저장

사용자가 일기를 입력하면, 일기 내용에서 날짜, 장소, 활동, 감정, 동행인, 날씨 같은 데이터를 추출.

⇒ 여기서 어떤거 추출할지 고민 필요

GPT-3.5 API를 사용하여 입력된 텍스트를 분석, JSON 형식으로 반환.

def extract_metadata(text):
    today = date.today().strftime("%Y-%m-%d")
    prompt = f"""
    다음 텍스트에서 날짜, 장소, 활동, 감정, 동행인, 날씨를 추출하세요. JSON 형식으로 반환하되, 날짜는 YYYY-MM-DD 형식으로 변환하세요. 오늘 날짜는 {today} 입니다.
    """

    response = openai_client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "user", "content": prompt}]
    )

    return json.loads(response.choices[0].message.content)

추출된 데이터와 일기 본문은 Chroma Vector DB에 임베딩과 함께 저장함.

def store_embedding_with_metadata(text):
    metadata = extract_metadata(text)

    diary_collection.add(
        documents=[text],
        metadatas=[metadata],
        ids=[datetime.datetime.now().isoformat()]
    )


일기 검색

사용자가 특정 검색어를 입력하면, 해당 검색어를 임베딩한 후,

Chroma Vector DB에서 메타데이터를 필터링하여 유사한 일기를 검색.

검색 결과는 날짜, 장소, 활동 등을 포함.

def query_diary(query_text, n_results=3):
    query_embedding = get_embedding(query_text)
    results = diary_collection.query(
        query_embeddings=[query_embedding],
        n_results=n_results
    )
    return results


메인 로직

def main():
    while True:
        print("\\\\n1. 일기 추가")
        print("2. 일기 검색")
        print("3. 종료")
        choice = input("선택하세요: ")

        if choice == '1':
            text = input("일기를 입력하세요: ")
            store_embedding_with_metadata(text)

        elif choice == '2':
            query = input("검색어를 입력하세요: ")
            results = query_diary(query)

            print("\\\\n검색 결과:")
            for i, (doc, metadata) in enumerate(zip(results['documents'][0], results['metadatas'][0])):
                print(f"\\\\n결과 {i+1}:")
                print(f"날짜: {metadata['date']}")
                print(f"장소: {metadata['location']}")
                print(f"활동: {metadata['activity']}")
                print(f"감정: {metadata['emotion']}")
                print(f"동행인: {metadata['companion']}")
                print(f"날씨: {metadata['weather']}")
                print(f"내용: {doc}")

        elif choice == '3':
            print("프로그램을 종료합니다.")
            break

        else:
            print("잘못된 선택입니다. 다시 선택해주세요.")


결론

날짜, 장소, 활동 등 어떠한 데이터를 넣어야 할지 고민 필요 !

감정에 대해서도 세분화 할 필요성 느낌