agents/BasicAgent.py

import os

from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool, HfApiModel, LiteLLMModel, Tool, GoogleSearchTool
from langchain_community.tools.wikidata.tool import WikidataAPIWrapper, WikidataQueryRun
from langchain_community.tools.brave_search.tool import BraveSearch

class BasicAgent(CodeAgent):
    def __init__(self):
        print("Initializing BasicAgent...")

        model_name = os.getenv("MODEL_NAME")

        if not model_name:
            raise ValueError(
                "Environment variable 'MODEL_NAME' is required but not set.")

        print(f"LLM Config: Model='{model_name}'")

        try:
            llm = LiteLLMModel(
            model_id=model_name,
            api_key=os.environ.get("GEMINI_API_KEY", "")
        )
        except Exception as e:
            print(f"Error initializing LLM: {e}")
            raise

        # TOOLS
        print("Initializing tools...")
        try:
            search_tool = Tool.from_langchain(BraveSearch.from_api_key(os.environ.get("SEARCH_API_KEY", "")))
            visit_tool = VisitWebpageTool()
            wikidata = Tool.from_langchain(WikidataQueryRun(api_wrapper=WikidataAPIWrapper())) 
            tools = [search_tool, visit_tool, wikidata]
            print(f"Tools initialized: {[tool.name for tool in tools]}")
        except Exception as e:
            print(f"Error initializing tools: {e}")
            raise

        try:
            super().__init__(
                model=llm,
                tools=tools,
            )
            print("BasicAgent (CodeAgent) initialized successfully.")
        except Exception as e:
            print(f"Error initializing CodeAgent: {e}")
            raise

    def __call__(self, question: str) -> str:
        print(
            f"BasicAgent received question (first 50 chars): {question[:50]}...")
        try:
            final_answer = self.run(question)
            print(
                f"BasicAgent returning answer (first 50 chars): {str(final_answer)[:50]}...")
            return str(final_answer)
        except Exception as e:
            print(
                f"Error during agent execution for question '{question[:50]}...': {e}")
            return f"AGENT ERROR: Failed to process the question due to: {e}"