import os

from huggingface_hub import InferenceClient
from langchain_community.tools import DuckDuckGoSearchRun
import config


def callWebSearch(query):
    return DuckDuckGo(query)

def callLLM(query):
    if "LOCALLLM" in os.environ:
        return callLocalLLM(query)
    else:
        return callHfInferenceClientLLM(query)

def DuckDuckGo(query):
    search_tool = DuckDuckGoSearchRun()
    results = search_tool.invoke(query)
    return results

def callLocalLLM(query):
    response = OllamaChat(model=os.environ["LOCALLLM"], messages=[ { 'role': 'user', 'content': query } ])
    return response['message']['content']

def callHfInferenceClientLLM(query):
    client = InferenceClient(config.hfMoldel)
    response = client.chat.completions.create(
        messages = [ {"role": "user", "content": query } ],
        stream=False, max_tokens=1024 )
    return response.choices[0].message.content

if __name__ == "__main__":
    os.environ["LOCALLLM"] = "llama3.2"
    from ollama import chat as OllamaChat
    response = callLLM("What is the capital of France?")
    print(response)
    response = callWebSearch("who is the president of France")
    print(response)
    response = callHfInferenceClientLLM("What is the capital of France?")
    print(response)