Neo4j와 Haystack으로
강력한 검색을 구현하다

# 텍스트 임베딩을 위한 Haystack + OpenAI 초기화
def initialize_haystack():
    # 문서 저장소 초기화 (지금은 인메모리, 다른 저장소로 구성 가능)
    document_store = InMemoryDocumentStore()
    # 텍스트 임베딩 생성을 위한 OpenAITextEmbedder 초기화
    embedder = OpenAITextEmbedder(
        api_key=Secret.from_env_var("OPENAI_API_KEY"),
        model="text-embedding-ada-002"
    )
    return embedder

# Neo4j에서 영화 줄거리와 제목을 가져온다
def retrieve_movie_plots():
    # 각 Movie 노드의 "title", "overview", "tmdbId" 프로퍼티를 가져오는 질의
    query = """
    MATCH (m:Movie)
    WHERE m.embedding IS NULL
    RETURN m.tmdbId AS tmdbId, m.title AS title, m.overview AS overview
    """
    with driver.session() as session:
        results = session.run(query)
        # 각 영화의 제목·줄거리(overview)·ID를 movies 리스트에 담는다
        movies = [
            {"tmdbId": row["tmdbId"],
             "title": row["title"],
             "overview": row["overview"]}
            for row in results
        ]
    return movies

# ThreadPoolExecutor를 이용한 병렬 임베딩 생성
def generate_and_store_embeddings(embedder, movies, max_workers=10):
    results_to_store = []
    def process_movie(movie):
        title = movie.get("title", "Unknown Title")
        overview = str(movie.get("overview", "")).strip()
        tmdbId = movie.get("tmdbId")
        if not overview:
            print(f"Skipping {title} — No overview available.")
            return None
        try:
            print(f"Generating embedding for: {title}")
            embedding_result = embedder.run(overview)
            embedding = embedding_result.get("embedding")
            if embedding:
                return (tmdbId, embedding)
            else:
                print(f"No embedding generated for: {title}")
        except Exception as e:
            print(f"Error processing {title}: {e}")
        return None

# 임베딩을 Neo4j에 다시 저장한다
def store_embedding_in_neo4j(tmdbId, embedding):
    query = """
    MATCH (m:Movie {tmdbId: $tmdbId})
    SET m.embedding = $embedding
    """
    with driver.session() as session:
        session.run(query, tmdbId=tmdbId, embedding=embedding)
    print(f"Stored embedding for TMDB ID: {tmdbId}")

# Neo4j에 저장된 임베딩 검증
def verify_embeddings():
    query = """
    MATCH (m:Movie)
    WHERE exists(m.embedding)
    RETURN m.title, m.embedding
    LIMIT 10
    """
    with driver.session() as session:
        results = session.run(query)
        for record in results:
            title = record["title"]
            embedding = np.array(record["embedding"])[:5]
            print(f" {title}: {embedding}...")

def create_or_reset_vector_index():
    with driver.session() as session:
        try:
            # 기존 벡터 인덱스가 있으면 삭제
            session.run("DROP INDEX overview_embeddings IF EXISTS")
            print("Old index dropped")
        except:
            print("No index to drop")
        # embedding 프로퍼티에 새 벡터 인덱스 생성
        print("Creating new vector index")
        query_index = """
        CREATE VECTOR INDEX overview_embeddings IF NOT EXISTS
        FOR (m:Movie) ON (m.embedding)
        OPTIONS {indexConfig: {
            `vector.dimensions`: 1536,
            `vector.similarity_function`: 'cosine'}}
        """
        session.run(query_index)
        print("Vector index created successfully")

text_embedder = OpenAITextEmbedder(
    api_key=Secret.from_env_var("OPENAI_API_KEY"),
    model="text-embedding-ada-002"
)
# 1단계: 질의의 임베딩 생성
query_embedding = text_embedder.run(query).get("embedding")

if query_embedding is None:
    print("Query embedding not created successfully.")
    return

print("Query embedding created successfully.")

# 2단계: 질의 임베딩으로 유사 문서 검색
similar_documents = document_store.query_by_embedding(
    query_embedding, top_k=3
)
if not similar_documents:
    print("No similar documents found.")
    return
print(f"Found {len(similar_documents)} similar documents.")

# 3단계: 결과 표시
for doc in similar_documents:
    title = doc.meta.get("title", "N/A")
    overview = doc.meta.get("overview", "N/A")
    score = doc.score
    print(f"Title: {title}\nOverview: {overview}\n"
          f"Score: {score:.2f}\n{'-'*40}")

