From a7466de41829a19a1d731a334ee4bd976ac1d013 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 24 Sep 2025 16:28:52 +0100 Subject: [PATCH] More LLMs --- .../model/text_completion/azure_openai/llm.py | 13 +++++++++---- .../trustgraph/model/text_completion/mistral/llm.py | 13 +++++++++---- .../trustgraph/model/text_completion/ollama/llm.py | 13 +++++++++---- .../trustgraph/model/text_completion/openai/llm.py | 13 +++++++++---- .../trustgraph/model/text_completion/vllm/llm.py | 13 +++++++++---- 5 files changed, 45 insertions(+), 20 deletions(-) diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py index 11376426..f0588e5a 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py @@ -54,7 +54,7 @@ class Processor(LlmService): self.temperature = temperature self.max_output = max_output - self.model = model + self.default_model = model self.openai = AzureOpenAI( api_key=token, @@ -62,14 +62,19 @@ class Processor(LlmService): 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 try: resp = self.openai.chat.completions.create( - model=self.model, + model=model_name, messages=[ { "role": "user", @@ -97,7 +102,7 @@ class Processor(LlmService): text = resp.choices[0].message.content, in_token = inputtokens, out_token = outputtokens, - model = self.model + model = model_name ) return r diff --git a/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py b/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py index 6dfd2656..c4ce26d5 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py @@ -41,21 +41,26 @@ class Processor(LlmService): } ) - self.model = model + self.default_model = model self.temperature = temperature self.max_output = max_output self.mistral = Mistral(api_key=api_key) 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 try: resp = self.mistral.chat.complete( - model=self.model, + model=model_name, messages=[ { "role": "user", @@ -87,7 +92,7 @@ class Processor(LlmService): text = resp.choices[0].message.content, in_token = inputtokens, out_token = outputtokens, - model = self.model + model = model_name ) return resp diff --git a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py index 97ed7d15..fc19ace3 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py @@ -33,16 +33,21 @@ class Processor(LlmService): } ) - self.model = model + self.default_model = model 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 try: - response = self.llm.generate(self.model, prompt) + response = self.llm.generate(model_name, prompt) response_text = response['response'] logger.debug("Sending response...") @@ -55,7 +60,7 @@ class Processor(LlmService): text = response_text, in_token = inputtokens, out_token = outputtokens, - model = self.model + model = model_name ) return resp diff --git a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py index 8aa8c6b9..79e0c86f 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py @@ -47,7 +47,7 @@ class Processor(LlmService): } ) - self.model = model + self.default_model = model self.temperature = temperature self.max_output = max_output @@ -58,14 +58,19 @@ class Processor(LlmService): 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 try: resp = self.openai.chat.completions.create( - model=self.model, + model=model_name, messages=[ { "role": "user", @@ -97,7 +102,7 @@ class Processor(LlmService): text = resp.choices[0].message.content, in_token = inputtokens, out_token = outputtokens, - model = self.model + model = model_name ) return resp diff --git a/trustgraph-flow/trustgraph/model/text_completion/vllm/llm.py b/trustgraph-flow/trustgraph/model/text_completion/vllm/llm.py index f194dc86..71b77d74 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/vllm/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/vllm/llm.py @@ -45,21 +45,26 @@ class Processor(LlmService): self.base_url = base_url self.temperature = temperature self.max_output = max_output - self.model = model + self.default_model = model self.session = aiohttp.ClientSession() logger.info(f"Using vLLM service at {base_url}") 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 = { "Content-Type": "application/json", } request = { - "model": self.model, + "model": model_name, "prompt": system + "\n\n" + prompt, "max_tokens": self.max_output, "temperature": self.temperature, @@ -91,7 +96,7 @@ class Processor(LlmService): text = ans, in_token = inputtokens, out_token = outputtokens, - model = self.model, + model = model_name, ) return resp