Update run_inference to use the provided LLM configuration params
This commit is contained in:
committed by
Paul Kompfner
parent
afa7573834
commit
21a55f6aae
@@ -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
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user