from smolagents import CodeAgent, DuckDuckGoSearchTool, InferenceClientModel

# 검색 도구 초기화
search_tool = DuckDuckGoSearchTool()

# 모델 초기화
model = InferenceClientModel()

agent = CodeAgent(
    model=model,
    tools=[search_tool],
)

# 예시 사용법
response = agent.run(
    "Search for luxury superhero-themed party ideas, including decorations, entertainment, and catering."
)
print(response)

from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from smolagents import Tool
from langchain_community.retrievers import BM25Retriever
from smolagents import CodeAgent, InferenceClientModel

class PartyPlanningRetrieverTool(Tool):
    name = "party_planning_retriever"
    description = "Alfred의 슈퍼히어로 테마 파티를 위한 아이디어를 의미적 검색으로 찾아줍니다."
    inputs = {
        "query": {
            "type": "string",
            "description": "파티 준비나 슈퍼히어로 테마와 관련된 쿼리입니다.",
        }
    }
    output_type = "string"

    def __init__(self, docs, **kwargs):
        super().__init__(**kwargs)
        self.retriever = BM25Retriever.from_documents(
            docs, k=5  # 상위 5개 문서 검색
        )

    def forward(self, query: str) -> str:
        assert isinstance(query, str), "검색 쿼리는 문자열이어야 합니다."

        docs = self.retriever.invoke(
            query,
        )
        return "\n검색된 아이디어:\n" + "".join(
            [
                f"\n\n===== 아이디어 {str(i)} =====\n" + doc.page_content
                for i, doc in enumerate(docs)
            ]
        )

# 파티 준비에 관한 지식 베이스 예시
party_ideas = [
    {"text": "고급 장식(금색 장식, 벨벳 커튼 등)이 있는 슈퍼히어로 테마 가면무도회.", "source": "Party Ideas 1"},
    {"text": "배트맨, 원더우먼 등 슈퍼히어로 테마 음악을 연주할 전문 DJ 고용.", "source": "Entertainment Ideas"},
    {"text": "케이터링 메뉴에 슈퍼히어로 이름을 붙인 요리 제공(예: '헐크의 그린 스무디', '아이언맨의 파워 스테이크').", "source": "Catering Ideas"},
    {"text": "장소 곳곳에 슈퍼히어로 로고와 고담 등 도시 프로젝션 장식.", "source": "Decoration Ideas"},
    {"text": "VR을 활용한 슈퍼히어로 시뮬레이션, 테마 게임 등 인터랙티브 체험 제공.", "source": "Entertainment Ideas"}
]

source_docs = [
    Document(page_content=doc["text"], metadata={"source": doc["source"]})
    for doc in party_ideas
]

# 문서를 작은 청크로 분할
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
    add_start_index=True,
    strip_whitespace=True,
    separators=["\n\n", "\n", ".", " ", ""],
)
docs_processed = text_splitter.split_documents(source_docs)

# 검색 도구 생성
party_planning_retriever = PartyPlanningRetrieverTool(docs_processed)

# 에이전트 초기화
agent = CodeAgent(tools=[party_planning_retriever], model=InferenceClientModel())

# 예시 사용법
response = agent.run(
    "Find ideas for a luxury superhero-themed party, including entertainment, catering, and decoration options."
)

print(response)