fix: api_key capturing + pytest errors

This commit is contained in:
Siddhant Rai
2024-04-15 22:32:24 +05:30
parent 60a670ce29
commit 77991896b4
11 changed files with 276 additions and 145 deletions

View File

@@ -4,7 +4,7 @@ from application.core.settings import settings
class LlamaCpp(BaseLLM):
def __init__(self, api_key, llm_name=settings.MODEL_PATH, *args, **kwargs):
def __init__(self, api_key=None, llm_name=settings.MODEL_PATH, *args, **kwargs):
global llama
try:
from llama_cpp import Llama
@@ -17,7 +17,7 @@ class LlamaCpp(BaseLLM):
self.api_key = api_key
llama = Llama(model_path=llm_name, n_ctx=2048)
def _raw_gen(self, model, messages, stream=False, **kwargs):
def _raw_gen(self, baseself, model, messages, stream=False, **kwargs):
context = messages[0]["content"]
user_question = messages[-1]["content"]
prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"
@@ -29,7 +29,7 @@ class LlamaCpp(BaseLLM):
return result["choices"][0]["text"].split("### Answer \n")[-1]
def _raw_gen_stream(self, model, messages, stream=True, **kwargs):
def _raw_gen_stream(self, baseself, model, messages, stream=True, **kwargs):
context = messages[0]["content"]
user_question = messages[-1]["content"]
prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"