GenAI 기반 검색시스템 구축Building a GenAI-Powered Search System

# 이 코드는 Streamlit 버전의 원본 코드입니다.
# ADK 버전으로의 변환은 아래 가이드를 참고하세요.
import streamlit as st
import requests
import re
import pandas as pd
from google.auth import default, exceptions as google_auth_exceptions
from google.auth.transport.requests import Request
from google.cloud import bigquery
from langchain_google_genai import ChatGoogleGenerativeAI
import os
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import logging

# ------------------------------
# 로깅 설정
# ------------------------------
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# ------------------------------
# 환경 설정
# ------------------------------
# 중요: 실제 운영 환경에서는 API 키와 같은 민감한 정보를 코드에 직접 하드코딩하지 마세요.
# Streamlit Secrets, 환경 변수, 또는 Google Secret Manager와 같은 안전한 방법을 사용하세요.
os.environ["GOOGLE_API_KEY"] = "YOUR_API_KEY"
PROJECT_NUMBER = "YOUR_PROJECT_NUMBER"
ENGINE_ID = "YOUR_ENGINE_ID"
BIGQUERY_TABLE = "YOUR_PROJECT_ID.YOUR_DATASET.YOUR_TABLE"

DISCOVERY_ENDPOINT = f"https://discoveryengine.googleapis.com/v1alpha/projects/{PROJECT_NUMBER}/locations/global/collections/default_collection/engines/{ENGINE_ID}/servingConfigs/default_search:search"

llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash-001", streaming=True)

# Streamlit 세션 상태 초기화
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "last_authors" not in st.session_state:
    st.session_state.last_authors = None
if "last_doc_ids" not in st.session_state:
    st.session_state.last_doc_ids = set()
if "context_docs" not in st.session_state:
    st.session_state.context_docs = []
if "last_search_query" not in st.session_state:
    st.session_state.last_search_query = ""

# ------------------------------
# 유틸 함수
# ------------------------------
def extract_id(question: str):
    match = re.search(r'id[=\s:]+([a-f0-9\-]+)', question, re.IGNORECASE)
    return match.group(1) if match else None

def extract_doi_or_url(question: str):
    doi_match = re.search(r'(10\.\d{4,9}/[^\s]+)', question)
    if doi_match: return doi_match.group(1)
    url_match = re.search(r'https://(?:www\.)?biorxiv\.org/content/([^\s]+)', question)
    if url_match: return url_match.group(1)
    return None

def extract_keywords_and_years(question: str):
    years = re.findall(r'(19|20)\d{2}', question)
    year = years[0] if years else None
    words = re.findall(r'\b([A-Z][a-zA-Z\-\'.]+(?:\s[A-Z][a-zA-Z\-\'.]+)*)\b', question)
    keyword = " ".join(words) if words else None
    return keyword, year

def extract_keywords_from_question(question: str) -> str:
    keywords = re.findall(r'\b\w+\b', question)
    return " ".join([word for word in keywords if len(word) > 2])

# ------------------------------
# 필터 함수
# ------------------------------
def filter_relevant_documents(df: pd.DataFrame, question: str) -> pd.DataFrame:
    logging.info(f"필터링 시작. 입력 행 개수: {len(df)}")
    if df.empty: return pd.DataFrame()

    keyword, year = extract_keywords_and_years(question)
    
    if not keyword and not year: return df

    def row_matches(row):
        keyword_match = True
        if keyword:
            keyword_match = any(keyword.lower() in str(row.get(col, "")).lower() for col in ["title", "abstract", "authors"])
        
        year_match = True
        if year:
            year_match = str(year) in str(row.get("date", ""))
            
        return keyword_match and year_match

    return df[df.apply(row_matches, axis=1)]

# ------------------------------
# 검색 함수
# ------------------------------
@st.cache_data(ttl=3600)
def retrieve_documents_as_dataframe(question: str, top_k: int = 20):
    logging.info(f"Discovery Engine 검색: '{question}'")
    try:
        credentials, _ = default(scopes=["https://www.googleapis.com/auth/cloud-platform"])
        credentials.refresh(Request())
    except Exception as e:
        st.error(f"❌ Google Cloud 인증 실패: {e}")
        return pd.DataFrame()

    headers = {"Authorization": f"Bearer {credentials.token}", "Content-Type": "application/json"}
    body = {
        "pageSize": top_k,
        "query": question,
        "languageCode": "ko",
        "queryExpansionSpec": {"condition": "AUTO"},
        "spellCorrectionSpec": {"mode": "AUTO"}
    }

    try:
        response = requests.post(DISCOVERY_ENDPOINT, headers=headers, json=body)
        response.raise_for_status()
        data = response.json()
        rows = [doc.get("document", {}).get("structData", {}) for doc in data.get("results", [])]
        logging.info(f"Discovery Engine 검색 성공. 결과 {len(rows)}개")
        return pd.DataFrame(rows)
    except requests.exceptions.RequestException as e:
        st.error(f"❌ Discovery Engine 검색 실패: {e}")
        return pd.DataFrame()

