feat: pass decoded_token to llm and retrievers

This commit is contained in:
Siddhant Rai
2025-03-18 23:46:02 +05:30
parent f4ab85a2bb
commit ab95d90284
9 changed files with 75 additions and 25 deletions

View File

@@ -17,6 +17,7 @@ class BraveRetSearch(BaseRetriever):
token_limit=150,
gpt_model="docsgpt",
user_api_key=None,
decoded_token=None,
):
self.question = question
self.source = source
@@ -35,6 +36,7 @@ class BraveRetSearch(BaseRetriever):
)
)
self.user_api_key = user_api_key
self.decoded_token = decoded_token
def _get_data(self):
if self.chunks == 0:
@@ -81,7 +83,10 @@ class BraveRetSearch(BaseRetriever):
messages_combine.append({"role": "user", "content": self.question})
llm = LLMCreator.create_llm(
settings.LLM_NAME, api_key=settings.API_KEY, user_api_key=self.user_api_key
settings.LLM_NAME,
api_key=settings.API_KEY,
user_api_key=self.user_api_key,
decoded_token=self.decoded_token,
)
completion = llm.gen_stream(model=self.gpt_model, messages=messages_combine)
@@ -100,5 +105,5 @@ class BraveRetSearch(BaseRetriever):
"chunks": self.chunks,
"token_limit": self.token_limit,
"gpt_model": self.gpt_model,
"user_api_key": self.user_api_key
"user_api_key": self.user_api_key,
}

View File

@@ -17,6 +17,7 @@ class ClassicRAG(BaseRetriever):
user_api_key=None,
llm_name=settings.LLM_NAME,
api_key=settings.API_KEY,
decoded_token=None,
):
self.original_question = ""
self.chat_history = chat_history if chat_history is not None else []
@@ -37,10 +38,14 @@ class ClassicRAG(BaseRetriever):
self.llm_name = llm_name
self.api_key = api_key
self.llm = LLMCreator.create_llm(
self.llm_name, api_key=self.api_key, user_api_key=self.user_api_key
self.llm_name,
api_key=self.api_key,
user_api_key=self.user_api_key,
decoded_token=decoded_token,
)
self.question = self._rephrase_query()
self.vectorstore = source["active_docs"] if "active_docs" in source else None
self.decoded_token = decoded_token
def _rephrase_query(self):
if (

View File

@@ -17,6 +17,7 @@ class DuckDuckSearch(BaseRetriever):
token_limit=150,
gpt_model="docsgpt",
user_api_key=None,
decoded_token=None,
):
self.question = question
self.source = source
@@ -35,6 +36,7 @@ class DuckDuckSearch(BaseRetriever):
)
)
self.user_api_key = user_api_key
self.decoded_token = decoded_token
def _parse_lang_string(self, input_string):
result = []
@@ -88,17 +90,20 @@ class DuckDuckSearch(BaseRetriever):
for doc in docs:
yield {"source": doc}
if len(self.chat_history) > 0:
if len(self.chat_history) > 0:
for i in self.chat_history:
if "prompt" in i and "response" in i:
messages_combine.append({"role": "user", "content": i["prompt"]})
messages_combine.append(
{"role": "assistant", "content": i["response"]}
)
if "prompt" in i and "response" in i:
messages_combine.append({"role": "user", "content": i["prompt"]})
messages_combine.append(
{"role": "assistant", "content": i["response"]}
)
messages_combine.append({"role": "user", "content": self.question})
llm = LLMCreator.create_llm(
settings.LLM_NAME, api_key=settings.API_KEY, user_api_key=self.user_api_key
settings.LLM_NAME,
api_key=settings.API_KEY,
user_api_key=self.user_api_key,
decoded_token=self.decoded_token,
)
completion = llm.gen_stream(model=self.gpt_model, messages=messages_combine)
@@ -107,7 +112,7 @@ class DuckDuckSearch(BaseRetriever):
def search(self):
return self._get_data()
def get_params(self):
return {
"question": self.question,
@@ -117,5 +122,5 @@ class DuckDuckSearch(BaseRetriever):
"chunks": self.chunks,
"token_limit": self.token_limit,
"gpt_model": self.gpt_model,
"user_api_key": self.user_api_key
"user_api_key": self.user_api_key,
}