apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-20 10:41:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Based AzureOpenAI support

Browse source
This commit is contained in:
JackColquitt 2024-10-03 16:54:37 -07:00 • committed by Cyber MacGeddon
parent d219b9a7cd
commit f861399a8d
2 changed files with 7 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

1
trustgraph-flow/setup.py
View file

@ -83,6 +83,7 @@ setuptools.setup(
"scripts/rows-write-cassandra", "scripts/rows-write-cassandra",
"scripts/run-processing", "scripts/run-processing",
"scripts/text-completion-azure", "scripts/text-completion-azure",
"scripts/text-completion-azure-openai",
"scripts/text-completion-claude", "scripts/text-completion-claude",
"scripts/text-completion-cohere", "scripts/text-completion-cohere",
"scripts/text-completion-llamafile", "scripts/text-completion-llamafile",

14
trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py
View file

@ -1,7 +1,7 @@
""" """
Simple LLM service, performs text prompt completion using the Azure Simple LLM service, performs text prompt completion using the Azure
serverless endpoint service. Input is prompt, output is response. OpenAI endpoit service. Input is prompt, output is response.
""" """
import requests import requests
@ -66,17 +66,14 @@ class Processor(ConsumerProducer):
] ]
) )
self.api = api
self.endpoint = endpoint
self.token = token
self.temperature = temperature self.temperature = temperature
self.max_output = max_output self.max_output = max_output
self.model = model self.model = model
self.openai = AzureOpenAI( self.openai = AzureOpenAI(
api_key=self.token, api_key=token,
api_version=self.api, api_version=api,
azure_endpoint = self.endpoint, azure_endpoint = endpoint,
) )
def handle(self, msg): def handle(self, msg):
@ -91,10 +88,11 @@ class Processor(ConsumerProducer):
prompt = v.prompt prompt = v.prompt
try: try:
with __class__.text_completion_metric.time(): with __class__.text_completion_metric.time():
resp = self.openai.completions.create( resp = self.openai.chat.completions.create(
model=self.model, model=self.model,
messages=[ messages=[
{ {