docsgpt llm provider

This commit is contained in:
Alex 2024-01-08 23:35:37 +00:00
parent a3e6239e6e
commit 033bcf80d0
3 changed files with 57 additions and 2 deletions

View File

@ -7,7 +7,7 @@ current_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__
class Settings(BaseSettings):
LLM_NAME: str = "openai"
LLM_NAME: str = "docsgpt"
EMBEDDINGS_NAME: str = "openai_text-embedding-ada-002"
CELERY_BROKER_URL: str = "redis://localhost:6379/0"
CELERY_RESULT_BACKEND: str = "redis://localhost:6379/1"

View File

@ -0,0 +1,53 @@
from application.llm.base import BaseLLM
from application.core.settings import settings
import json
import io
import requests
class DocsGPTAPILLM(BaseLLM):
def __init__(self, *args, **kwargs):
self.endpoint = "https://llm.docsgpt.co.uk"
def gen(self, model, engine, 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"
response = requests.post(
f"{self.endpoint}/answer",
json={
"prompt": prompt,
"max_new_tokens": 30
}
)
response_clean = response.json()['a'].split("###")[0]
return response_clean
def gen_stream(self, model, engine, 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"
# send prompt to endpoint /stream
response = requests.post(
f"{self.endpoint}/stream",
json={
"prompt": prompt,
"max_new_tokens": 256
},
stream=True
)
for line in response.iter_lines():
import sys
print(line, file=sys.stderr)
if line:
#data = json.loads(line)
data_str = line.decode('utf-8')
if data_str.startswith("data: "):
data = json.loads(data_str[6:])
yield data['a']

View File

@ -3,6 +3,7 @@ from application.llm.sagemaker import SagemakerAPILLM
from application.llm.huggingface import HuggingFaceLLM
from application.llm.llama_cpp import LlamaCpp
from application.llm.anthropic import AnthropicLLM
from application.llm.docsgpt_provider import DocsGPTAPILLM
@ -13,7 +14,8 @@ class LLMCreator:
'sagemaker': SagemakerAPILLM,
'huggingface': HuggingFaceLLM,
'llama.cpp': LlamaCpp,
'anthropic': AnthropicLLM
'anthropic': AnthropicLLM,
'docsgpt': DocsGPTAPILLM
}
@classmethod