!pip install -q smolagents

!pip install markdownify duckduckgo-search spaces gradio-tools langchain langchain-community langchain-huggingface faiss-cpu --upgrade -q

from huggingface_hub import notebook_login

notebook_login()

from smolagents import HfApiModel

model_id = "Qwen/Qwen2.5-72B-Instruct"
model = HfApiModel(model_id)

from smolagents import CodeAgent, ToolCallingAgent, ManagedAgent, DuckDuckGoSearchTool, VisitWebpageTool

web_agent = ToolCallingAgent(tools=[DuckDuckGoSearchTool(), VisitWebpageTool()], model=model)

managed_web_agent = ManagedAgent(
    agent=web_agent,
    name="search_agent",
    description="Runs web searches for you. Give it your query as an argument.",
)

import datasets

knowledge_base = datasets.load_dataset("m-ric/huggingface_doc", split="train")

from tqdm import tqdm
from transformers import AutoTokenizer
from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores.utils import DistanceStrategy

source_docs = [
    Document(page_content=doc["text"], metadata={"source": doc["source"].split("/")[1]}) for doc in knowledge_base
]

text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
    AutoTokenizer.from_pretrained("thenlper/gte-small"),
    chunk_size=200,
    chunk_overlap=20,
    add_start_index=True,
    strip_whitespace=True,
    separators=["\n\n", "\n", ".", " ", ""],
)

# Split docs and keep only unique ones
print("Splitting documents...")
docs_processed = []
unique_texts = {}
for doc in tqdm(source_docs):
    new_docs = text_splitter.split_documents([doc])
    for new_doc in new_docs:
        if new_doc.page_content not in unique_texts:
            unique_texts[new_doc.page_content] = True
            docs_processed.append(new_doc)

print("Embedding documents...")
embedding_model = HuggingFaceEmbeddings(model_name="thenlper/gte-small")
huggingface_doc_vector_db = FAISS.from_documents(
    documents=docs_processed,
    embedding=embedding_model,
    distance_strategy=DistanceStrategy.COSINE,
)

from smolagents import Tool
from langchain_core.vectorstores import VectorStore


class RetrieverTool(Tool):
    name = "retriever"
    description = "Using semantic similarity, retrieves some documents from the knowledge base that have the closest embeddings to the input query."
    inputs = {
        "query": {
            "type": "string",
            "description": "The query to perform. This should be semantically close to your target documents. Use the affirmative form rather than a question.",
        }
    }
    output_type = "string"

    def __init__(self, vectordb: VectorStore, **kwargs):
        super().__init__(**kwargs)
        self.vectordb = vectordb

    def forward(self, query: str) -> str:
        assert isinstance(query, str), "Your search query must be a string"

        docs = self.vectordb.similarity_search(
            query,
            k=7,
        )

        return "\nRetrieved documents:\n" + "".join(
            [f"===== Document {str(i)} =====\n" + doc.page_content for i, doc in enumerate(docs)]
        )

huggingface_doc_retriever_tool = RetrieverTool(huggingface_doc_vector_db)

from google.colab import userdata

GITHUB_ACCESS_TOKEN = userdata.get("GITHUB_PERSONAL_TOKEN")

from langchain.document_loaders import GitHubIssuesLoader

loader = GitHubIssuesLoader(repo="huggingface/peft", access_token=GITHUB_ACCESS_TOKEN, include_prs=False, state="all")
docs = loader.load()

splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=30)
chunked_docs = splitter.split_documents(docs)

peft_issues_vector_db = FAISS.from_documents(chunked_docs, embedding=embedding_model)

peft_issues_retriever_tool = RetrieverTool(peft_issues_vector_db)

retriever_agent = ToolCallingAgent(
    tools=[huggingface_doc_retriever_tool, peft_issues_retriever_tool], model=model, max_iterations=4, verbose=2
)

managed_retriever_agent = ManagedAgent(
    agent=retriever_agent,
    name="retriever_agent",
    description="Retrieves documents from the knowledge base for you that are close to the input query. Give it your query as an argument. The knowledge base includes Hugging Face documentation and PEFT issues.",
)

from transformers import load_tool, CodeAgent

prompt_generator_tool = Tool.from_space(
    "sergiopaniego/Promptist", name="generator_tool", description="Optimizes user input into model-preferred prompts"
)

image_generation_tool = load_tool("m-ric/text-to-image", trust_remote_code=True)
image_generation_agent = CodeAgent(tools=[prompt_generator_tool, image_generation_tool], model=model)

managed_image_generation_agent = ManagedAgent(
    agent=image_generation_agent,
    name="image_generation_agent",
    description="Generates images from text prompts. Give it your prompt as an argument.",
    additional_prompting="\n\nYour final answer MUST BE only the generated image location.",
)

manager_agent = CodeAgent(
    tools=[],
    model=model,
    managed_agents=[managed_web_agent, managed_retriever_agent, managed_image_generation_agent],
    additional_authorized_imports=["time", "datetime", "PIL"],
)

manager_agent.run("How many years ago was Stripe founded?")

result = manager_agent.run(
    "Improve this prompt, then generate an image of it.", prompt="A rabbit wearing a space suit"
)

>>> from IPython.display import Image, display

>>> display(Image(filename=result))

manager_agent.run("How can I push a model to the Hub?")

manager_agent.run("How do you combine multiple adapters in peft?")