Add stop_sequences to Google models (#1935)

Author: Kludex

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -366,6 +366,8 @@ def _settings_to_generation_config(model_settings: GeminiModelSettings) -> _Gemi
     config: _GeminiGenerationConfig = {}
     if (max_tokens := model_settings.get('max_tokens')) is not None:
         config['max_output_tokens'] = max_tokens
+    if (stop_sequences := model_settings.get('stop_sequences')) is not None:
+        config['stop_sequences'] = stop_sequences  # pragma: no cover
     if (temperature := model_settings.get('temperature')) is not None:
         config['temperature'] = temperature
     if (top_p := model_settings.get('top_p')) is not None:

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -260,6 +260,7 @@ async def _generate_content(
             temperature=model_settings.get('temperature'),
             top_p=model_settings.get('top_p'),
             max_output_tokens=model_settings.get('max_tokens'),
+            stop_sequences=model_settings.get('stop_sequences'),
             presence_penalty=model_settings.get('presence_penalty'),
             frequency_penalty=model_settings.get('frequency_penalty'),
             safety_settings=model_settings.get('google_safety_settings'),

pydantic_ai_slim/pydantic_ai/settings.py

@@ -139,6 +139,7 @@ class ModelSettings(TypedDict, total=False):
     * Mistral
     * Groq
     * Cohere
+    * Google
     """
 
     extra_headers: dict[str, str]

tests/cassettes/test_settings/test_stop_settings[anthropic].yaml

@@ -8,7 +8,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '193'
+      - '272'
       content-type:
       - application/json
       host:
@@ -18,7 +18,7 @@ interactions:
       max_tokens: 1024
       messages:
       - content:
-        - text: What is the capital of France?
+        - text: What is the capital of France? Give me an answer that contains the word "Paris", but is not the first word.
           type: text
         role: user
       model: claude-3-5-sonnet-latest
@@ -31,7 +31,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '359'
+      - '356'
       content-type:
       - application/json
       strict-transport-security:
@@ -40,9 +40,9 @@ interactions:
       - chunked
     parsed_body:
       content:
-      - text: 'The capital of France is '
+      - text: 'The beautiful city of '
         type: text
-      id: msg_01SVFZqkSLhUPkAA4w6rhzFD
+      id: msg_01376yZQxHcw9pER2Ab2SvQb
       model: claude-3-5-sonnet-20241022
       role: assistant
       stop_reason: stop_sequence
@@ -51,8 +51,8 @@ interactions:
       usage:
         cache_creation_input_tokens: 0
         cache_read_input_tokens: 0
-        input_tokens: 14
-        output_tokens: 6
+        input_tokens: 32
+        output_tokens: 5
         service_tier: standard
     status:
       code: 200

tests/cassettes/test_settings/test_stop_settings[bedrock].yaml

@@ -1,16 +1,16 @@
 interactions:
 - request:
-    body: '{"messages": [{"role": "user", "content": [{"text": "What is the capital of France?"}]}], "system": [], "inferenceConfig":
-      {"stopSequences": ["Paris"]}}'
+    body: '{"messages": [{"role": "user", "content": [{"text": "What is the capital of France? Give me an answer that contains
+      the word \"Paris\", but is not the first word."}]}], "system": [], "inferenceConfig": {"stopSequences": ["Paris"]}}'
     headers:
       amz-sdk-invocation-id:
       - !!binary |
-        YzVkZjljOTMtMDQ1Zi00NWE0LWJhY2YtMDAwMjdjYTg1NmRl
+        YWEyMWExMDYtNjJjMy00MmNiLTlkNjUtYTEwYzYzYTIxNGFk
       amz-sdk-request:
       - !!binary |
         YXR0ZW1wdD0x
       content-length:
-      - '152'
+      - '231'
       content-type:
       - !!binary |
         YXBwbGljYXRpb24vanNvbg==
@@ -21,22 +21,22 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '209'
+      - '197'
       content-type:
       - application/json
     parsed_body:
       metrics:
-        latencyMs: 179
+        latencyMs: 189
       output:
         message:
           content:
-          - text: The capital of France is Paris
+          - text: The city of Paris
           role: assistant
       stopReason: end_turn
       usage:
-        inputTokens: 7
-        outputTokens: 6
-        totalTokens: 13
+        inputTokens: 25
+        outputTokens: 4
+        totalTokens: 29
     status:
       code: 200
       message: OK

tests/cassettes/test_settings/test_stop_settings[cohere].yaml

@@ -8,15 +8,15 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '140'
+      - '219'
       content-type:
       - application/json
       host:
       - api.cohere.com
     method: POST
     parsed_body:
       messages:
-      - content: What is the capital of France?
+      - content: What is the capital of France? Give me an answer that contains the word "Paris", but is not the first word.
         role: user
       model: command-r-plus
       stop_sequences:
@@ -32,34 +32,34 @@ interactions:
       cache-control:
       - no-cache, no-store, no-transform, must-revalidate, private, max-age=0
       content-length:
-      - '280'
+      - '297'
       content-type:
       - application/json
       expires:
       - Thu, 01 Jan 1970 00:00:00 UTC
       num_chars:
-      - '1200'
+      - '521'
       num_tokens:
-      - '12'
+      - '34'
       pragma:
       - no-cache
       vary:
       - Origin
     parsed_body:
       finish_reason: STOP_SEQUENCE
-      id: 2ed6908c-e5cb-4063-b2a3-0ac4990d1b85
+      id: 02069ccc-65d6-4286-ba8c-097e317431c5
       message:
         content:
-        - text: The capital of France is
+        - text: The capital of France is a city known as
           type: text
         role: assistant
       usage:
         billed_units:
-          input_tokens: 7
-          output_tokens: 5
+          input_tokens: 25
+          output_tokens: 9
         tokens:
-          input_tokens: 200
-          output_tokens: 7
+          input_tokens: 91
+          output_tokens: 11
     status:
       code: 200
       message: OK

tests/cassettes/test_settings/test_stop_settings[google].yaml

@@ -0,0 +1,65 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '216'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+    method: POST
+    parsed_body:
+      contents:
+      - parts:
+        - text: What is the capital of France? Give me an answer that contains the word "Paris", but is not the first word.
+        role: user
+      generationConfig:
+        stopSequences:
+        - Paris
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent
+  response:
+    headers:
+      alt-svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      content-length:
+      - '677'
+      content-type:
+      - application/json; charset=UTF-8
+      server-timing:
+      - gfet4t7; dur=362
+      transfer-encoding:
+      - chunked
+      vary:
+      - Origin
+      - X-Origin
+      - Referer
+    parsed_body:
+      candidates:
+      - avgLogprobs: -0.5998413562774658
+        content:
+          parts:
+          - text: 'The most iconic city in France is '
+          role: model
+        finishReason: STOP
+      modelVersion: gemini-1.5-flash
+      responseId: UB5DaMfEN7jFnvgPocrJaA
+      usageMetadata:
+        candidatesTokenCount: 8
+        candidatesTokensDetails:
+        - modality: TEXT
+          tokenCount: 8
+        promptTokenCount: 25
+        promptTokensDetails:
+        - modality: TEXT
+          tokenCount: 25
+        totalTokenCount: 33
+    status:
+      code: 200
+      message: OK
+version: 1

tests/cassettes/test_settings/test_stop_settings[groq].yaml

@@ -8,15 +8,15 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '136'
+      - '215'
       content-type:
       - application/json
       host:
       - api.groq.com
     method: POST
     parsed_body:
       messages:
-      - content: What is the capital of France?
+      - content: What is the capital of France? Give me an answer that contains the word "Paris", but is not the first word.
         role: user
       model: llama3-8b-8192
       n: 1
@@ -33,38 +33,38 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '562'
+      - '603'
       content-type:
       - application/json
       transfer-encoding:
       - chunked
       vary:
-      - Origin, Accept-Encoding
+      - Origin
     parsed_body:
       choices:
       - finish_reason: stop
         index: 0
         logprobs: null
         message:
-          content: 'The capital of France is '
+          content: "Bien sûr!\n\nThe lovely city that is the capital of France is "
           role: assistant
-      created: 1744190899
-      id: chatcmpl-9de98446-dccf-439b-bcb8-996ee77471d8
+      created: 1749229135
+      id: chatcmpl-69930f13-fe23-4584-9d09-c1e1612a2183
       model: llama3-8b-8192
       object: chat.completion
-      system_fingerprint: fp_dadc9d6142
+      system_fingerprint: fp_179b0f92c9
       usage:
-        completion_time: 0.006666667
-        completion_tokens: 8
-        prompt_time: 0.003914844
-        prompt_tokens: 17
-        queue_time: 0.5322581550000001
-        total_time: 0.010581511
-        total_tokens: 25
+        completion_time: 0.020833333
+        completion_tokens: 25
+        prompt_time: 0.004837532
+        prompt_tokens: 35
+        queue_time: 0.022124504000000003
+        total_time: 0.025670865
+        total_tokens: 60
       usage_breakdown:
         models: null
       x_groq:
-        id: req_01jrcy20hceg1t9cpsg7cmj7f9
+        id: req_01jx32wsk4fwf98jtfjm1agphs
     status:
       code: 200
       message: OK

tests/cassettes/test_settings/test_stop_settings[mistral].yaml

@@ -8,15 +8,15 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '153'
+      - '232'
       content-type:
       - application/json
       host:
       - api.mistral.ai
     method: POST
     parsed_body:
       messages:
-      - content: What is the capital of France?
+      - content: What is the capital of France? Give me an answer that contains the word "Paris", but is not the first word.
         role: user
       model: ministral-8b-latest
       n: 1
@@ -37,14 +37,6 @@ interactions:
       - '322'
       content-type:
       - application/json
-      ratelimitbysize-limit:
-      - '500000'
-      ratelimitbysize-query-cost:
-      - '32008'
-      ratelimitbysize-remaining:
-      - '467992'
-      ratelimitbysize-reset:
-      - '42'
       transfer-encoding:
       - chunked
     parsed_body:
@@ -55,14 +47,14 @@ interactions:
           content: 'The capital of France is '
           role: assistant
           tool_calls: null
-      created: 1744190898
-      id: 0bbc8cf8fc76455fae759fb9109f8547
+      created: 1749229172
+      id: 59bee455a928496eb22bbf8b8b817252
       model: ministral-8b-latest
       object: chat.completion
       usage:
         completion_tokens: 6
-        prompt_tokens: 10
-        total_tokens: 16
+        prompt_tokens: 28
+        total_tokens: 34
     status:
       code: 200
       message: OK