from application.retriever.base import BaseRetriever from application.core.settings import settings from application.llm.llm_creator import LLMCreator from application.utils import count_tokens from langchain_community.tools import DuckDuckGoSearchResults from langchain_community.utilities import DuckDuckGoSearchAPIWrapper class DuckDuckSearch(BaseRetriever): def __init__(self, question, source, chat_history, prompt, chunks=2, gpt_model='docsgpt'): self.question = question self.source = source self.chat_history = chat_history self.prompt = prompt self.chunks = chunks self.gpt_model = gpt_model def _parse_lang_string(self, input_string): result = [] current_item = "" inside_brackets = False for char in input_string: if char == "[": inside_brackets = True elif char == "]": inside_brackets = False result.append(current_item) current_item = "" elif inside_brackets: current_item += char if inside_brackets: result.append(current_item) return result def _get_data(self): if self.chunks == 0: docs = [] else: wrapper = DuckDuckGoSearchAPIWrapper(max_results=self.chunks) search = DuckDuckGoSearchResults(api_wrapper=wrapper) results = search.run(self.question) results = self._parse_lang_string(results) docs = [] for i in results: try: text = i.split("title:")[0] title = i.split("title:")[1].split("link:")[0] link = i.split("link:")[1] docs.append({"text": text, "title": title, "link": link}) except IndexError: pass if settings.LLM_NAME == "llama.cpp": docs = [docs[0]] return docs def gen(self): docs = self._get_data() # join all page_content together with a newline docs_together = "\n".join([doc["text"] for doc in docs]) p_chat_combine = self.prompt.replace("{summaries}", docs_together) messages_combine = [{"role": "system", "content": p_chat_combine}] for doc in docs: yield {"source": doc} if len(self.chat_history) > 1: tokens_current_history = 0 # count tokens in history self.chat_history.reverse() for i in self.chat_history: if "prompt" in i and "response" in i: tokens_batch = count_tokens(i["prompt"]) + count_tokens(i["response"]) if tokens_current_history + tokens_batch < settings.TOKENS_MAX_HISTORY: tokens_current_history += tokens_batch messages_combine.append({"role": "user", "content": i["prompt"]}) messages_combine.append({"role": "system", "content": i["response"]}) messages_combine.append({"role": "user", "content": self.question}) llm = LLMCreator.create_llm(settings.LLM_NAME, api_key=settings.API_KEY) completion = llm.gen_stream(model=self.gpt_model, messages=messages_combine) for line in completion: yield {"answer": str(line)} def search(self): return self._get_data()