[AI] Server Prompt Templates #15402
[AI] Server Prompt Templates #15402
Conversation
Generated by 🚫 Danger |
| | ||
| /// A chat session that allows for conversation with a model. | ||
| @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) | ||
| public class TemplateChatSession { |
There was a problem hiding this comment.
Just a quick thing I noticed: public final class TemplateChatSession: Sendable. This would likely require some refactoring with the history. I'm wondering if we should extract
firebase-ios-sdk/FirebaseAI/Sources/Chat.swift
Lines 29 to 52 in 455d291
| /gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces server-side prompt templates, a significant feature addition. The changes include new models (TemplateGenerativeModel, TemplateImagenModel), request types, and chat session logic to support using templates with variables. The code is well-structured, with good separation of concerns, such as moving history management into its own History class. I've identified a couple of areas where force unwraps on URL creation could be made safer. Overall, this is a solid implementation with good test coverage.
| } | ||
| let templateName = template.hasSuffix(".prompt") ? template : "\(template).prompt" | ||
| urlString += "/templates/\(templateName):\(ImageAPIMethod.generateImages.rawValue)" | ||
| return URL(string: urlString)! |
There was a problem hiding this comment.
Using a force unwrap ! for URL creation can lead to a runtime crash if the URL string is malformed for any reason. It's safer to use guard let to safely unwrap the optional URL. This makes the code more robust and easier to debug by providing a more descriptive error message if the URL construction fails.
guard let url = URL(string: urlString) else { fatalError("Could not construct URL for GenerateImagesRequest. URL string was: \(urlString)") } return urlThere was a problem hiding this comment.
Don't want a fatal error and the force unwrap is consistent with GenerateContentRequest
| // TODO: Fix this. | ||
| urlString += "?alt=sse" | ||
| } | ||
| return URL(string: urlString)! |
There was a problem hiding this comment.
Using a force unwrap ! for URL creation can lead to a runtime crash if the URL string is malformed. It's safer to use guard let to safely unwrap the optional URL and provide a descriptive error message if it fails. This improves code robustness and debuggability.
guard let url = URL(string: urlString) else { fatalError("Could not construct URL for TemplateGenerateContentRequest. URL string was: \(urlString)") } return url | /gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces support for server-side prompt templates for both generative and image generation models. This is a significant feature addition, including new model types like TemplateGenerativeModel and TemplateImagenModel, new request structures, and a TemplateChatSession for conversational AI with templates. The changes also involve some nice refactoring, such as abstracting history management into a dedicated History class. The overall implementation is well-structured and comes with a good set of unit and integration tests. My feedback is focused on improving robustness by avoiding force-unwraps in URL construction, ensuring proper URL encoding of user-provided data, and making the TemplateVariable type more versatile by handling Float values. I've also included a minor suggestion to clean up test code.
| public var url: URL { | ||
| var urlString = | ||
| "\(apiConfig.service.endpoint.rawValue)/\(apiConfig.version.rawValue)/projects/\(projectID)" | ||
| if case let .vertexAI(_, location) = apiConfig.service { | ||
| urlString += "/locations/\(location)" | ||
| } | ||
| let templateName = template.hasSuffix(".prompt") ? template : "\(template).prompt" | ||
| urlString += "/templates/\(templateName):\(ImageAPIMethod.generateImages.rawValue)" | ||
| return URL(string: urlString)! | ||
| } |
There was a problem hiding this comment.
The url computed property uses force unwrapping (!) to create a URL object. This can lead to a runtime crash if urlString is not a valid URL. The template parameter, which is part of the URL, comes from user input and is not being percent-encoded. If it contains characters that are not allowed in a URL path, URL(string: urlString) will return nil, causing a crash.
It's safer to percent-encode the template name and handle the potential failure of URL creation gracefully instead of force-unwrapping.
| public var url: URL { | |
| var urlString = | |
| "\(apiConfig.service.endpoint.rawValue)/\(apiConfig.version.rawValue)/projects/\(projectID)" | |
| if case let .vertexAI(_, location) = apiConfig.service { | |
| urlString += "/locations/\(location)" | |
| } | |
| let templateName = template.hasSuffix(".prompt") ? template : "\(template).prompt" | |
| urlString += "/templates/\(templateName):\(ImageAPIMethod.generateImages.rawValue)" | |
| return URL(string: urlString)! | |
| } | |
| public var url: URL { | |
| var urlString = | |
| "\(apiConfig.service.endpoint.rawValue)/\(apiConfig.version.rawValue)/projects/\(projectID)" | |
| if case let .vertexAI(_, location) = apiConfig.service { | |
| urlString += "/locations/\(location)" | |
| } | |
| let templateName = template.hasSuffix(".prompt") ? template : "\(template).prompt" | |
| guard let encodedTemplateName = templateName.addingPercentEncoding(withAllowedCharacters: .urlPathAllowed) else { | |
| fatalError("Failed to percent-encode template name: \(templateName)") | |
| } | |
| urlString += "/templates/\(encodedTemplateName):\(ImageAPIMethod.generateImages.rawValue)" | |
| guard let url = URL(string: urlString) else { | |
| fatalError("Could not construct a valid URL from string: \(urlString). This is a bug.") | |
| } | |
| return url | |
| } |
| var url: URL { | ||
| var urlString = | ||
| "\(apiConfig.service.endpoint.rawValue)/\(apiConfig.version.rawValue)/projects/\(projectID)" | ||
| if case let .vertexAI(_, location) = apiConfig.service { | ||
| urlString += "/locations/\(location)" | ||
| } | ||
| let templateName = template.hasSuffix(".prompt") ? template : "\(template).prompt" | ||
| urlString += "/templates/\(templateName):templateGenerateContent" | ||
| if stream { | ||
| urlString += "?alt=sse" | ||
| } | ||
| return URL(string: urlString)! | ||
| } |
There was a problem hiding this comment.
The url computed property uses force unwrapping (!) which can cause a crash if the URL string is invalid. The template name is part of the URL and comes from user input, but it's not percent-encoded. This could lead to an invalid URL if the template name contains special characters.
Please percent-encode the template name and avoid force-unwrapping to make the URL creation more robust.
var url: URL { var urlString = "\(apiConfig.service.endpoint.rawValue)/\(apiConfig.version.rawValue)/projects/\(projectID)" if case let .vertexAI(_, location) = apiConfig.service { urlString += "/locations/\(location)" } let templateName = template.hasSuffix(".prompt") ? template : "\(template).prompt" guard let encodedTemplateName = templateName.addingPercentEncoding(withAllowedCharacters: .urlPathAllowed) else { fatalError("Failed to percent-encode template name: \(templateName)") } urlString += "/templates/\(encodedTemplateName):templateGenerateContent" if stream { urlString += "?alt=sse" } guard let url = URL(string: urlString) else { fatalError("Could not construct a valid URL from string: \(urlString). This is a bug.") } return url } FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift Outdated Show resolved Hide resolved
| /gemini summarize |
| @andrewheard @cynthiajoan Still some test tweaks once the backend supports both API providers with the same templates and this won't merge until after M172, but should now be ready for review. |
Updated `GenerateImagesRequest` to encode the `variables` property as `inputs`, similar to how `TemplateGenerateContentRequest` handles it. This ensures consistency across the API.
Updated `GenerateImagesRequest` to correctly construct the URL, including the `projectID` and `location`. This mirrors the URL construction in `TemplateGenerateContentRequest` and fixes a 404 error that was occurring in the `testGenerateImages` integration test. The `TemplateImagenModel` was also updated to pass the `projectID` to the `GenerateImagesRequest` initializer.
3 participants
Apple implementation of Server Prompt Templates inspired by Flutterfire PR.
Googlers: see API review at go/firebase-ai-server-prompt-template
Summary of Changes
This pull request significantly enhances the Firebase AI SDK by introducing support for server-side prompt templates. This allows developers to define and use reusable templates for generative AI models, making it easier to manage complex prompts and integrate dynamic data. The changes include new model types for generative and image AI, a dedicated chat session for template-based conversations, and a refactored, thread-safe history management system. These additions provide a more robust and flexible framework for building AI-powered applications.
Highlights
TemplateGenerativeModelandTemplateImagenModelclasses to support server-side prompt templates, allowing for more flexible and dynamic AI interactions.TemplateChatSessionto enable conversational interactions with generative models using predefined templates and dynamic variables, supporting both single messages and streaming responses.Historyclass, improving modularity, reusability, and thread safety for managing conversation history in both standard and template-based chat sessions.TemplateVariableenum to handle various data types (String, Int, Double, Bool, Array, Dictionary) for substituting values into prompt templates, including automatic conversion forFloattoDouble.TemplateGenerateContentRequest,TemplateGenerateImagesRequest) and internal API method enums (APIMethod,ImageAPIMethod) to support the new template-based API endpoints.Changelog
Historyclass for managing chat history, removing internal history management logic and related locks.appendHistorycalls to use the_historyinstance of the newHistoryclass.templateGenerativeModel()andtemplateImagenModel()to initialize new template-based AI models.APIMethodenum, which has been moved to a new, dedicated internal file.firebaseInfoproperty fromprivate lettoletto allow access from newly introduced template models.Historyclass to encapsulate chat history management, including thread-safe appending ofModelContentand aggregation of content chunks.TemplateChatSessionfor managing chat conversations with server-side prompt templates, providingsendMessageandsendMessageStreammethods.TemplateGenerativeModelfor interacting with generative AI models using templates, including methods for content generation and starting chat sessions.TemplateImagenModelfor interacting with image generation models using templates.TemplateVariableenum to handle various data types for template variables, including a conversion fromFloattoDoubleduring initialization.APIMethodenum (generateContent, streamGenerateContent, countTokens) here fromGenerateContentRequest.swift.ImageAPIMethodenum specifically for image generation API methods, such astemplatePredict.TemplateGenerativeModelandTemplateImagenModel, covering text generation, image generation, and chat sessions with templates and media.TemplateChatSessionclass, verifyingsendMessageandsendMessageStreamfunctionality.TemplateGenerativeModelclass, testinggenerateContentandgenerateContentStreammethods.TemplateImagenModelclass, specifically testing thegenerateImagesmethod.httpRequestHandlerto include anisTemplateRequestparameter, allowing for conditional URL path assertions to support template-based requests in mock tests.collectTextFromStreamto simplify collecting text fromAsyncThrowingStreamresponses in tests.Activity
andrewheardsuggested extracting chat history management into a separate type, which was subsequently implemented bypaulb777.gemini-code-assist[bot]highlighted potential issues with force unwrapping URLs in new request types and suggested percent-encoding template names for robustness.paulb777acknowledged these, noting consistency with existing code.gemini-code-assist[bot]to handleFloatvalues inTemplateVariableinitialization was addressed and implemented in the changes.gemini-code-assist[bot]included aTODOfor streaming and a hardcoded path in a test utility.