tools/search_helios_documentation.py

from smolagents import Tool
from typing import Any, Optional

class SimpleTool(Tool):
    name = "search_helios_documentation"
    description = "Recherche dans la documentation Helios - fonction simplifi\u00e9e."
    inputs = {'query': {'type': 'string', 'description': 'requête de recherche'}, 'max_results': {'type': 'integer', 'nullable': True, 'description': 'Nombre maximum de résultats'}}
    output_type = "string"

    def forward(self, query: str, max_results: int = 5) -> str:
        """
        Recherche dans la documentation Helios - fonction simplifiée.

        Args:
            query: requête de recherche
            max_results: Nombre maximum de résultats

        Returns:
            Liste des documents trouvés
        """
        from langchain_chroma import Chroma

        from langchain_openai import OpenAIEmbeddings, ChatOpenAI
        embedding_function=OpenAIEmbeddings()
        persist_directory = "chroma_db"
        vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embedding_function)

        retriever = vectorstore.as_retriever(search_kwargs={"k": max_results})
        # Récupérer les documents pertinents
        relevant_docs = retriever.get_relevant_documents(query, filter={"source": "bugs"})
        unique_docs = []
        seen = set()
        for doc in relevant_docs:
            if doc.page_content not in seen:
                unique_docs.append({'source':doc.metadata.get('source', 'Non spécifiée'), 'content':doc.page_content})
                seen.add(doc.page_content)
        return unique_docs