mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-17 17:21:02 +02:00
LLM base class working
This commit is contained in:
Cyber MacGeddon
parent
1047cd2fa1
commit
6a9469de34
4 changed files with 139 additions and 98 deletions
|
|
@ -12,4 +12,7 @@ from . setting_spec import SettingSpec
|
|||
from . producer_spec import ProducerSpec
|
||||
from . subscriber_spec import SubscriberSpec
|
||||
from . request_response_spec import RequestResponseSpec
|
||||
from . llm_service import LlmService, LlmResult
|
||||
|
||||
|
||||
|
||||
|
|
|
|||
116
trustgraph-base/trustgraph/base/llm_service.py
Executable file
116
trustgraph-base/trustgraph/base/llm_service.py
Executable file
|
|
@ -0,0 +1,116 @@
|
|||
|
||||
"""
|
||||
LLM text completion base class
|
||||
"""
|
||||
|
||||
import time
|
||||
from prometheus_client import Histogram
|
||||
|
||||
from .. schema import TextCompletionRequest, TextCompletionResponse, Error
|
||||
from .. exceptions import TooManyRequests
|
||||
from .. base import FlowProcessor, ConsumerSpec, ProducerSpec
|
||||
|
||||
default_ident = "text-completion"
|
||||
|
||||
class LlmResult:
|
||||
__slots__ = ["text", "in_token", "out_token", "model"]
|
||||
|
||||
class LlmService(FlowProcessor):
|
||||
|
||||
def __init__(self, **params):
|
||||
|
||||
id = params.get("id")
|
||||
|
||||
super(LlmService, self).__init__(**params | { "id": id })
|
||||
|
||||
self.register_specification(
|
||||
ConsumerSpec(
|
||||
name = "request",
|
||||
schema = TextCompletionRequest,
|
||||
handler = self.on_request
|
||||
)
|
||||
)
|
||||
|
||||
self.register_specification(
|
||||
ProducerSpec(
|
||||
name = "response",
|
||||
schema = TextCompletionResponse
|
||||
)
|
||||
)
|
||||
|
||||
if not hasattr(__class__, "text_completion_metric"):
|
||||
__class__.text_completion_metric = Histogram(
|
||||
'text_completion_duration',
|
||||
'Text completion duration (seconds)',
|
||||
["id", "flow"],
|
||||
buckets=[
|
||||
0.25, 0.5, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0,
|
||||
8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0,
|
||||
17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0, 25.0,
|
||||
30.0, 35.0, 40.0, 45.0, 50.0, 60.0, 80.0, 100.0,
|
||||
120.0
|
||||
]
|
||||
)
|
||||
|
||||
async def on_request(self, msg, consumer, flow):
|
||||
|
||||
try:
|
||||
|
||||
request = msg.value()
|
||||
|
||||
# Sender-produced ID
|
||||
|
||||
id = msg.properties()["id"]
|
||||
|
||||
prompt = request.system + "\n\n" + request.prompt
|
||||
|
||||
with __class__.text_completion_metric.labels(
|
||||
id=self.id,
|
||||
flow=f"{flow.name}-{consumer.name}",
|
||||
).time():
|
||||
|
||||
response = await self.generate_content(
|
||||
request.system, request.prompt
|
||||
)
|
||||
|
||||
await flow.producer["response"].send(
|
||||
TextCompletionResponse(
|
||||
error=None,
|
||||
response=response.text,
|
||||
in_token=response.in_token,
|
||||
out_token=response.out_token,
|
||||
model=response.model
|
||||
),
|
||||
properties={"id": id}
|
||||
)
|
||||
|
||||
except TooManyRequests as e:
|
||||
raise e
|
||||
|
||||
except Exception as e:
|
||||
|
||||
# Apart from rate limits, treat all exceptions as unrecoverable
|
||||
|
||||
print(f"Exception: {e}")
|
||||
|
||||
print("Send error response...", flush=True)
|
||||
|
||||
await flow.producer["response"].send(
|
||||
TextCompletionResponse(
|
||||
error=Error(
|
||||
type = "llm-error",
|
||||
message = str(e),
|
||||
),
|
||||
response=None,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=None,
|
||||
),
|
||||
properties={"id": id}
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
FlowProcessor.add_args(parser)
|
||||
|
||||
|
|
@ -13,8 +13,6 @@ import uuid
|
|||
from .... schema import Chunk, Triple, Triples, Metadata, Value
|
||||
from .... schema import EntityContext, EntityContexts
|
||||
from .... schema import PromptRequest, PromptResponse
|
||||
from .... log_level import LogLevel
|
||||
from .... clients.prompt_client import PromptClient
|
||||
from .... rdf import TRUSTGRAPH_ENTITIES, DEFINITION, RDF_LABEL, SUBJECT_OF
|
||||
|
||||
from .... base import FlowProcessor, RequestResponseSpec, ConsumerSpec
|
||||
|
|
|
|||
|
|
@ -4,22 +4,17 @@ Simple LLM service, performs text prompt completion using VertexAI on
|
|||
Google Cloud. Input is prompt, output is response.
|
||||
"""
|
||||
|
||||
import vertexai
|
||||
import time
|
||||
from prometheus_client import Histogram
|
||||
import os
|
||||
|
||||
from google.oauth2 import service_account
|
||||
import google
|
||||
import vertexai
|
||||
|
||||
from vertexai.preview.generative_models import (
|
||||
Content, FunctionDeclaration, GenerativeModel, GenerationConfig,
|
||||
HarmCategory, HarmBlockThreshold, Part, Tool,
|
||||
)
|
||||
|
||||
from .... schema import TextCompletionRequest, TextCompletionResponse, Error
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import FlowProcessor, ConsumerSpec, ProducerSpec
|
||||
from .... base import LlmService, LlmResult
|
||||
|
||||
default_ident = "text-completion"
|
||||
|
||||
|
|
@ -29,11 +24,10 @@ default_temperature = 0.0
|
|||
default_max_output = 8192
|
||||
default_private_key = "private.json"
|
||||
|
||||
class Processor(FlowProcessor):
|
||||
class Processor(LlmService):
|
||||
|
||||
def __init__(self, **params):
|
||||
|
||||
id = params.get("id")
|
||||
region = params.get("region", default_region)
|
||||
model = params.get("model", default_model)
|
||||
private_key = params.get("private_key", default_private_key)
|
||||
|
|
@ -43,41 +37,7 @@ class Processor(FlowProcessor):
|
|||
if private_key is None:
|
||||
raise RuntimeError("Private key file not specified")
|
||||
|
||||
super(Processor, self).__init__(
|
||||
**params | {
|
||||
"request_schema": TextCompletionRequest,
|
||||
"response_schema": TextCompletionResponse,
|
||||
}
|
||||
)
|
||||
|
||||
self.register_specification(
|
||||
ConsumerSpec(
|
||||
name = "request",
|
||||
schema = TextCompletionRequest,
|
||||
handler = self.on_request
|
||||
)
|
||||
)
|
||||
|
||||
self.register_specification(
|
||||
ProducerSpec(
|
||||
name = "response",
|
||||
schema = TextCompletionResponse
|
||||
)
|
||||
)
|
||||
|
||||
if not hasattr(__class__, "text_completion_metric"):
|
||||
__class__.text_completion_metric = Histogram(
|
||||
'text_completion_duration',
|
||||
'Text completion duration (seconds)',
|
||||
["id", "flow"],
|
||||
buckets=[
|
||||
0.25, 0.5, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0,
|
||||
8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0,
|
||||
17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0, 25.0,
|
||||
30.0, 35.0, 40.0, 45.0, 50.0, 60.0, 80.0, 100.0,
|
||||
120.0
|
||||
]
|
||||
)
|
||||
super(Processor, self).__init__(**params)
|
||||
|
||||
self.parameters = {
|
||||
"temperature": temperature,
|
||||
|
|
@ -134,48 +94,29 @@ class Processor(FlowProcessor):
|
|||
|
||||
print("Initialisation complete", flush=True)
|
||||
|
||||
async def on_request(self, msg, consumer, flow):
|
||||
async def generate_content(self, system, prompt):
|
||||
|
||||
try:
|
||||
|
||||
request = msg.value()
|
||||
prompt = system + "\n\n" + prompt
|
||||
|
||||
# Sender-produced ID
|
||||
response = self.llm.generate_content(
|
||||
prompt, generation_config=self.generation_config,
|
||||
safety_settings=self.safety_settings
|
||||
)
|
||||
|
||||
id = msg.properties()["id"]
|
||||
resp = LlmResult()
|
||||
resp.text = response.text
|
||||
resp.in_token = response.usage_metadata.prompt_token_count
|
||||
resp.out_token = response.usage_metadata.candidates_token_count
|
||||
resp.model = self.model
|
||||
|
||||
prompt = request.system + "\n\n" + request.prompt
|
||||
|
||||
with __class__.text_completion_metric.labels(
|
||||
id=self.id,
|
||||
flow=f"{flow.name}-{consumer.name}",
|
||||
).time():
|
||||
|
||||
response = self.llm.generate_content(
|
||||
prompt, generation_config=self.generation_config,
|
||||
safety_settings=self.safety_settings
|
||||
)
|
||||
|
||||
resp = response.text
|
||||
inputtokens = int(response.usage_metadata.prompt_token_count)
|
||||
outputtokens = int(response.usage_metadata.candidates_token_count)
|
||||
print(resp, flush=True)
|
||||
print(f"Input Tokens: {inputtokens}", flush=True)
|
||||
print(f"Output Tokens: {outputtokens}", flush=True)
|
||||
print(f"Input Tokens: {resp.in_token}", flush=True)
|
||||
print(f"Output Tokens: {resp.out_token}", flush=True)
|
||||
|
||||
print("Send response...", flush=True)
|
||||
|
||||
|
||||
await flow.producer["response"].send(
|
||||
TextCompletionResponse(
|
||||
error=None,
|
||||
response=resp,
|
||||
in_token=inputtokens,
|
||||
out_token=outputtokens,
|
||||
model=self.model
|
||||
),
|
||||
properties={"id": id}
|
||||
)
|
||||
return resp
|
||||
|
||||
except google.api_core.exceptions.ResourceExhausted as e:
|
||||
|
||||
|
|
@ -187,29 +128,13 @@ class Processor(FlowProcessor):
|
|||
except Exception as e:
|
||||
|
||||
# Apart from rate limits, treat all exceptions as unrecoverable
|
||||
|
||||
print(f"Exception: {e}")
|
||||
|
||||
print("Send error response...", flush=True)
|
||||
|
||||
await flow.producer["response"].send(
|
||||
TextCompletionResponse(
|
||||
error=Error(
|
||||
type = "llm-error",
|
||||
message = str(e),
|
||||
),
|
||||
response=None,
|
||||
in_token=None,
|
||||
out_token=None,
|
||||
model=None,
|
||||
),
|
||||
properties={"id": id}
|
||||
)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
FlowProcessor.add_args(parser)
|
||||
LlmService.add_args(parser)
|
||||
|
||||
parser.add_argument(
|
||||
'-m', '--model',
|
||||
|
|
@ -243,6 +168,5 @@ class Processor(FlowProcessor):
|
|||
)
|
||||
|
||||
def run():
|
||||
|
||||
Processor.launch(default_ident, __doc__)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue