Update run_inference to use the provided LLM configuration params

This commit is contained in:
Mark Backman
2025-12-09 13:33:05 -05:00
committed by Paul Kompfner
parent afa7573834
commit 21a55f6aae
6 changed files with 366 additions and 65 deletions

View File

@@ -267,26 +267,41 @@ class AnthropicLLMService(LLMService):
"""
messages = []
system = NOT_GIVEN
tools = []
if isinstance(context, LLMContext):
adapter: AnthropicLLMAdapter = self.get_llm_adapter()
params = adapter.get_llm_invocation_params(
invocation_params = adapter.get_llm_invocation_params(
context, enable_prompt_caching=self._settings["enable_prompt_caching"]
)
messages = params["messages"]
system = params["system"]
messages = invocation_params["messages"]
system = invocation_params["system"]
tools = invocation_params["tools"]
else:
context = AnthropicLLMContext.upgrade_to_anthropic(context)
messages = context.messages
system = getattr(context, "system", NOT_GIVEN)
tools = context.tools or []
# Build params using the same method as streaming completions
params = {
"model": self.model_name,
"max_tokens": self._settings["max_tokens"],
"stream": False,
"temperature": self._settings["temperature"],
"top_k": self._settings["top_k"],
"top_p": self._settings["top_p"],
"messages": messages,
"system": system,
"tools": tools,
"betas": ["interleaved-thinking-2025-05-14"],
}
if self._settings["thinking"]:
params["thinking"] = self._settings["thinking"].model_dump(exclude_unset=True)
params.update(self._settings["extra"])
# LLM completion
response = await self._client.messages.create(
model=self.model_name,
messages=messages,
system=system,
max_tokens=8192,
stream=False,
)
response = await self._client.beta.messages.create(**params)
return response.content[0].text

View File

@@ -840,15 +840,13 @@ class AWSBedrockLLMService(LLMService):
messages = context.messages
system = getattr(context, "system", None) # [{"text": "system message"}]
# Determine if we're using Claude or Nova based on model ID
model_id = self.model_name
# Prepare request parameters
# Prepare request parameters using the same method as streaming
inference_config = self._build_inference_config()
request_params = {
"modelId": model_id,
"modelId": self.model_name,
"messages": messages,
"additionalModelRequestFields": self._settings["additional_model_request_fields"],
}
if inference_config:

View File

@@ -798,17 +798,25 @@ class GoogleLLMService(LLMService):
"""
messages = []
system = []
tools = []
if isinstance(context, LLMContext):
adapter = self.get_llm_adapter()
params: GeminiLLMInvocationParams = adapter.get_llm_invocation_params(context)
messages = params["messages"]
system = params["system_instruction"]
tools = params["tools"]
else:
context = GoogleLLMContext.upgrade_to_google(context)
messages = context.messages
system = getattr(context, "system_message", None)
tools = context.tools or []
generation_config = GenerateContentConfig(system_instruction=system)
# Build generation config using the same method as streaming
generation_params = self._build_generation_params(
system_instruction=system, tools=tools if tools else None
)
generation_config = GenerateContentConfig(**generation_params)
# Use the new google-genai client's async method
response = await self._client.aio.models.generate_content(
@@ -825,6 +833,48 @@ class GoogleLLMService(LLMService):
return None
def _build_generation_params(
self,
system_instruction: Optional[str] = None,
tools: Optional[List] = None,
tool_config: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
"""Build generation parameters for Google AI API.
Args:
system_instruction: Optional system instruction to use.
tools: Optional list of tools to include.
tool_config: Optional tool configuration.
Returns:
Dictionary of generation parameters with None values filtered out.
"""
# Filter out None values and create GenerationContentConfig
generation_params = {
k: v
for k, v in {
"system_instruction": system_instruction,
"temperature": self._settings["temperature"],
"top_p": self._settings["top_p"],
"top_k": self._settings["top_k"],
"max_output_tokens": self._settings["max_tokens"],
"tools": tools,
"tool_config": tool_config,
}.items()
if v is not None
}
# Add thinking parameters if configured
if self._settings["thinking"]:
generation_params["thinking_config"] = self._settings["thinking"].model_dump(
exclude_unset=True
)
if self._settings["extra"]:
generation_params.update(self._settings["extra"])
return generation_params
def _maybe_unset_thinking_budget(self, generation_params: Dict[str, Any]):
try:
# There's no way to introspect on model capabilities, so
@@ -862,36 +912,15 @@ class GoogleLLMService(LLMService):
if self._tool_config:
tool_config = self._tool_config
# Filter out None values and create GenerationContentConfig
generation_params = {
k: v
for k, v in {
"system_instruction": self._system_instruction,
"temperature": self._settings["temperature"],
"top_p": self._settings["top_p"],
"top_k": self._settings["top_k"],
"max_output_tokens": self._settings["max_tokens"],
"tools": tools,
"tool_config": tool_config,
}.items()
if v is not None
}
# Add thinking parameters if configured
if self._settings["thinking"]:
generation_params["thinking_config"] = self._settings["thinking"].model_dump(
exclude_unset=True
)
if self._settings["extra"]:
generation_params.update(self._settings["extra"])
# Build generation parameters
generation_params = self._build_generation_params(
system_instruction=self._system_instruction, tools=tools, tool_config=tool_config
)
# possibly modify generation_params (in place) to set thinking to off by default
self._maybe_unset_thinking_budget(generation_params)
generation_config = (
GenerateContentConfig(**generation_params) if generation_params else None
)
generation_config = GenerateContentConfig(**generation_params)
await self.start_ttfb_metrics()
return await self._client.aio.models.generate_content_stream(
@@ -1166,6 +1195,14 @@ class GoogleLLMService(LLMService):
# Do nothing - we're shutting down anyway
pass
async def _update_settings(self, settings):
"""Override to handle ThinkingConfig validation."""
# Convert thinking dict to ThinkingConfig if needed
if "thinking" in settings and isinstance(settings["thinking"], dict):
settings = dict(settings) # Make a copy to avoid modifying the original
settings["thinking"] = self.ThinkingConfig(**settings["thinking"])
await super()._update_settings(settings)
def create_context_aggregator(
self,
context: OpenAILLMContext,

View File

@@ -276,17 +276,23 @@ class BaseOpenAILLMService(LLMService):
"""
if isinstance(context, LLMContext):
adapter = self.get_llm_adapter()
params: OpenAILLMInvocationParams = adapter.get_llm_invocation_params(context)
messages = params["messages"]
invocation_params: OpenAILLMInvocationParams = adapter.get_llm_invocation_params(
context
)
else:
messages = context.messages
invocation_params = OpenAILLMInvocationParams(
messages=context.messages, tools=context.tools, tool_choice=context.tool_choice
)
# Build params using the same method as streaming completions
params = self.build_chat_completion_params(invocation_params)
# Override for non-streaming
params["stream"] = False
params.pop("stream_options", None)
# LLM completion
response = await self._client.chat.completions.create(
model=self.model_name,
messages=messages,
stream=False,
)
response = await self._client.chat.completions.create(**params)
return response.choices[0].message.content