apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-04-25 00:16:23 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Use model in Azure LLM integration

Browse source
This commit is contained in:
Cyber MacGeddon 2026-03-04 12:05:40 +00:00
parent e19ea8667d
commit ae27eccbed
1 changed files with 16 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

21
trustgraph-flow/trustgraph/model/text_completion/azure/llm.py
View file

@ -55,11 +55,13 @@ class Processor(LlmService):
self.max_output = max_output
self.default_model = model
def build_prompt(self, system, content, temperature=None, stream=False):
def build_prompt(self, system, content, temperature=None, stream=False, model=None):
# Use provided temperature or fall back to default
effective_temperature = temperature if temperature is not None else self.temperature
model_name = model or self.default_model
data = {
"model": model_name,
"messages": [
{
"role": "system", "content": system
@ -100,7 +102,8 @@ class Processor(LlmService):
raise TooManyRequests()
if resp.status_code != 200:
raise RuntimeError("LLM failure")
logger.error(f"Azure API error: status={resp.status_code}, body={resp.text}")
raise RuntimeError(f"LLM failure: HTTP {resp.status_code}")
result = resp.json()
@ -121,7 +124,8 @@ class Processor(LlmService):
prompt = self.build_prompt(
system,
prompt,
effective_temperature
effective_temperature,
model=model_name
)
response = self.call_llm(prompt)
@ -174,7 +178,7 @@ class Processor(LlmService):
logger.debug(f"Using temperature: {effective_temperature}")
try:
body = self.build_prompt(system, prompt, effective_temperature, stream=True)
body = self.build_prompt(system, prompt, effective_temperature, stream=True, model=model_name)
url = self.endpoint
api_key = self.token
@ -190,7 +194,8 @@ class Processor(LlmService):
raise TooManyRequests()
if response.status_code != 200:
raise RuntimeError("LLM failure")
logger.error(f"Azure API error: status={response.status_code}, body={response.text}")
raise RuntimeError(f"LLM failure: HTTP {response.status_code}")
total_input_tokens = 0
total_output_tokens = 0
@ -279,6 +284,12 @@ class Processor(LlmService):
help=f'LLM max output tokens (default: {default_max_output})'
)
parser.add_argument(
'-m', '--model',
default=default_model,
help=f'LLM model name (default: {default_model})'
)
def run():
Processor.launch(default_ident, __doc__)