cypher_query = """
    CALL db.index.vector.queryNodes("overview_embeddings",
        $top_k, $query_embedding)
    YIELD node AS movie, score
    MATCH (movie:Movie)
    RETURN movie.title AS title, movie.overview AS overview, score
"""

text_embedder = OpenAITextEmbedder(
    api_key=Secret.from_env_var("OPENAI_API_KEY"),
    model="text-embedding-ada-002"
)

retriever = Neo4jDynamicDocumentRetriever(
    client_config=client_config,
    runtime_parameters=["query_embedding"],
    compose_doc_from_result=True,
    verify_connectivity=True,
)
pipeline = Pipeline()
pipeline.add_component("query_embedder", text_embedder)
pipeline.add_component("retriever", retriever)
pipeline.connect(
    "query_embedder.embedding", "retriever.query_embedding"
)
result = pipeline.run(
    {
        "query_embedder": {"text": query},
        "retriever": {
            "query": cypher_query,
            "parameters": {
                "index": "overview_embeddings", "top_k": 3
            },
        },
    }
)

# 리트리버 결과에서 문서 추출
documents = result["retriever"]["documents"]
for doc in documents:
    # 문서 메타데이터에서 제목·개요 추출
    title = doc.meta.get("title", "N/A")
    overview = doc.meta.get("overview", "N/A")
    # 문서에서 점수 추출
    score = getattr(doc, "score", None)
    score_display = f"{score:.2f}" if score is not None else "N/A"
    print(f"Title: {title}\nOverview: {overview}\n"
          f"Score: {score_display}\n{'-'*40}\n")

import gradio as gr

# Gradio 챗봇 인터페이스 정의
def chatbot(user_input):
    return perform_vector_search_cypher(user_input)

# Gradio 인터페이스 생성
chat_interface = gr.Interface(
    fn=chatbot,
    inputs=gr.Textbox(
        placeholder="What kind of movie would you like to watch?",
        lines=3,
        label="Your movie preference"
    ),
    outputs=gr.Textbox(label="Recommendations", lines=12),
    title="AI Movie Recommendation System",
    description="Ask me about movies! I can recommend movies
        based on your preferences.",
    examples=[
        ["I want to watch a sci-fi movie with time travel"],
        ["Recommend me a romantic comedy with a happy ending"],
        ["I'm in the mood for something with superheroes
          but not too serious"],
        ["I want a thriller that keeps me on the edge of my seat"],
        ["Show me movies about artificial intelligence
          taking over the world"]
    ],
    flagging_mode="never"
)

# Cypher 기반 검색과 Haystack을 쓰는 대화형 챗봇 핸들러
def perform_vector_search(query):
    print("MESSAGES RECEIVED:", user_input)
    cypher_query = """
        CALL db.index.vector.queryNodes("overview_embeddings",
            $top_k, $query_embedding)
        YIELD node AS movie, score
        MATCH (movie:Movie)
        RETURN movie.title AS title, movie.overview AS overview, score
    """
    # 임베더
    embedder = OpenAITextEmbedder(
        api_key=Secret.from_env_var("OPENAI_API_KEY"),
        model="text-embedding-ada-002"
    )
    # 리트리버
    retriever = Neo4jDynamicDocumentRetriever(
        client_config=client_config,
        runtime_parameters=["query_embedding"],
        compose_doc_from_result=True,
        verify_connectivity=True,
    )
    # 파이프라인
    pipeline = Pipeline()
    pipeline.add_component("query_embedder", embedder)
    pipeline.add_component("retriever", retriever)
    pipeline.connect(
        "query_embedder.embedding", "retriever.query_embedding"
    )

# 전체 과정을 오케스트레이션하는 메인 함수
def main():
    # 1단계: Neo4j AuraDB에 벡터 인덱스 생성 또는 리셋
    create_or_reset_vector_index()
    # 2단계: Gradio 챗봇 인터페이스 시작
    chat_interface.launch()

if __name__ == "__main__":
    main()

embedder = OpenAITextEmbedder(
    api_key=Secret.from_env_var("OPENAI_API_KEY"),
    model="text-embedding-babbage-001"  # 다른 모델로 실험
)

def create_additional_indexes():
    with driver.session() as session:
        session.run("CREATE INDEX IF NOT EXISTS movie_title_index
            FOR (m:Movie) ON (m.title)")
        session.run("CREATE INDEX IF NOT EXISTS movie_tmdbId_index
            FOR (m:Movie) ON (m.tmdbId)")
        print("Additional indexes created successfully")

import logging
logging.basicConfig(filename='chatbot_queries.log', level=logging.INFO)

def log_query(query):
    logging.info(f"User query: {query}")