# ------------------------------
# BigQuery 직접 쿼리 함수
# ------------------------------
@st.cache_data(ttl=3600)
def query_bigquery_by_id(doc_id: str) -> pd.DataFrame:
    logging.info(f"BigQuery 쿼리: ID='{doc_id}'")
    try:
        bq_client = bigquery.Client()
        query = f"SELECT * FROM `{BIGQUERY_TABLE}` WHERE id = @doc_id"
        job_config = bigquery.QueryJobConfig(query_parameters=[bigquery.ScalarQueryParameter("doc_id", "STRING", doc_id)])
        return bq_client.query(query, job_config=job_config).to_dataframe()
    except Exception as e:
        st.error(f"❌ BigQuery 쿼리 오류: {e}")
        return pd.DataFrame()

# ------------------------------
# 관련 논문 찾기 함수
# ------------------------------
def preprocess_text(text):
    return re.sub(r'[^\w\s]', '', str(text)).lower() if text else ""

def find_related_documents(target_doc, all_docs_df: pd.DataFrame, top_k=5):
    if target_doc is None or all_docs_df.empty: return pd.DataFrame()

    all_docs_df["combined_text"] = all_docs_df["title"].fillna("") + " " + all_docs_df["abstract"].fillna("")
    all_docs_df["combined_text"] = all_docs_df["combined_text"].apply(preprocess_text)
    target_text = preprocess_text(target_doc.get("title", "") + " " + target_doc.get("abstract", ""))
    
    tfidf = TfidfVectorizer()
    all_docs_tfidf = tfidf.fit_transform(all_docs_df["combined_text"])
    target_tfidf = tfidf.transform([target_text])
    
    similarities = cosine_similarity(target_tfidf, all_docs_tfidf).flatten()
    related_indices = similarities.argsort()[::-1]
    related_indices = [i for i in related_indices if all_docs_df.iloc[i].get("id") != target_doc.get("id")]
    
    return all_docs_df.iloc[related_indices[:top_k]]

# ------------------------------
# 복합 질의 처리 함수
# ------------------------------
def handle_complex_query(question: str, context_df: pd.DataFrame) -> pd.DataFrame:
    if "관련된 논문" in question.lower() and not context_df.empty:
        target_doc = context_df.iloc[0].to_dict()
        search_query = target_doc.get("title", "")
        all_docs_df = retrieve_documents_as_dataframe(search_query)
        return find_related_documents(target_doc, all_docs_df)
    
    return filter_relevant_documents(context_df, question)

# ------------------------------
# LLM 응답 생성
# ------------------------------
def generate_answer_stream(question: str, docs: list):
    context = "\n\n---\n\n".join([str(doc) for doc in docs])
    history = "\n".join([f"Q: {turn['question']}\nA: {turn['answer']}" for turn in st.session_state.chat_history])
    prompt = f"""
대화 기록:
{history}

참고 문서:
{context}

위 대화 기록과 참고 문서를 바탕으로 다음 질문에 대해 상세하고 친절하게 답변해 주세요.
질문: {question}
답변:
"""
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""
        for chunk in llm.stream(prompt):
            full_response += chunk.content
            message_placeholder.markdown(full_response + "▌")
        message_placeholder.markdown(full_response)
    return full_response

# ------------------------------
# Streamlit UI
# ------------------------------
st.set_page_config(page_title="QA (Discovery Engine)", layout="wide")
st.title("🔍 VertexAI 기반 BigQuery QA")
question_input = st.text_input("질문을 입력하세요", placeholder="예: 논문 ID, 저자, 키워드 등으로 검색")

if question_input:
    with st.spinner("🔎 검색 및 LLM 분석 중..."):
        extracted_id = extract_id(question_input)
        is_complex = st.session_state.chat_history and any(k in question_input.lower() for k in ["다른", "관련된"])

        if extracted_id:
            filtered_df = query_bigquery_by_id(extracted_id)
        elif is_complex:
            context_df = pd.DataFrame(st.session_state.context_docs)
            filtered_df = handle_complex_query(question_input, context_df)
        else:
            df = retrieve_documents_as_dataframe(question_input)
            filtered_df = filter_relevant_documents(df, question_input)
            st.session_state.last_search_query = question_input
        
        if not filtered_df.empty:
            st.session_state.context_docs = filtered_df.to_dict('records')
        
        docs = st.session_state.context_docs
        with st.chat_message("user"):
            st.markdown(question_input)
        
        answer = generate_answer_stream(question_input, docs)
        st.session_state.chat_history.append({"question": question_input, "answer": answer, "df": filtered_df})

# 대화 히스토리 표시
if st.session_state.chat_history:
    st.subheader("💬 대화 히스토리")
    for i, turn in enumerate(reversed(st.session_state.chat_history)):
        with st.expander(f"#{len(st.session_state.chat_history)-i} {turn['question'][:60]}..."):
            with st.chat_message("user"): st.markdown(turn["question"])
            with st.chat_message("assistant"): st.markdown(turn['answer'])
            
            df = turn.get("df")
            if df is not None and not df.empty:
                st.markdown("#### 📊 검색 결과")
                st.dataframe(df, use_container_width=True)