mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-28 01:46:22 +02:00
65cda7b276
* Add system prompt to LLM invocation * Added system parameter to LLMs * Added to Bedrock and VertexAI
42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
|
|
import _pulsar
|
|
|
|
from .. schema import TextCompletionRequest, TextCompletionResponse
|
|
from .. schema import text_completion_request_queue
|
|
from .. schema import text_completion_response_queue
|
|
from . base import BaseClient
|
|
|
|
# Ugly
|
|
ERROR=_pulsar.LoggerLevel.Error
|
|
WARN=_pulsar.LoggerLevel.Warn
|
|
INFO=_pulsar.LoggerLevel.Info
|
|
DEBUG=_pulsar.LoggerLevel.Debug
|
|
|
|
class LlmClient(BaseClient):
|
|
|
|
def __init__(
|
|
self, log_level=ERROR,
|
|
subscriber=None,
|
|
input_queue=None,
|
|
output_queue=None,
|
|
pulsar_host="pulsar://pulsar:6650",
|
|
):
|
|
|
|
if input_queue is None: input_queue = text_completion_request_queue
|
|
if output_queue is None: output_queue = text_completion_response_queue
|
|
|
|
super(LlmClient, self).__init__(
|
|
log_level=log_level,
|
|
subscriber=subscriber,
|
|
input_queue=input_queue,
|
|
output_queue=output_queue,
|
|
pulsar_host=pulsar_host,
|
|
input_schema=TextCompletionRequest,
|
|
output_schema=TextCompletionResponse,
|
|
)
|
|
|
|
def request(self, system, prompt, timeout=300):
|
|
return self.call(
|
|
system=system, prompt=prompt, timeout=timeout
|
|
).response
|
|
|