fed8b2eed7
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
from application.llm.base import BaseLLM
|
|
from application.core.settings import settings
|
|
|
|
|
|
class PremAILLM(BaseLLM):
|
|
provider_name = "premai"
|
|
|
|
def __init__(self, api_key=None, user_api_key=None, *args, **kwargs):
|
|
from premai import Prem
|
|
|
|
super().__init__(*args, **kwargs)
|
|
self.client = Prem(api_key=api_key)
|
|
self.api_key = api_key
|
|
self.user_api_key = user_api_key
|
|
self.project_id = settings.PREMAI_PROJECT_ID
|
|
|
|
def _raw_gen(self, baseself, model, messages, stream=False, **kwargs):
|
|
response = self.client.chat.completions.create(
|
|
model=model,
|
|
project_id=self.project_id,
|
|
messages=messages,
|
|
stream=stream,
|
|
**kwargs
|
|
)
|
|
|
|
return response.choices[0].message["content"]
|
|
|
|
def _raw_gen_stream(self, baseself, model, messages, stream=True, **kwargs):
|
|
response = self.client.chat.completions.create(
|
|
model=model,
|
|
project_id=self.project_id,
|
|
messages=messages,
|
|
stream=stream,
|
|
**kwargs
|
|
)
|
|
|
|
for line in response:
|
|
if line.choices[0].delta["content"] is not None:
|
|
yield line.choices[0].delta["content"]
|