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

More LLMs

Browse source
This commit is contained in:
Cyber MacGeddon 2025-09-24 16:28:52 +01:00
parent e065c270fa
commit a7466de418
5 changed files with 45 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -54,7 +54,7 @@ class Processor(LlmService):
self.temperature = temperature self.temperature = temperature
self.max_output = max_output self.max_output = max_output
self.model = model self.default_model = model
self.openai = AzureOpenAI( self.openai = AzureOpenAI(
api_key=token, api_key=token,
@ -62,14 +62,19 @@ class Processor(LlmService):
azure_endpoint = endpoint, azure_endpoint = endpoint,
) )
async def generate_content(self, system, prompt): async def generate_content(self, system, prompt, model=None):
# Use provided model or fall back to default
model_name = model or self.default_model
logger.debug(f"Using model: {model_name}")
prompt = system + "\n\n" + prompt prompt = system + "\n\n" + prompt
try: try:
resp = self.openai.chat.completions.create( resp = self.openai.chat.completions.create(
model=self.model, model=model_name,
messages=[ messages=[
{ {
"role": "user", "role": "user",
@ -97,7 +102,7 @@ class Processor(LlmService):
text = resp.choices[0].message.content, text = resp.choices[0].message.content,
in_token = inputtokens, in_token = inputtokens,
out_token = outputtokens, out_token = outputtokens,
model = self.model model = model_name
) )
return r return r

13
trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py
View file

@ -41,21 +41,26 @@ class Processor(LlmService):
} }
) )
self.model = model self.default_model = model
self.temperature = temperature self.temperature = temperature
self.max_output = max_output self.max_output = max_output
self.mistral = Mistral(api_key=api_key) self.mistral = Mistral(api_key=api_key)
logger.info("Mistral LLM service initialized") logger.info("Mistral LLM service initialized")
async def generate_content(self, system, prompt): async def generate_content(self, system, prompt, model=None):
# Use provided model or fall back to default
model_name = model or self.default_model
logger.debug(f"Using model: {model_name}")
prompt = system + "\n\n" + prompt prompt = system + "\n\n" + prompt
try: try:
resp = self.mistral.chat.complete( resp = self.mistral.chat.complete(
model=self.model, model=model_name,
messages=[ messages=[
{ {
"role": "user", "role": "user",
@ -87,7 +92,7 @@ class Processor(LlmService):
text = resp.choices[0].message.content, text = resp.choices[0].message.content,
in_token = inputtokens, in_token = inputtokens,
out_token = outputtokens, out_token = outputtokens,
model = self.model model = model_name
) )
return resp return resp

13
trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py
View file

@ -33,16 +33,21 @@ class Processor(LlmService):
} }
) )
self.model = model self.default_model = model
self.llm = Client(host=ollama) self.llm = Client(host=ollama)
async def generate_content(self, system, prompt): async def generate_content(self, system, prompt, model=None):
# Use provided model or fall back to default
model_name = model or self.default_model
logger.debug(f"Using model: {model_name}")
prompt = system + "\n\n" + prompt prompt = system + "\n\n" + prompt
try: try:
response = self.llm.generate(self.model, prompt) response = self.llm.generate(model_name, prompt)
response_text = response['response'] response_text = response['response']
logger.debug("Sending response...") logger.debug("Sending response...")
@ -55,7 +60,7 @@ class Processor(LlmService):
text = response_text, text = response_text,
in_token = inputtokens, in_token = inputtokens,
out_token = outputtokens, out_token = outputtokens,
model = self.model model = model_name
) )
return resp return resp

13
trustgraph-flow/trustgraph/model/text_completion/openai/llm.py
View file

@ -47,7 +47,7 @@ class Processor(LlmService):
} }
) )
self.model = model self.default_model = model
self.temperature = temperature self.temperature = temperature
self.max_output = max_output self.max_output = max_output
@ -58,14 +58,19 @@ class Processor(LlmService):
logger.info("OpenAI LLM service initialized") logger.info("OpenAI LLM service initialized")
async def generate_content(self, system, prompt): async def generate_content(self, system, prompt, model=None):
# Use provided model or fall back to default
model_name = model or self.default_model
logger.debug(f"Using model: {model_name}")
prompt = system + "\n\n" + prompt prompt = system + "\n\n" + prompt
try: try:
resp = self.openai.chat.completions.create( resp = self.openai.chat.completions.create(
model=self.model, model=model_name,
messages=[ messages=[
{ {
"role": "user", "role": "user",
@ -97,7 +102,7 @@ class Processor(LlmService):
text = resp.choices[0].message.content, text = resp.choices[0].message.content,
in_token = inputtokens, in_token = inputtokens,
out_token = outputtokens, out_token = outputtokens,
model = self.model model = model_name
) )
return resp return resp

13
trustgraph-flow/trustgraph/model/text_completion/vllm/llm.py
View file

@ -45,21 +45,26 @@ class Processor(LlmService):
self.base_url = base_url self.base_url = base_url
self.temperature = temperature self.temperature = temperature
self.max_output = max_output self.max_output = max_output
self.model = model self.default_model = model
self.session = aiohttp.ClientSession() self.session = aiohttp.ClientSession()
logger.info(f"Using vLLM service at {base_url}") logger.info(f"Using vLLM service at {base_url}")
logger.info("vLLM LLM service initialized") logger.info("vLLM LLM service initialized")
async def generate_content(self, system, prompt): async def generate_content(self, system, prompt, model=None):
# Use provided model or fall back to default
model_name = model or self.default_model
logger.debug(f"Using model: {model_name}")
headers = { headers = {
"Content-Type": "application/json", "Content-Type": "application/json",
} }
request = { request = {
"model": self.model, "model": model_name,
"prompt": system + "\n\n" + prompt, "prompt": system + "\n\n" + prompt,
"max_tokens": self.max_output, "max_tokens": self.max_output,
"temperature": self.temperature, "temperature": self.temperature,
@ -91,7 +96,7 @@ class Processor(LlmService):
text = ans, text = ans,
in_token = inputtokens, in_token = inputtokens,
out_token = outputtokens, out_token = outputtokens,
model = self.model, model = model_name,
) )
return resp return resp