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

Add model detection for inference profiles in US and EU

Browse source
This commit is contained in:
Cyber MacGeddon 2025-03-13 14:46:46 +00:00
parent b6dcc484fc
commit 3967d92b8d
1 changed files with 16 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

18
trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py
View file

@ -35,6 +35,7 @@ default_profile = os.getenv("AWS_PROFILE", None)
default_region = os.getenv("AWS_DEFAULT_REGION", None)
# Variant API handling depends on the model type
# FIXME: Missing, Amazon models, Deepseek
class ModelVariant(enum.Enum):
MISTRAL = enum.auto() # Mistral
META = enum.auto() # Llama 3.1
@ -122,6 +123,9 @@ class Processor(ConsumerProducer):
def determine_model(self, model):
# FIXME: Missing, Amazon models, Deepseek
# This set of conditions deals with normal bedrock on-demand usage
if self.model.startswith("mistral"):
return ModelVariant.MISTRAL
elif self.model.startswith("meta"):
@ -132,8 +136,18 @@ class Processor(ConsumerProducer):
return ModelVariant.AI21
elif self.model.startswith("cohere"):
return ModelVariant.COHERE
else:
return ModelVariant.DEFAULT
# The inference profiles
if self.model.startswith("us.meta"):
return ModelVariant.META
elif self.model.startswith("us.anthropic"):
return ModelVariant.ANTHROPIC
elif self.model.startswith("eu.meta"):
return ModelVariant.META
elif self.model.startswith("eu.anthropic"):
return ModelVariant.ANTHROPIC
return ModelVariant.DEFAULT
async def handle(self, msg):