# 🤖 벡터 검색 에이전트: 허깅페이스 허브를 백엔드로 하는 지능형 검색 엔진

_참조: [Martin Elstner](https://github.com/MartinEls)
_작성자: [안정](https://github.com/ahnjj)

검색 엔진은 크게 키워드 검색과 벡터 검색으로 나눌 수 있습니다.
키워드 검색과 달리, 벡터 검색을 사용할 때는 두 가지를 고려해야 합니다.
1. 적합한 임베딩 모델로 데이터셋과 쿼리를 임베딩하는 작업
2. 임베딩 데이터를 처리할 수 있는 DB

하지만 임베딩 값을 기반으로 한 벡터 검색만으로는 ‘사용자가 원하는 답변’을 보장하기 어렵습니다.

그래서 검색 과정에서 에이전트가 자율적으로 판단하고 최적화한다면, 사용자 의도에 더 가까운 검색 결과를 얻을 수 있습니다.

### Agentic 접근 방식의 차별점
기존 벡터 검색 워크플로우
```
데이터 ➡ 데이터 임베딩(고정 모델) ➡ 인덱스 생성 ➡ 사용자 질의 ➡ 유사도 검색 ➡ 답변
```
똑똑한 Agentic 방식
```
사용자 질의 분석(검색 전략 세움) ➡ 최적 임베딩 모델 선택 ➡ 데이터 임베딩 ➡ 인덱스 생성 ➡ 유사도 검색 ➡ 검색결과 바탕으로 답변 정제
```

### DuckDB?
허깅페이스의 데이터셋은 파켓(parquet) 파일을 기반으로 동작하는데, 이때 빠른 인메모리 데이터베이스 시스템인 [DuckDB를 사용하면 이 파일들과 상호작용](https://huggingface.co/docs/hub/en/datasets-duckdb)할 수 있습니다. 또한 DuckDB의 기능 중 하나는 [벡터 유사도 검색](https://duckdb.org/docs/extensions/vss.html)으로, 인덱스 유무에 관계없이 사용할 수 있습니다.

이번 노트북에서는 단일 Agent에 여러가지 도구를 주어 수행하는 간단한 Agentic 벡터 검색 엔진을 만들어 보겠습니다.

필요한 의존성을 설치합니다 :

```python
# 본 예제 파일은 Python 3.10 이상 버전에서만 실행할 수 있습니다.
!pip install -U smolagents datasets sentence-transformers duckdb openai
```

HuggingFace의 추론 API를 사용하기 위해 로그인합니다 : 

```python
from huggingface_hub import notebook_login

notebook_login()
```

```python
from smolagents import tool
from datasets import Dataset
import os

# 도구 사용을 위해 OPENAI 키를 발급 받아야합니다.
os.environ["OPENAI_API_KEY"] = "YOUR KEY"
```

# 도구 정의
정의할 도구는 아래와 같습니다.
- 임베딩 생성 도구
- 인덱스 생성 도구
- 유사도 검색 도구
- 답변 생성 도구

### 도구1 : 임베딩 생성
일반적으로, 임베딩 작업에서는 작은 배치 사이즈로 데이터를 청킹하지만 여기서는 단순히 데이터셋을 임베딩으로 변환하는 과정만 수행하겠습니다.

```python
@tool
def create_embeddings(
    dataset: Dataset,
    model_id: str,
    column_name: str,
    ) -> Dataset:
  """
    주어진 데이터셋에 대해 임베딩을 생성합니다.

    Args:
        dataset: 임베딩을 생성할 대상 데이터셋
        model_id: 임베딩에 사용할 모델
        column_name: 임베딩할 열 이름

    Returns:
        임베딩이 추가된 데이터셋
  """
  from sentence_transformers import SentenceTransformer

  model = SentenceTransformer(model_id)

  def embed_batch(batch):
    embeddings = model.encode(batch[column_name], convert_to_numpy=True)
    batch["embeddings"] = embeddings.tolist()
    return batch

  dataset = dataset.map(embed_batch, batched=True)

  return dataset
```

## DuckDB로 벡터 검색 수행하기
`duckdb`를 사용하여 데이터셋에서 벡터 검색을 수행할 수 있습니다.
인덱스는 사용하거나 사용하지 않을 수 있습니다. 인덱스를 **활용하지 않고** 검색하는 것은 더 느리지만 더 정확하고, 인덱스를 **활용하여** 검색하는 것은 더 빠르지만 덜 정확합니다.

### 인덱스를 활용하지 않고 검색하기
인덱스를 활용하지 않고 검색하는 것은 느린 작업지만 일반적으로 약 10만 행까지의 작은 데이터셋에서는 충분히 빠르게 동작합니다. 하지만 이번 노트북에서는 DuckDB의 인덱스를 활용해서 검색해보겠습니다.

### 인덱스를 활용하여 검색하기

이 접근법은 데이터셋의 로컬 복사본을 생성하고 이를 사용하여 인덱스를 생성합니다. 약간의 오버헤드가 있지만 한번 인덱스를 생성한 후에는 검색 속도가 개선될 것입니다.

### 도구2 : DuckDB 인덱스 만들기

```python
@tool
def create_db_index(
    dataset_with_embeddings: Dataset,  # 임베딩이 포함된 데이터셋
    table_name: str,
    embedding_column: str = "embeddings"
) -> None:
    """
    임베딩이 포함된 데이터셋에 대해 DuckDB 인덱스를 생성합니다.

    Args:
        dataset_with_embeddings: 이미 임베딩이 포함된 데이터셋
        table_name: 생성할 테이블 이름
        embedding_column: 임베딩 열 이름

    Returns:
        None
    """
    import duckdb
    
    # VSS 확장 설치 및 로드
    duckdb.sql("INSTALL vss; LOAD vss;")
    duckdb.sql(f"DROP TABLE IF EXISTS {table_name};")

    # 데이터셋을 pandas DataFrame으로 변환
    df = dataset_with_embeddings.to_pandas()

    # DuckDB에 DataFrame 등록
    duckdb.register(f"{table_name}_temp", df)

    # 모델에서 임베딩 차원 가져오기
    embedding_dim = len(df[embedding_column].iloc[0])
    # embedding_dim = model.get_sentence_embedding_dimension()

    # 테이블 생성 (임베딩을 FLOAT 배열로 변환)
    duckdb.sql(f"""
        CREATE TABLE {table_name} AS
        SELECT *, {embedding_column}::FLOAT[{embedding_dim}] AS {embedding_column}_float
        FROM {table_name}_temp;
    """)

    # HNSW 인덱스 생성
    duckdb.sql(f"""
        CREATE INDEX idx_{embedding_column} ON {table_name}
        USING HNSW ({embedding_column}_float) WITH (metric = 'cosine');
    """)

    # 임시 테이블 정리
    duckdb.sql(f"DROP VIEW IF EXISTS {table_name}_temp;")
```

이 도구를 통해 인덱스를 사용하여 벡터 검색을 수행할 수 있으며, 결과는 즉시 반환됩니다.

### 도구3 : 벡터 검색 수행하기

```python
@tool
def similarity_search_with_duckdb_index(
    query: str,
    table_name: str,
    model_id: str,
    k: int = 5,
    embedding_column: str = "embeddings"
)-> dict:
    """
    DuckDB 인덱스를 이용해 벡터 검색을 수행합니다.

    Args:
        query: 검색할 쿼리 문자열
        model_id: 임베딩에 사용할 모델
        k: 반환할 결과 수
        table_name: 검색할 테이블 이름
        embedding_column: 임베딩 컬럼 이름

    Returns:
        dict: 검색 결과
    """
    from sentence_transformers import SentenceTransformer
    import duckdb
    
    model = SentenceTransformer(model_id)
    embedding = model.encode(query).tolist()
    return duckdb.sql(
        query=f"""
        SELECT *, array_cosine_distance({embedding_column}_float, {embedding}::FLOAT[{model.get_sentence_embedding_dimension()}]) as distance
        FROM {table_name}
        ORDER BY distance
        LIMIT {k};
    """
    ).to_df()
```

무거운 벡터 검색 엔진을 따로 배포할 필요 없고, 저장소는 허브에서 처리됩니다.

### 도구4 : 답변 생성 도구
유사도 검색 결과 청크를 기반으로, LLM이 사용자가 원할 만한 답변을 생성합니다.

```python
@tool
def generate_answer(chunks: list, query: str) -> str:
    """
    쿼리를 기반으로 주어진 텍스트 청크 목록에서 답변을 생성합니다.

    Args:
        chunks: 답변 생성에 사용할 텍스트 청크 목록
        query: 답변할 쿼리 문자열

    Returns:
        str: 생성된 답변
    """
    import openai   # OPENAI 키 발급이 필요합니다.
    
    context = "\n\n".join(chunks)
    prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
    response = openai.ChatCompletion.create(
        model="gpt-4",
        messages=[{"role": "user", "content": prompt}]
    )
    return response["choices"][0]["message"]["content"]
```

도구를 정의하였으니, 데이터셋을 로드하고 에이전트를 동작해보겠습니다.

사용한 데이터셋은 [huggingface-KREW/KoCultre-Descriptions](https://huggingface.co/datasets/huggingface-KREW/KoCultre-Descriptions)으로, 한국어 밈(meme)에 대한 데이터입니다.

```python
from datasets import load_dataset

dataset = load_dataset("huggingface-KREW/KoCultre-Descriptions")
```

에이전트를 정의합니다. 사용할 모델은 Qwen/Qwen2.5-Coder-32B-Instruct 입니다.

```python
from smolagents import CodeAgent, InferenceClientModel

model = InferenceClientModel(
    "Qwen/Qwen2.5-Coder-32B-Instruct",
    provider="together",
    max_tokens=2048
)

tools = [
    create_embeddings,                    # 임베딩 생성 도구
    create_db_index,                      # 인덱스 생성 도구
    similarity_search_with_duckdb_index,  # 유사도 검색 도구
    generate_answer                       # 답변 생성 도구
]

agent = CodeAgent(
    model=model,
    tools=tools,
    verbosity_level=2,
    max_steps=15
)
```

에이전트가 자율적으로 임베딩 모델을 택할 수 있도록 프롬프트를 구성하겠습니다.

```python
def agentic_prompt(query: str):
  return f"""
      당신은 지능형 검색 전문가입니다. {query}:

      검색을 위해 먼저 쿼리와 데이터셋을 분석하고, 분석 결과에 따라 가장 적합한 임베딩 모델을 선택하세요
        - sentence-transformers/all-MiniLM-L6-v2 (fast, lightweight, good for simple queries)
        - sentence-transformers/all-mpnet-base-v2 (balanced performance)
        - intfloat/e5-large-v2 (high quality for complex tasks)
        - minishlab/potion-base-8M (very efficient)
      """
```

데이터셋에 대한 설명과, 검색어를 입력하고 검색엔진 에이전트를 실행합니다.

```python
result = agent.run(
    agentic_prompt(query="아이돌 관련 밈 알려주세요"),
    additional_args={"dataset": dataset,
                     "dataset_description": "한국어 밈(meme)에 대한 데이터셋으로, meme과(title), meme의 뜻(content)을 알려준다.",
                     "column_name": "content"}
)
```

```python
>>> print(f"Final result: {result}")
```

Final result: 엄마가되는 밈은 최근 SNS와 온라인 커뮤니티에서 유행하는 표현으로, 주로 누군가가 매우 귀엽거나 사랑스러워서 모성애적 감정을 느낄 때 사용됩니다. 이 표현은 아이돌 팬덤 문화에서 자주 사용되며, '엄마가 돼'가 올바른 표기이지만 의도적인 맞춤법 파괴를 통해 밈으로서의 독특한 성격을 갖게 되었습니다.

```
Final result: 엄마가되는 밈은 최근 SNS와 온라인 커뮤니티에서 유행하는 표현으로, 주로 누군가가 매우 귀엽거나 사랑스러워서 모성애적 감정을 느낄 때 사용됩니다. 이 표현은 아이돌 팬덤 문화에서 자주 사용되며, '엄마가 돼'가 올바른 표기이지만 의도적인 맞춤법 파괴를 통해 밈으로서의 독특한 성격을 갖게 되었습니다.
```

### **Conclusion**

단순히 검색 결과만을 가져오는 것이 아니라, 검색결과를 바탕으로 쿼리에 따른 답변을 정제해서 보여주었습니다.✌🏻

지금까지 간단한 에이전트 시스템을 만들어보았습니다. 여기에 품질 평가, 분석 등을 추가한다면 진정한 Agentic 검색 엔진을 구현할 수 있습니다.