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

from huggingface_hub import notebook_login

notebook_login()

from smolagents import tool
from datasets import Dataset
import os

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

@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

@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;")

@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()

@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"]

from datasets import load_dataset

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

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
)

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)
      """

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

print(f"Final result: {result}")