feat: enable inline context in grounding to TextGenerationModel predict.

PiperOrigin-RevId: 597296033
googleapis · Jan 10, 2024 · a75e81c · a75e81c
1 parent ad8d9c1
commit a75e81c
Showing 1 changed file with 52 additions and 6 deletions.
diff --git a/vertexai/language_models/_language_models.py b/vertexai/language_models/_language_models.py
@@ -749,6 +749,27 @@ def _to_grounding_source_dict(self) -> Dict[str, Any]:
         }
 
 
+@dataclasses.dataclass
+class InlineContext(_GroundingSourceBase):
+    """InlineContext represents a grounding source using provided inline context.
+    Attributes:
+        inline_context: The content used as inline context.
+    """
+
+    inline_context: str
+    _type: str = dataclasses.field(default="INLINE", init=False, repr=False)
+
+    def _to_grounding_source_dict(self) -> Dict[str, Any]:
+        return {
+            "sources": [
+                {
+                    "type": self._type,
+                }
+            ],
+            "inlineContext": self.inline_context,
+        }
+
+
 @dataclasses.dataclass
 class VertexAISearch(_GroundingSourceBase):
     """VertexAISearchDatastore represents a grounding source using Vertex AI Search datastore
@@ -792,6 +813,7 @@ class GroundingSource:
 
     WebSearch = WebSearch
     VertexAISearch = VertexAISearch
+    InlineContext = InlineContext
 
 
 @dataclasses.dataclass
@@ -976,7 +998,11 @@ def predict(
         stop_sequences: Optional[List[str]] = None,
         candidate_count: Optional[int] = None,
         grounding_source: Optional[
-            Union[GroundingSource.WebSearch, GroundingSource.VertexAISearch]
+            Union[
+                GroundingSource.WebSearch,
+                GroundingSource.VertexAISearch,
+                GroundingSource.InlineContext,
+            ]
         ] = None,
         logprobs: Optional[int] = None,
         presence_penalty: Optional[float] = None,
@@ -1053,7 +1079,11 @@ async def predict_async(
         stop_sequences: Optional[List[str]] = None,
         candidate_count: Optional[int] = None,
         grounding_source: Optional[
-            Union[GroundingSource.WebSearch, GroundingSource.VertexAISearch]
+            Union[
+                GroundingSource.WebSearch,
+                GroundingSource.VertexAISearch,
+                GroundingSource.InlineContext,
+            ]
         ] = None,
         logprobs: Optional[int] = None,
         presence_penalty: Optional[float] = None,
@@ -1284,7 +1314,11 @@ def _create_text_generation_prediction_request(
     stop_sequences: Optional[List[str]] = None,
     candidate_count: Optional[int] = None,
     grounding_source: Optional[
-        Union[GroundingSource.WebSearch, GroundingSource.VertexAISearch]
+        Union[
+            GroundingSource.WebSearch,
+            GroundingSource.VertexAISearch,
+            GroundingSource.InlineContext,
+        ]
     ] = None,
     logprobs: Optional[int] = None,
     presence_penalty: Optional[float] = None,
@@ -2136,7 +2170,11 @@ def _prepare_request(
         stop_sequences: Optional[List[str]] = None,
         candidate_count: Optional[int] = None,
         grounding_source: Optional[
-            Union[GroundingSource.WebSearch, GroundingSource.VertexAISearch]
+            Union[
+                GroundingSource.WebSearch,
+                GroundingSource.VertexAISearch,
+                GroundingSource.InlineContext,
+            ]
         ] = None,
     ) -> _PredictionRequest:
         """Prepares a request for the language model.
@@ -2289,7 +2327,11 @@ def send_message(
         stop_sequences: Optional[List[str]] = None,
         candidate_count: Optional[int] = None,
         grounding_source: Optional[
-            Union[GroundingSource.WebSearch, GroundingSource.VertexAISearch]
+            Union[
+                GroundingSource.WebSearch,
+                GroundingSource.VertexAISearch,
+                GroundingSource.InlineContext,
+            ]
         ] = None,
     ) -> "MultiCandidateTextGenerationResponse":
         """Sends message to the language model and gets a response.
@@ -2352,7 +2394,11 @@ async def send_message_async(
         stop_sequences: Optional[List[str]] = None,
         candidate_count: Optional[int] = None,
         grounding_source: Optional[
-            Union[GroundingSource.WebSearch, GroundingSource.VertexAISearch]
+            Union[
+                GroundingSource.WebSearch,
+                GroundingSource.VertexAISearch,
+                GroundingSource.InlineContext,
+            ]
         ] = None,
     ) -> "MultiCandidateTextGenerationResponse":
         """Asynchronously sends message to the language model and gets a response.