slandl/Agents_Course_Final_Assignment

1import os2from time import sleep3 4from llama_index.llms.ollama import Ollama5from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI6from llama_index.llms.openai import OpenAI7 8from llama_index.tools.duckduckgo.base import DuckDuckGoSearchToolSpec9from llama_index.tools.wikipedia import WikipediaToolSpec10from llama_index.tools.playwright import PlaywrightToolSpec11 12import tiktoken13from llama_index.core.callbacks import CallbackManager, TokenCountingHandler14from llama_index.core import Settings15from dotenv import load_dotenv16 17load_dotenv()18 19 20def get_llm_backend(structured_output_cls=None):21    model_type = os.getenv("MODEL_TYPE", "no_model_type").lower()22 23    if model_type == "ollama":24        llm = get_ollama_backend()25    elif model_type == "huggingface":26        llm = get_huggingface_backend()27    elif model_type == "openai":28        llm = get_openai_backend()29    else:30        raise Exception(f"Invalid model type: {model_type}")31 32    if structured_output_cls:33        print(34            f"Creating structured LLM with output class: {structured_output_cls.__name__}"35        )36        return llm.as_structured_llm(output_cls=structured_output_cls)37 38    return llm39 40 41def get_huggingface_backend():42    model_name = os.getenv("HF_MODEL")43    temperature = float(os.getenv("HF_TEMPERATURE", "0.15"))44    print(f"Using Hugging Face model: {model_name} with temperature: {temperature}")45    return HuggingFaceInferenceAPI(46        model_name=model_name,47        token=os.getenv("HF_TOKEN"),48        temperature=temperature,49        provider="auto",50    )51 52 53def get_openai_backend():54    model_name = os.getenv("OPENAI_MODEL")55    api_key = os.getenv("OPENAI_KEY")56    temperature = float(os.getenv("OPENAI_TEMPERATURE", "0.15"))57    base_url = os.getenv("OPENAI_BASE_URL")58 59    max_retries = int(os.getenv("OPENAI_MAX_RETRIES", "10"))60    timeout = float(os.getenv("OPENAI_TIMEOUT", "120.0"))61 62    if not api_key:63        raise Exception(64            "OPENAI_KEY environment variable is required for OpenAI backend"65        )66 67    openai_kwargs = {68        "model": model_name,69        "api_key": api_key,70        "temperature": temperature,71        "max_retries": max_retries,72        "timeout": timeout,73    }74 75    if base_url:76        openai_kwargs["base_url"] = base_url77        print(78            f"Using OpenAI model: {model_name} with temperature: {temperature}, max_retries: {max_retries}, timeout: {timeout}s and base_url: {base_url}"79        )80    else:81        print(82            f"Using OpenAI model: {model_name} with temperature: {temperature}, max_retries: {max_retries}, timeout: {timeout}s"83        )84 85    return OpenAI(**openai_kwargs)86 87 88def get_ollama_backend():89    model = os.getenv("OLLAMA_MODEL")90    temperature = float(os.getenv("OLLAMA_TEMPERATURE", "0.15"))91    context_window = int(os.getenv("OLLAMA_CONTEXT_WINDOW", "4000"))92    print(93        f"Using Ollama model: {model} with temperature: {temperature} and context window: {context_window}"94    )95    return Ollama(96        model=model,97        temperature=temperature,98        context_window=context_window,99    )100 101 102async def get_tools():103    tools = []104 105    # ddg_tool_spec = DuckDuckGoSearchToolSpec()106    # tools.extend(ddg_tool_spec.to_tool_list())107 108    # google_tool_spec = GoogleSearchToolSpec()109    # tools.extend(google_tool_spec.to_tool_list())110 111    # wikipedia_tool_spec = WikipediaToolSpec()112    # tools.extend(wikipedia_tool_spec.to_tool_list())113 114    browser = await PlaywrightToolSpec.create_async_playwright_browser(headless=False)115    playwright_tool_spec = PlaywrightToolSpec(async_browser=browser)116    tools.extend(playwright_tool_spec.to_tool_list())117 118    return tools119 120 121def init_token_counter(model=None):122    if model:123        tokenizer = tiktoken.encoding_for_model(model).encode124        print(f"Token counter initializing with encoding_for_model: {model}")125    else:126        tokenizer = tiktoken.get_encoding("cl100k_base").encode127        print(f"Token counter initializing with get_encoding: cl100k_base")128 129    token_counter = TokenCountingHandler(tokenizer=tokenizer)130    callback_manager = CallbackManager([token_counter])131 132    Settings.callback_manager = callback_manager133 134    return token_counter135