fix imports

Author: MarkDaoust

google/generativeai/types/__init__.py

@@ -17,6 +17,11 @@
 from google.generativeai.types.discuss_types import *
 from google.generativeai.types.model_types import *
 from google.generativeai.types.text_types import *
+from google.generativeai.types.citation_types import *
+from google.generativeai.types.safety_types import *
 
 del discuss_types
 del model_types
+del text_types
+del citation_types
+del safety_types

google/generativeai/types/citation_types.py

@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import Optional, List
 
 from google.ai import generativelanguage as glm
 from google.generativeai import docstring_utils

google/generativeai/types/discuss_types.py

@@ -20,6 +20,7 @@
 
 import google.ai.generativelanguage as glm
 from google.generativeai.types import safety_types
+from google.generativeai.types import citation_types
 
 __all__ = [
     "MessageDict",
@@ -155,17 +156,17 @@ class ChatResponse(abc.ABC):
     model: str
     context: str
     examples: List[ExampleDict]
-    messages: List[MessageDict]
+    messages: List[Optional[MessageDict]]
     temperature: Optional[float]
     candidate_count: Optional[int]
     candidates: List[MessageDict]
     top_p: Optional[float] = None
     top_k: Optional[float] = None
-    filters: List[safety.ContentFilter]
+    filters: List[safety_types.ContentFilter]
 
     @property
     @abc.abstractmethod
-    def last(self) -> str:
+    def last(self) -> Optional[str]:
         """A settable property that provides simple access to the last response string
 
         A shortcut for `response.messages[0]['content']`.

google/generativeai/types/safety_types.py

@@ -15,6 +15,7 @@
 
 import enum
 from google.ai import generativelanguage as glm
+from google.generativeai import docstring_utils
 from typing import TypedDict
 
 __all__ = [
@@ -31,7 +32,7 @@
 # These are basic python enums, it's okay to expose them
 HarmCategory = glm.HarmCategory
 HarmProbability = glm.SafetyRating.HarmProbability
-HarmBlockThreshold = glm.SafetySetting
+HarmBlockThreshold = glm.SafetySetting.HarmBlockThreshold
 BlockReason = glm.ContentFilter.BlockedReason
 
 

google/generativeai/types/text_types.py

@@ -17,28 +17,38 @@
 import dataclasses
 from typing import Any, Dict, Optional, List, Iterator, TypedDict
 
-__all__ = ["Completion"]
+from google.generativeai.types import safety_types
+from google.generativeai.types import citation_types
 
 
-class TextCandidate(TypedDict, total=False):
+__all__ = ["TextResponse"]
+
+
+class TextCompletion(TypedDict, total=False):
     output: str
+    safety_ratings: Optional[List[safety_types.SafetyRatingDict]]
+    citation_metadata: Optional[citation_types.CitationMetadataDict]
 
 
 @dataclasses.dataclass(init=False)
-class Completion(abc.ABC):
+class TextResponse(abc.ABC):
     """A text completion given a prompt from the model.
 
-    * Use `completion.candidates` to access all of the text completion options generated by the model.
+    Use `GenerateTextResponse.candidates` to access all the completions generated by the model.
 
     Attributes:
         candidates: A list of candidate text completions generated by the model.
     """
 
-    candidates: List[TextCandidate]
+    candidates: List[TextCompletion]
     result: Optional[str]
+    filters: Optional[list[safety_types.ContentFilter]]
+    safety_feedback: Optional[list[safety_types.SafetyFeedbackDict]]
 
     def to_dict(self) -> Dict[str, Any]:
         result = {
             "candidates": self.candidates,
+            "filters": self.filters,
+            "safety_feedback": self.safety_feedback,
         }
         return result