LLM/DocsGPT
1
0
Fork 0
mirror of https://github.com/arc53/DocsGPT.git synced 2025-11-29 16:43:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: api_key capturing + pytest errors

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

6
application/llm/sagemaker.py
View File

@@ -60,7 +60,7 @@ class LineIterator:
class SagemakerAPILLM(BaseLLM):
def __init__(self, api_key, *args, **kwargs):
def __init__(self, api_key=None, *args, **kwargs):
import boto3
runtime = boto3.client(
@@ -75,7 +75,7 @@ class SagemakerAPILLM(BaseLLM):
self.endpoint = settings.SAGEMAKER_ENDPOINT
self.runtime = runtime
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"
@@ -104,7 +104,7 @@ class SagemakerAPILLM(BaseLLM):
print(result[0]["generated_text"], file=sys.stderr)
return result[0]["generated_text"][len(prompt) :]
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"