From e065c270fa9bc2cab013ab3156328214a4287b84 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 24 Sep 2025 16:24:38 +0100 Subject: [PATCH] Ported some LLMs to dynamic models --- .../model/text_completion/bedrock/llm.py | 55 +++++++++++++------ .../model/text_completion/claude/llm.py | 13 +++-- .../model/text_completion/cohere/llm.py | 15 +++-- .../text_completion/googleaistudio/llm.py | 43 ++++++++++----- 4 files changed, 86 insertions(+), 40 deletions(-) diff --git a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py index 292a2282..25584780 100755 --- a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py @@ -183,13 +183,13 @@ class Processor(LlmService): } ) - self.model = model + # Store default configuration + self.default_model = model self.temperature = temperature self.max_output = max_output - self.variant = self.determine_variant(self.model)() - self.variant.set_temperature(temperature) - self.variant.set_max_output(max_output) + # Cache for model variants to avoid re-initialization + self.model_variants = {} self.session = boto3.Session( aws_access_key_id=aws_access_key_id, @@ -208,47 +208,66 @@ class Processor(LlmService): # FIXME: Missing, Amazon models, Deepseek # This set of conditions deals with normal bedrock on-demand usage - if self.model.startswith("mistral"): + if model.startswith("mistral"): return Mistral - elif self.model.startswith("meta"): + elif model.startswith("meta"): return Meta - elif self.model.startswith("anthropic"): + elif model.startswith("anthropic"): return Anthropic - elif self.model.startswith("ai21"): + elif model.startswith("ai21"): return Ai21 - elif self.model.startswith("cohere"): + elif model.startswith("cohere"): return Cohere # The inference profiles - if self.model.startswith("us.meta"): + if model.startswith("us.meta"): return Meta - elif self.model.startswith("us.anthropic"): + elif model.startswith("us.anthropic"): return Anthropic - elif self.model.startswith("eu.meta"): + elif model.startswith("eu.meta"): return Meta - elif self.model.startswith("eu.anthropic"): + elif model.startswith("eu.anthropic"): return Anthropic return Default - async def generate_content(self, system, prompt): + def _get_or_create_variant(self, model_name): + """Get cached model variant or create new one""" + if model_name not in self.model_variants: + logger.info(f"Creating model variant for '{model_name}'") + variant_class = self.determine_variant(model_name) + variant = variant_class() + variant.set_temperature(self.temperature) + variant.set_max_output(self.max_output) + self.model_variants[model_name] = variant + + return self.model_variants[model_name] + + 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}") try: + # Get the appropriate variant for this model + variant = self._get_or_create_variant(model_name) - promptbody = self.variant.encode_request(system, prompt) + promptbody = variant.encode_request(system, prompt) accept = 'application/json' contentType = 'application/json' response = self.bedrock.invoke_model( body=promptbody, - modelId=self.model, + modelId=model_name, accept=accept, contentType=contentType ) # Response structure decode - outputtext = self.variant.decode_response(response) + outputtext = variant.decode_response(response) metadata = response['ResponseMetadata']['HTTPHeaders'] inputtokens = int(metadata['x-amzn-bedrock-input-token-count']) @@ -262,7 +281,7 @@ class Processor(LlmService): text = outputtext, in_token = inputtokens, out_token = outputtokens, - model = self.model + model = model_name ) return resp diff --git a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py index 87b611f4..b6180038 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py @@ -41,19 +41,24 @@ class Processor(LlmService): } ) - self.model = model + self.default_model = model self.claude = anthropic.Anthropic(api_key=api_key) self.temperature = temperature self.max_output = max_output logger.info("Claude 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}") try: response = message = self.claude.messages.create( - model=self.model, + model=model_name, max_tokens=self.max_output, temperature=self.temperature, system = system, @@ -81,7 +86,7 @@ class Processor(LlmService): text = resp, in_token = inputtokens, out_token = outputtokens, - model = self.model + model = model_name ) return resp diff --git a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py index df2c1143..a5b1deda 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py @@ -39,18 +39,23 @@ class Processor(LlmService): } ) - self.model = model + self.default_model = model self.temperature = temperature self.cohere = cohere.Client(api_key=api_key) logger.info("Cohere 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}") try: - output = self.cohere.chat( - model=self.model, + output = self.cohere.chat( + model=model_name, message=prompt, preamble = system, temperature=self.temperature, @@ -71,7 +76,7 @@ class Processor(LlmService): text = resp, in_token = inputtokens, out_token = outputtokens, - model = self.model + model = model_name ) return resp diff --git a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py index 6170490a..c1814129 100644 --- a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py @@ -53,10 +53,13 @@ class Processor(LlmService): ) self.client = genai.Client(api_key=api_key) - self.model = model + self.default_model = model self.temperature = temperature self.max_output = max_output + # Cache for generation configs per model + self.generation_configs = {} + block_level = HarmBlockThreshold.BLOCK_ONLY_HIGH self.safety_settings = [ @@ -83,22 +86,36 @@ class Processor(LlmService): logger.info("GoogleAIStudio LLM service initialized") - async def generate_content(self, system, prompt): + def _get_or_create_config(self, model_name): + """Get cached generation config or create new one""" + if model_name not in self.generation_configs: + logger.info(f"Creating generation config for '{model_name}'") + self.generation_configs[model_name] = types.GenerateContentConfig( + temperature = self.temperature, + top_p = 1, + top_k = 40, + max_output_tokens = self.max_output, + response_mime_type = "text/plain", + safety_settings = self.safety_settings, + ) - generation_config = types.GenerateContentConfig( - temperature = self.temperature, - top_p = 1, - top_k = 40, - max_output_tokens = self.max_output, - response_mime_type = "text/plain", - system_instruction = system, - safety_settings = self.safety_settings, - ) + return self.generation_configs[model_name] + + 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}") + + generation_config = self._get_or_create_config(model_name) + # Set system instruction per request (can't be cached) + generation_config.system_instruction = system try: response = self.client.models.generate_content( - model=self.model, + model=model_name, config=generation_config, contents=prompt, ) @@ -114,7 +131,7 @@ class Processor(LlmService): text = resp, in_token = inputtokens, out_token = outputtokens, - model = self.model + model = model_name ) return resp