[WIP] Universal (LLM-agnostic) context machinery to support runtime LLM switching.

- Add to Google LLM service support for universal LLM context
This commit is contained in:
Paul Kompfner
2025-08-13 11:41:46 -04:00
parent 809c4c1bc5
commit 688b136141
5 changed files with 620 additions and 75 deletions

View File

@@ -6,21 +6,65 @@
"""Gemini LLM adapter for Pipecat."""
from typing import Any, Dict, List, Union
import base64
import json
from dataclasses import dataclass
from typing import Any, List, Optional, TypedDict
from loguru import logger
from pipecat.adapters.base_llm_adapter import BaseLLMAdapter
from pipecat.adapters.schemas.tools_schema import AdapterType, ToolsSchema
from pipecat.processors.aggregators.llm_context import LLMContext, LLMContextMessage
try:
from google.genai.types import (
Blob,
Content,
FunctionCall,
FunctionResponse,
Part,
)
except ModuleNotFoundError as e:
logger.error(f"Exception: {e}")
logger.error("In order to use Google AI, you need to `pip install pipecat-ai[google]`.")
raise Exception(f"Missing module: {e}")
class GeminiLLMAdapter(BaseLLMAdapter):
"""LLM adapter for Google's Gemini service.
class GeminiLLMInvocationParams(TypedDict):
"""Context-based parameters for invoking Gemini LLM."""
Provides tool schema conversion functionality to transform standard tool
definitions into Gemini's specific function-calling format for use with
Gemini LLM models.
system_instruction: Optional[str]
messages: List[Content]
tools: List[Any]
class GeminiLLMAdapter(BaseLLMAdapter[GeminiLLMInvocationParams]):
"""Gemini-specific adapter for Pipecat.
Handles:
- Extracting parameters for Gemini's API from a universal LLM context
- Converting Pipecat's standardized tools schema to Gemini's function-calling format.
- Extracting and sanitizing messages from the LLM context for logging with Gemini.
"""
def to_provider_tools_format(self, tools_schema: ToolsSchema) -> List[Dict[str, Any]]:
def get_llm_invocation_params(self, context: LLMContext) -> GeminiLLMInvocationParams:
"""Get Gemini-specific LLM invocation parameters from a universal LLM context.
Args:
context: The LLM context containing messages, tools, etc.
Returns:
Dictionary of parameters for Gemini's API.
"""
messages = self._from_standard_messages(context.messages)
return {
"system_instruction": messages.system_instruction,
"messages": messages.messages,
"tools": self.from_standard_tools(context.tools),
}
def to_provider_tools_format(self, tools_schema: ToolsSchema) -> List[Any]:
"""Convert tool schemas to Gemini's function-calling format.
Args:
@@ -39,3 +83,227 @@ class GeminiLLMAdapter(BaseLLMAdapter):
custom_gemini_tools = tools_schema.custom_tools.get(AdapterType.GEMINI, [])
return formatted_standard_tools + custom_gemini_tools
def get_messages_for_logging(self, context: LLMContext) -> List[dict[str, Any]]:
"""Get messages from a universal LLM context in a format ready for logging about Gemini.
Removes or truncates sensitive data like image content for safe logging.
Args:
context: The LLM context containing messages.
Returns:
List of messages in a format ready for logging about Gemini.
"""
# Get messages in Gemini's format
messages = self._from_standard_messages(context.messages).messages
# Sanitize messages for logging
messages_for_logging = []
for message in messages:
obj = message.to_json_dict()
try:
if "parts" in obj:
for part in obj["parts"]:
if "inline_data" in part:
part["inline_data"]["data"] = "..."
except Exception as e:
logger.debug(f"Error: {e}")
messages_for_logging.append(obj)
return messages_for_logging
@dataclass
class ConvertedMessages:
"""Container for converted messages.
Holds the converted messages in a format suitable for Gemini's API.
"""
messages: List[Content]
system_instruction: Optional[str] = None
def _from_standard_messages(
self, standard_messages: List[LLMContextMessage]
) -> ConvertedMessages:
"""Restructures messages to ensure proper Google format and message ordering.
This method handles conversion of OpenAI-formatted messages to Google format,
with special handling for function calls, function responses, and system messages.
System messages are added back to the context as user messages when needed.
The final message order is preserved as:
1. Function calls (from model)
2. Function responses (from user)
3. Text messages (converted from system messages)
Note:
System messages are only added back when there are no regular text
messages in the context, ensuring proper conversation continuity
after function calls.
"""
system_instruction = None
messages = []
# Process each message, preserving Google-formatted messages and converting others
for message in standard_messages:
if isinstance(message, Content):
# Keep existing Google-formatted messages (e.g., function calls/responses)
# TODO: this branch is probably not needed anymore, since LLMContext contains a universal format
messages.append(message)
continue
# Convert standard format to Google format
converted = self._from_standard_message(message)
if isinstance(converted, Content):
# Regular (non-system) message
messages.append(converted)
else:
# System instruction
system_instruction = converted
# Check if we only have function-related messages (no regular text)
has_regular_messages = any(
len(msg.parts) == 1
and getattr(msg.parts[0], "text", None)
and not getattr(msg.parts[0], "function_call", None)
and not getattr(msg.parts[0], "function_response", None)
for msg in messages
)
# Add system instruction back as a user message if we only have function messages
if system_instruction and not has_regular_messages:
messages.append(Content(role="user", parts=[Part(text=system_instruction)]))
# Remove any empty messages
messages = [m for m in messages if m.parts]
return self.ConvertedMessages(messages=messages, system_instruction=system_instruction)
def _from_standard_message(self, message: LLMContextMessage) -> Content | str:
"""Convert standard format message to Google Content object.
Handles conversion of text, images, and function calls to Google's
format.
System instructions are returned as a plain string.
Args:
message: Message in standard format.
Returns:
Content object with role and parts, or a plain string for system
messages.
Examples:
Standard text message::
{
"role": "user",
"content": "Hello there"
}
Converts to Google Content with::
Content(
role="user",
parts=[Part(text="Hello there")]
)
Standard function call message::
{
"role": "assistant",
"tool_calls": [
{
"function": {
"name": "search",
"arguments": '{"query": "test"}'
}
}
]
}
Converts to Google Content with::
Content(
role="model",
parts=[Part(function_call=FunctionCall(name="search", args={"query": "test"}))]
)
"""
role = message["role"]
content = message.get("content", [])
if role == "system":
# System instructions are returned as plain text
# TODO: here we've always assumed that system instructions are plain text...is that a safe assumption?
return content
elif role == "assistant":
role = "model"
parts = []
if message.get("tool_calls"):
for tc in message["tool_calls"]:
parts.append(
Part(
function_call=FunctionCall(
name=tc["function"]["name"],
args=json.loads(tc["function"]["arguments"]),
)
)
)
elif role == "tool":
role = "model"
try:
response = json.loads(message["content"])
if isinstance(response, dict):
response_dict = response
else:
response_dict = {"value": response}
except Exception as e:
# Response might not be JSON-deserializable.
# This occurs with a UserImageFrame, for example, where we get a plain "COMPLETED" string.
response_dict = {"value": message["content"]}
parts.append(
Part(
function_response=FunctionResponse(
name="tool_call_result", # seems to work to hard-code the same name every time
response=response_dict,
)
)
)
elif isinstance(content, str):
parts.append(Part(text=content))
elif isinstance(content, list):
for c in content:
if c["type"] == "text":
parts.append(Part(text=c["text"]))
elif c["type"] == "image_url":
parts.append(
Part(
inline_data=Blob(
mime_type="image/jpeg",
data=base64.b64decode(c["image_url"]["url"].split(",")[1]),
)
)
)
elif c["type"] == "input_audio":
input_audio = c["input_audio"]
parts.append(
Part(
inline_data=Blob(
mime_type="audio/wav",
data=(
bytes(
self.create_wav_header(
input_audio["sample_rate"],
input_audio["num_channels"],
16,
len(input_audio["data"]),
)
+ input_audio["data"]
)
),
)
)
)
message = Content(role=role, parts=parts)
return message

