Updated to latest API release

graphlit · Jan 5, 2025 · b940a92 · b940a92
1 parent fc407cb
commit b940a92
Show file tree

Hide file tree

Showing 7 changed files with 28 additions and 12 deletions.
diff --git a/documents/specification/QueryModels.graphql b/documents/specification/QueryModels.graphql
@@ -1,10 +1,12 @@
-query QueryModels {
-  models {
+query QueryModels($filter: ModelFilter) {
+  models(filter: $filter) {
     results {
+      uri
       name
       type
       serviceType
       model
+      modelType
       description
       availableOn
       features {

diff --git a/graphlit_api/__init__.py b/graphlit_api/__init__.py
@@ -1178,6 +1178,7 @@
     MistralModelPropertiesInput,
     MistralModelPropertiesUpdateInput,
     ModelDocumentPreparationPropertiesInput,
+    ModelFilter,
     ModelImageExtractionPropertiesInput,
     ModelTextExtractionPropertiesInput,
     NamedEntityReferenceInput,
@@ -3482,6 +3483,7 @@
     "MistralModelPropertiesUpdateInput",
     "MistralModels",
     "ModelDocumentPreparationPropertiesInput",
+    "ModelFilter",
     "ModelImageExtractionPropertiesInput",
     "ModelServiceTypes",
     "ModelTextExtractionPropertiesInput",

diff --git a/graphlit_api/client.py b/graphlit_api/client.py
@@ -256,6 +256,7 @@
     MedicalTherapyUpdateInput,
     MicrosoftTeamsChannelsInput,
     MicrosoftTeamsTeamsInput,
+    ModelFilter,
     ObservationInput,
     ObservationReferenceInput,
     ObservationUpdateInput,
@@ -4694,8 +4695,10 @@ async def prompt_specifications(
         data = self.get_data(response)
         return PromptSpecifications.model_validate(data)
 
-    async def query_models(self, **kwargs: Any) -> QueryModels:
-        variables: Dict[str, object] = {}
+    async def query_models(
+        self, filter: Union[Optional[ModelFilter], UnsetType] = UNSET, **kwargs: Any
+    ) -> QueryModels:
+        variables: Dict[str, object] = {"filter": filter}
         response = await self.execute(
             query=QUERY_MODELS_GQL,
             operation_name="QueryModels",

diff --git a/graphlit_api/enums.py b/graphlit_api/enums.py
@@ -172,7 +172,7 @@ class ContentTypes(str, Enum):
 
 
 class CerebrasModels(str, Enum):
-    LLAMA_3_1_70B = "LLAMA_3_1_70B"
+    LLAMA_3_3_70B = "LLAMA_3_3_70B"
     LLAMA_3_1_8B = "LLAMA_3_1_8B"
     CUSTOM = "CUSTOM"
 
@@ -606,8 +606,10 @@ class FileTypes(str, Enum):
 
 class VoyageModels(str, Enum):
     VOYAGE = "VOYAGE"
+    VOYAGE_3_0_LARGE = "VOYAGE_3_0_LARGE"
     VOYAGE_3_0 = "VOYAGE_3_0"
     VOYAGE_LITE_3_0 = "VOYAGE_LITE_3_0"
+    VOYAGE_CODE_3_0 = "VOYAGE_CODE_3_0"
     VOYAGE_FINANCE_2_0 = "VOYAGE_FINANCE_2_0"
     VOYAGE_MULTILINGUAL_2_0 = "VOYAGE_MULTILINGUAL_2_0"
     VOYAGE_LAW_2_0 = "VOYAGE_LAW_2_0"
@@ -616,17 +618,12 @@ class VoyageModels(str, Enum):
 
 
 class GroqModels(str, Enum):
-    LLAVA_1_5_7B_PREVIEW = "LLAVA_1_5_7B_PREVIEW"
     MIXTRAL_8X7B_INSTRUCT = "MIXTRAL_8X7B_INSTRUCT"
     LLAMA_3_3_70B = "LLAMA_3_3_70B"
-    LLAMA_3_2_90B_TEXT_PREVIEW = "LLAMA_3_2_90B_TEXT_PREVIEW"
     LLAMA_3_2_90B_VISION_PREVIEW = "LLAMA_3_2_90B_VISION_PREVIEW"
-    LLAMA_3_2_11B_TEXT_PREVIEW = "LLAMA_3_2_11B_TEXT_PREVIEW"
     LLAMA_3_2_11B_VISION_PREVIEW = "LLAMA_3_2_11B_VISION_PREVIEW"
     LLAMA_3_2_3B_PREVIEW = "LLAMA_3_2_3B_PREVIEW"
     LLAMA_3_2_1B_PREVIEW = "LLAMA_3_2_1B_PREVIEW"
-    LLAMA_3_1_405B = "LLAMA_3_1_405B"
-    LLAMA_3_1_70B = "LLAMA_3_1_70B"
     LLAMA_3_1_8B = "LLAMA_3_1_8B"
     LLAMA_3_70B = "LLAMA_3_70B"
     LLAMA_3_8B = "LLAMA_3_8B"

diff --git a/graphlit_api/input_types.py b/graphlit_api/input_types.py
@@ -65,6 +65,7 @@
     MetadataTypes,
     MistralModels,
     ModelServiceTypes,
+    ModelTypes,
     NotionTypes,
     ObservableTypes,
     OccurrenceTypes,
@@ -710,6 +711,13 @@ class H3IndexFilter(BaseModel):
     key: Optional[str] = None
 
 
+class ModelFilter(BaseModel):
+    types: Optional[List[Optional[ModelTypes]]] = None
+    service_types: Optional[List[Optional[ModelServiceTypes]]] = Field(
+        alias="serviceTypes", default=None
+    )
+
+
 class CollectionUpdateInput(BaseModel):
     id: str
     name: Optional[str] = None

diff --git a/graphlit_api/operations.py b/graphlit_api/operations.py
@@ -8505,13 +8505,15 @@
 """
 
 QUERY_MODELS_GQL = """
-query QueryModels {
-  models {
+query QueryModels($filter: ModelFilter) {
+  models(filter: $filter) {
     results {
+      uri
       name
       type
       serviceType
       model
+      modelType
       description
       availableOn
       features {

diff --git a/graphlit_api/query_models.py b/graphlit_api/query_models.py
@@ -18,10 +18,12 @@ class QueryModelsModels(BaseModel):
 
 
 class QueryModelsModelsResults(BaseModel):
+    uri: Optional[Any]
     name: str
     type: Optional[ModelTypes]
     service_type: Optional[ModelServiceTypes] = Field(alias="serviceType")
     model: Optional[str]
+    model_type: Optional[str] = Field(alias="modelType")
     description: Optional[str]
     available_on: Optional[List[Optional[str]]] = Field(alias="availableOn")
     features: Optional["QueryModelsModelsResultsFeatures"]