jproman's picture
Refactor to use langchain
522e6c3
import os
import config
if "LOCALLLM" in os.environ and os.environ["LOCALLLM"]:
from ollama import chat as OllamaChat
#################################################
from langchain_community.tools import DuckDuckGoSearchRun
def callWebSearch(query):
return DuckDuckGo(query)
def DuckDuckGo(query):
search_tool = DuckDuckGoSearchRun()
results = search_tool.invoke(query)
return results
#################################################
from langchain_community.tools import WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper
import requests
from bs4 import BeautifulSoup
def callWikipediaSearch(query):
return callWikipediaLangchain(query)
def callWikipediaLangchain(query):
wikipedia = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
response = wikipedia.run(query)
return response
def callCustomWikipediaSearch(query):
searchURL = config.wikipediaSearchURL + query
response = requests.get(searchURL, timeout=60)
response.raise_for_status()
searchResult = response.json()
for pageId in searchResult['query']['pages']:
if searchResult['query']['pages'][pageId]['index'] == 1:
page = searchResult['query']['pages'][pageId]['title']
# response3 = requests.get('https://en.wikipedia.org/w/api.php',
# params={'action': 'parse','page': page,'format': 'json'}).json()
# raw_html3 = response3['parse']['text']['*']
response2 = requests.get(config.wikipediaRetrieveURL + page)
response2.raise_for_status()
raw_html = response2.text
soup = BeautifulSoup(raw_html, 'html.parser')
raw_text = soup.get_text(separator=" ",strip=True)
return raw_text
# response2 = requests.get(config.wikipediaRetrieveURL + page)
# response2.raise_for_status()
# return response2.text
# response2 = requests.get('https://en.wikipedia.org/w/api.php',
# params={'action': 'query', 'format': 'json', 'titles': page, 'prop': 'extracts', 'exintro': True, 'explaintext': True })
# searchResult2 = response.json()
# for pageId2 in searchResult2['query']['pages']:
# if searchResult2['query']['pages'][pageId2]['index'] == 1:
# return searchResult2['query']['pages'][pageId2]['extract']
return "No result found in wikipedia. Search elsewhere!!"
#################################################
from huggingface_hub import InferenceClient
def callLLM(query):
if "LOCALLLM" in os.environ:
return callLocalLLM(query)
else:
return callHfInferenceClientLLM(query)
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)
print(callWikipediaSearch("Mercedes Sosa discography"))