View File

@@ -763,7 +763,7 @@ class LLMAssistantAggregator(LLMContextAggregator):
del self._function_calls_in_progress[frame.request.tool_call_id]
# Update context with the image frame
await self._update_function_call_result(
self._update_function_call_result(
frame.request.function_name, frame.request.tool_call_id, "COMPLETED"
)
self._context.add_image_frame_message(

View File

@@ -16,19 +16,20 @@ import json
import os
import uuid
from dataclasses import dataclass
from typing import Any, Dict, List, Optional
from typing import Any, AsyncIterator, Dict, List, Optional
from loguru import logger
from PIL import Image
from pydantic import BaseModel, Field
from pipecat.adapters.services.gemini_adapter import GeminiLLMAdapter
from pipecat.adapters.services.gemini_adapter import GeminiLLMAdapter, GeminiLLMInvocationParams
from pipecat.frames.frames import (
AudioRawFrame,
Frame,
FunctionCallCancelFrame,
FunctionCallInProgressFrame,
FunctionCallResultFrame,
LLMContextFrame,
LLMFullResponseEndFrame,
LLMFullResponseStartFrame,
LLMMessagesFrame,
@@ -38,6 +39,7 @@ from pipecat.frames.frames import (
VisionImageRawFrame,
)
from pipecat.metrics.metrics import LLMTokenUsage
from pipecat.processors.aggregators.llm_context import LLMContext
from pipecat.processors.aggregators.llm_response import (
LLMAssistantAggregatorParams,
LLMUserAggregatorParams,
@@ -67,6 +69,7 @@ try:
FunctionCall,
FunctionResponse,
GenerateContentConfig,
GenerateContentResponse,
HttpOptions,
Part,
)
@@ -436,11 +439,20 @@ class GoogleLLMContext(OpenAILLMContext):
)
elif role == "tool":
role = "model"
try:
response = json.loads(message["content"])
if isinstance(response, dict):
response_dict = response
else:
response_dict = {"value": response}
except Exception as e:
# Response might not be JSON-deserializable (e.g. plain text).
response_dict = {"value": message["content"]}
parts.append(
Part(
function_response=FunctionResponse(
name="tool_call_result", # seems to work to hard-code the same name every time
response=json.loads(message["content"]),
response=response_dict,
)
)
)
@@ -636,9 +648,8 @@ class GoogleLLMService(LLMService):
"""Google AI (Gemini) LLM service implementation.
This class implements inference with Google's AI models, translating internally
from OpenAILLMContext to the messages format expected by the Google AI model.
We use OpenAILLMContext as a lingua franca for all LLM services to enable
easy switching between different LLMs.
from an OpenAILLMContext or a universal LLMContext to the messages format
expected by the Google AI model.
"""
# Overriding the default adapter to use the Gemini one.
@@ -740,8 +751,89 @@ class GoogleLLMService(LLMService):
except Exception as e:
logger.exception(f"Failed to unset thinking budget: {e}")
async def _stream_content(
self, params_from_context: GeminiLLMInvocationParams
) -> AsyncIterator[GenerateContentResponse]:
messages = params_from_context["messages"]
if (
params_from_context["system_instruction"]
and self._system_instruction != params_from_context["system_instruction"]
):
logger.debug(f"System instruction changed: {params_from_context['system_instruction']}")
self._system_instruction = params_from_context["system_instruction"]
tools = []
if params_from_context["tools"]:
tools = params_from_context["tools"]
elif self._tools:
tools = self._tools
tool_config = None
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
}
if self._settings["extra"]:
generation_params.update(self._settings["extra"])
# 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
)
await self.start_ttfb_metrics()
return await self._client.aio.models.generate_content_stream(
model=self._model_name,
contents=messages,
config=generation_config,
)
async def _stream_content_specific_context(
self, context: OpenAILLMContext
) -> AsyncIterator[GenerateContentResponse]:
logger.debug(
# f"{self}: Generating chat [{self._system_instruction}] | [{context.get_messages_for_logging()}]"
f"{self}: Generating chat from OpenAI context [{context.get_messages_for_logging()}]"
)
params = GeminiLLMInvocationParams(
messages=context.messages,
system_instruction=context.system_message,
tools=context.tools,
)
return await self._stream_content(params)
async def _stream_content_universal_context(
self, context: LLMContext
) -> AsyncIterator[GenerateContentResponse]:
adapter = self.get_llm_adapter()
logger.debug(
# f"{self}: Generating chat [{self._system_instruction}] | [{context.get_messages_for_logging()}]"
f"{self}: Generating chat from universal context [{adapter.get_messages_for_logging(context)}]"
)
params: GeminiLLMInvocationParams = adapter.get_llm_invocation_params(context)
return await self._stream_content(params)
@traced_llm
async def _process_context(self, context: OpenAILLMContext):
async def _process_context(self, context: OpenAILLMContext | LLMContext):
await self.push_frame(LLMFullResponseStartFrame())
prompt_tokens = 0
@@ -754,55 +846,11 @@ class GoogleLLMService(LLMService):
search_result = ""
try:
logger.debug(
# f"{self}: Generating chat [{self._system_instruction}] | [{context.get_messages_for_logging()}]"
f"{self}: Generating chat [{context.get_messages_for_logging()}]"
)
messages = context.messages
if context.system_message and self._system_instruction != context.system_message:
logger.debug(f"System instruction changed: {context.system_message}")
self._system_instruction = context.system_message
tools = []
if context.tools:
tools = context.tools
elif self._tools:
tools = self._tools
tool_config = None
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
}
if self._settings["extra"]:
generation_params.update(self._settings["extra"])
# 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
)
await self.start_ttfb_metrics()
response = await self._client.aio.models.generate_content_stream(
model=self._model_name,
contents=messages,
config=generation_config,
# Generate content using either OpenAILLMContext or universal LLMContext
response = await (
self._stream_content_specific_context(context)
if isinstance(context, OpenAILLMContext)
else self._stream_content_universal_context(context)
)
function_calls = []
@@ -915,7 +963,12 @@ class GoogleLLMService(LLMService):
if isinstance(frame, OpenAILLMContextFrame):
context = GoogleLLMContext.upgrade_to_google(frame.context)
elif isinstance(frame, LLMContextFrame):
# Handle universal (LLM-agnostic) LLM context frames
context = frame.context
elif isinstance(frame, LLMMessagesFrame):
# NOTE: LLMMessagesFrame is deprecated, so we don't support the newer universal
# LLMContext with it
context = GoogleLLMContext(frame.messages)
elif isinstance(frame, VisionImageRawFrame):
# This is only useful in very simple pipelines because it creates

View File

@@ -285,7 +285,6 @@ class BaseOpenAILLMService(LLMService):
)
params: OpenAILLMInvocationParams = adapter.get_llm_invocation_params(context)
chunks = await self.get_chat_completions(params)
return chunks
@@ -302,16 +301,12 @@ class BaseOpenAILLMService(LLMService):
await self.start_ttfb_metrics()
if isinstance(context, OpenAILLMContext):
# Use OpenAI-specific context
chunk_stream: AsyncStream[ChatCompletionChunk] = await self._stream_chat_completions(
context
)
else:
# Use universal (LLM-agnostic) context
chunk_stream: AsyncStream[
ChatCompletionChunk
] = await self._stream_chat_completions_universal_context(context)
# Generate chat completions using either OpenAILLMContext or universal LLMContext
chunk_stream = await (
self._stream_chat_completions(context)
if isinstance(context, OpenAILLMContext)
else self._stream_chat_completions_universal_context(context)
)
async for chunk in chunk_stream:
if chunk.usage: