Skip to content

Commit 36c07df

Browse files
kirk-marplekirk-marple
committed
Updated to latest API release.
1 parent 7105913 commit 36c07df

11 files changed

+218
-0
lines changed

‎documents/project/QueryTokens.graphql

+20
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
query QueryTokens($startDate: DateTime!, $duration: TimeSpan!) {
2+
tokens(startDate: $startDate, duration: $duration) {
3+
correlationId
4+
ownerId
5+
embeddingInputTokens
6+
embeddingModelServices
7+
completionInputTokens
8+
completionOutputTokens
9+
completionModelServices
10+
preparationInputTokens
11+
preparationOutputTokens
12+
preparationModelServices
13+
extractionInputTokens
14+
extractionOutputTokens
15+
extractionModelServices
16+
generationInputTokens
17+
generationOutputTokens
18+
generationModelServices
19+
}
20+
}

‎documents/specification/GetSpecification.graphql

+10
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,16 @@ query GetSpecification($id: ID!, $correlationId: String) {
132132
probability
133133
chunkTokenLimit
134134
}
135+
bedrock {
136+
tokenLimit
137+
completionTokenLimit
138+
model
139+
accessKey
140+
secretAccessKey
141+
modelName
142+
temperature
143+
probability
144+
}
135145
groq {
136146
tokenLimit
137147
completionTokenLimit

‎documents/specification/QuerySpecifications.graphql

+10
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,16 @@ query QuerySpecifications($filter: SpecificationFilter, $correlationId: String)
133133
probability
134134
chunkTokenLimit
135135
}
136+
bedrock {
137+
tokenLimit
138+
completionTokenLimit
139+
model
140+
accessKey
141+
secretAccessKey
142+
modelName
143+
temperature
144+
probability
145+
}
136146
groq {
137147
tokenLimit
138148
completionTokenLimit

‎graphlit_api/__init__.py

+15
Original file line numberDiff line numberDiff line change
@@ -515,6 +515,7 @@
515515
AzureDocumentIntelligenceModels,
516516
AzureDocumentIntelligenceVersions,
517517
AzureOpenAIModels,
518+
BedrockModels,
518519
BillableMetrics,
519520
CategoryFacetTypes,
520521
CerebrasModels,
@@ -924,6 +925,7 @@
924925
GetSpecificationSpecificationAnthropic,
925926
GetSpecificationSpecificationAzureAi,
926927
GetSpecificationSpecificationAzureOpenAi,
928+
GetSpecificationSpecificationBedrock,
927929
GetSpecificationSpecificationCerebras,
928930
GetSpecificationSpecificationCohere,
929931
GetSpecificationSpecificationDeepseek,
@@ -1098,6 +1100,8 @@
10981100
AzureOpenAIModelPropertiesInput,
10991101
AzureOpenAIModelPropertiesUpdateInput,
11001102
AzureTextExtractionPropertiesInput,
1103+
BedrockModelPropertiesInput,
1104+
BedrockModelPropertiesUpdateInput,
11011105
BoundingBoxInput,
11021106
BoxFeedPropertiesInput,
11031107
BoxFeedPropertiesUpdateInput,
@@ -1644,6 +1648,7 @@
16441648
QUERY_SLACK_CHANNELS_GQL,
16451649
QUERY_SOFTWARES_GQL,
16461650
QUERY_SPECIFICATIONS_GQL,
1651+
QUERY_TOKENS_GQL,
16471652
QUERY_USAGE_GQL,
16481653
QUERY_USERS_GQL,
16491654
QUERY_WORKFLOWS_GQL,
@@ -2129,6 +2134,7 @@
21292134
QuerySpecificationsSpecificationsResultsAnthropic,
21302135
QuerySpecificationsSpecificationsResultsAzureAi,
21312136
QuerySpecificationsSpecificationsResultsAzureOpenAi,
2137+
QuerySpecificationsSpecificationsResultsBedrock,
21322138
QuerySpecificationsSpecificationsResultsCerebras,
21332139
QuerySpecificationsSpecificationsResultsCohere,
21342140
QuerySpecificationsSpecificationsResultsDeepseek,
@@ -2147,6 +2153,7 @@
21472153
QuerySpecificationsSpecificationsResultsStrategy,
21482154
QuerySpecificationsSpecificationsResultsVoyage,
21492155
)
2156+
from .query_tokens import QueryTokens, QueryTokensTokens
21502157
from .query_usage import QueryUsage, QueryUsageUsage
21512158
from .query_users import (
21522159
QueryUsers,
@@ -2561,6 +2568,9 @@
25612568
"AzureOpenAIModels",
25622569
"AzureTextExtractionPropertiesInput",
25632570
"BaseModel",
2571+
"BedrockModelPropertiesInput",
2572+
"BedrockModelPropertiesUpdateInput",
2573+
"BedrockModels",
25642574
"BillableMetrics",
25652575
"BoundingBoxInput",
25662576
"BoxFeedPropertiesInput",
@@ -3549,6 +3559,7 @@
35493559
"GetSpecificationSpecificationAnthropic",
35503560
"GetSpecificationSpecificationAzureAi",
35513561
"GetSpecificationSpecificationAzureOpenAi",
3562+
"GetSpecificationSpecificationBedrock",
35523563
"GetSpecificationSpecificationCerebras",
35533564
"GetSpecificationSpecificationCohere",
35543565
"GetSpecificationSpecificationDeepseek",
@@ -4018,6 +4029,7 @@
40184029
"QUERY_SLACK_CHANNELS_GQL",
40194030
"QUERY_SOFTWARES_GQL",
40204031
"QUERY_SPECIFICATIONS_GQL",
4032+
"QUERY_TOKENS_GQL",
40214033
"QUERY_USAGE_GQL",
40224034
"QUERY_USERS_GQL",
40234035
"QUERY_WORKFLOWS_GQL",
@@ -4312,6 +4324,7 @@
43124324
"QuerySpecificationsSpecificationsResultsAnthropic",
43134325
"QuerySpecificationsSpecificationsResultsAzureAi",
43144326
"QuerySpecificationsSpecificationsResultsAzureOpenAi",
4327+
"QuerySpecificationsSpecificationsResultsBedrock",
43154328
"QuerySpecificationsSpecificationsResultsCerebras",
43164329
"QuerySpecificationsSpecificationsResultsCohere",
43174330
"QuerySpecificationsSpecificationsResultsDeepseek",
@@ -4329,6 +4342,8 @@
43294342
"QuerySpecificationsSpecificationsResultsRevisionStrategy",
43304343
"QuerySpecificationsSpecificationsResultsStrategy",
43314344
"QuerySpecificationsSpecificationsResultsVoyage",
4345+
"QueryTokens",
4346+
"QueryTokensTokens",
43324347
"QueryUsage",
43334348
"QueryUsageUsage",
43344349
"QueryUsers",

‎graphlit_api/client.py

+15
Original file line numberDiff line numberDiff line change
@@ -559,6 +559,7 @@
559559
QUERY_SLACK_CHANNELS_GQL,
560560
QUERY_SOFTWARES_GQL,
561561
QUERY_SPECIFICATIONS_GQL,
562+
QUERY_TOKENS_GQL,
562563
QUERY_USAGE_GQL,
563564
QUERY_USERS_GQL,
564565
QUERY_WORKFLOWS_GQL,
@@ -656,6 +657,7 @@
656657
from .query_slack_channels import QuerySlackChannels
657658
from .query_softwares import QuerySoftwares
658659
from .query_specifications import QuerySpecifications
660+
from .query_tokens import QueryTokens
659661
from .query_usage import QueryUsage
660662
from .query_users import QueryUsers
661663
from .query_workflows import QueryWorkflows
@@ -4857,6 +4859,19 @@ async def query_credits(
48574859
data = self.get_data(response)
48584860
return QueryCredits.model_validate(data)
48594861

4862+
async def query_tokens(
4863+
self, start_date: Any, duration: Any, **kwargs: Any
4864+
) -> QueryTokens:
4865+
variables: Dict[str, object] = {"startDate": start_date, "duration": duration}
4866+
response = await self.execute(
4867+
query=QUERY_TOKENS_GQL,
4868+
operation_name="QueryTokens",
4869+
variables=variables,
4870+
**kwargs
4871+
)
4872+
data = self.get_data(response)
4873+
return QueryTokens.model_validate(data)
4874+
48604875
async def query_usage(
48614876
self,
48624877
start_date: Any,

‎graphlit_api/enums.py

+5
Original file line numberDiff line numberDiff line change
@@ -246,6 +246,7 @@ class ModelServiceTypes(str, Enum):
246246
DEEPSEEK = "DEEPSEEK"
247247
JINA = "JINA"
248248
VOYAGE = "VOYAGE"
249+
BEDROCK = "BEDROCK"
249250

250251

251252
class MedicalStudyFacetTypes(str, Enum):
@@ -537,6 +538,10 @@ class OccurrenceTypes(str, Enum):
537538
TEXT = "TEXT"
538539

539540

541+
class BedrockModels(str, Enum):
542+
CUSTOM = "CUSTOM"
543+
544+
540545
class ContentSourceTypes(str, Enum):
541546
FRAME = "FRAME"
542547
DOCUMENT = "DOCUMENT"

‎graphlit_api/get_specification.py

+13
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from .enums import (
1010
AnthropicModels,
1111
AzureOpenAIModels,
12+
BedrockModels,
1213
CerebrasModels,
1314
CohereModels,
1415
ConversationSearchTypes,
@@ -78,6 +79,7 @@ class GetSpecificationSpecification(BaseModel):
7879
google: Optional["GetSpecificationSpecificationGoogle"]
7980
replicate: Optional["GetSpecificationSpecificationReplicate"]
8081
mistral: Optional["GetSpecificationSpecificationMistral"]
82+
bedrock: Optional["GetSpecificationSpecificationBedrock"]
8183
groq: Optional["GetSpecificationSpecificationGroq"]
8284
cerebras: Optional["GetSpecificationSpecificationCerebras"]
8385
deepseek: Optional["GetSpecificationSpecificationDeepseek"]
@@ -222,6 +224,17 @@ class GetSpecificationSpecificationMistral(BaseModel):
222224
chunk_token_limit: Optional[int] = Field(alias="chunkTokenLimit")
223225

224226

227+
class GetSpecificationSpecificationBedrock(BaseModel):
228+
token_limit: Optional[int] = Field(alias="tokenLimit")
229+
completion_token_limit: Optional[int] = Field(alias="completionTokenLimit")
230+
model: BedrockModels
231+
access_key: Optional[str] = Field(alias="accessKey")
232+
secret_access_key: Optional[str] = Field(alias="secretAccessKey")
233+
model_name: Optional[str] = Field(alias="modelName")
234+
temperature: Optional[float]
235+
probability: Optional[float]
236+
237+
225238
class GetSpecificationSpecificationGroq(BaseModel):
226239
token_limit: Optional[int] = Field(alias="tokenLimit")
227240
completion_token_limit: Optional[int] = Field(alias="completionTokenLimit")

‎graphlit_api/input_types.py

+29
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
AzureDocumentIntelligenceModels,
1515
AzureDocumentIntelligenceVersions,
1616
AzureOpenAIModels,
17+
BedrockModels,
1718
CategoryFacetTypes,
1819
CerebrasModels,
1920
CohereModels,
@@ -852,6 +853,7 @@ class SpecificationInput(BaseModel):
852853
google: Optional["GoogleModelPropertiesInput"] = None
853854
replicate: Optional["ReplicateModelPropertiesInput"] = None
854855
mistral: Optional["MistralModelPropertiesInput"] = None
856+
bedrock: Optional["BedrockModelPropertiesInput"] = None
855857
groq: Optional["GroqModelPropertiesInput"] = None
856858
cerebras: Optional["CerebrasModelPropertiesInput"] = None
857859
deepseek: Optional["DeepseekModelPropertiesInput"] = None
@@ -1028,6 +1030,19 @@ class IntegrationConnectorUpdateInput(BaseModel):
10281030
twitter: Optional["TwitterIntegrationPropertiesInput"] = None
10291031

10301032

1033+
class BedrockModelPropertiesInput(BaseModel):
1034+
model: BedrockModels
1035+
model_name: Optional[str] = Field(alias="modelName", default=None)
1036+
access_key: Optional[str] = Field(alias="accessKey", default=None)
1037+
secret_access_key: Optional[str] = Field(alias="secretAccessKey", default=None)
1038+
temperature: Optional[float] = None
1039+
probability: Optional[float] = None
1040+
token_limit: Optional[int] = Field(alias="tokenLimit", default=None)
1041+
completion_token_limit: Optional[int] = Field(
1042+
alias="completionTokenLimit", default=None
1043+
)
1044+
1045+
10311046
class MedicalDeviceFilter(BaseModel):
10321047
search: Optional[str] = None
10331048
order_by: Optional[OrderByTypes] = Field(alias="orderBy", default=None)
@@ -1095,6 +1110,19 @@ class AzureFileFeedPropertiesUpdateInput(BaseModel):
10951110
prefix: Optional[str] = None
10961111

10971112

1113+
class BedrockModelPropertiesUpdateInput(BaseModel):
1114+
model: Optional[BedrockModels] = None
1115+
model_name: Optional[str] = Field(alias="modelName", default=None)
1116+
access_key: Optional[str] = Field(alias="accessKey", default=None)
1117+
secret_access_key: Optional[str] = Field(alias="secretAccessKey", default=None)
1118+
temperature: Optional[float] = None
1119+
probability: Optional[float] = None
1120+
token_limit: Optional[int] = Field(alias="tokenLimit", default=None)
1121+
completion_token_limit: Optional[int] = Field(
1122+
alias="completionTokenLimit", default=None
1123+
)
1124+
1125+
10981126
class PromptClassificationRuleInput(BaseModel):
10991127
if_: Optional[str] = Field(alias="if", default=None)
11001128
then: Optional[str] = None
@@ -1487,6 +1515,7 @@ class SpecificationUpdateInput(BaseModel):
14871515
google: Optional["GoogleModelPropertiesUpdateInput"] = None
14881516
replicate: Optional["ReplicateModelPropertiesUpdateInput"] = None
14891517
mistral: Optional["MistralModelPropertiesUpdateInput"] = None
1518+
bedrock: Optional["BedrockModelPropertiesUpdateInput"] = None
14901519
groq: Optional["GroqModelPropertiesUpdateInput"] = None
14911520
cerebras: Optional["CerebrasModelPropertiesUpdateInput"] = None
14921521
deepseek: Optional["DeepseekModelPropertiesUpdateInput"] = None

‎graphlit_api/operations.py

+44
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,7 @@
245245
"QUERY_SLACK_CHANNELS_GQL",
246246
"QUERY_SOFTWARES_GQL",
247247
"QUERY_SPECIFICATIONS_GQL",
248+
"QUERY_TOKENS_GQL",
248249
"QUERY_USAGE_GQL",
249250
"QUERY_USERS_GQL",
250251
"QUERY_WORKFLOWS_GQL",
@@ -8015,6 +8016,29 @@
80158016
}
80168017
"""
80178018

8019+
QUERY_TOKENS_GQL = """
8020+
query QueryTokens($startDate: DateTime!, $duration: TimeSpan!) {
8021+
tokens(startDate: $startDate, duration: $duration) {
8022+
correlationId
8023+
ownerId
8024+
embeddingInputTokens
8025+
embeddingModelServices
8026+
completionInputTokens
8027+
completionOutputTokens
8028+
completionModelServices
8029+
preparationInputTokens
8030+
preparationOutputTokens
8031+
preparationModelServices
8032+
extractionInputTokens
8033+
extractionOutputTokens
8034+
extractionModelServices
8035+
generationInputTokens
8036+
generationOutputTokens
8037+
generationModelServices
8038+
}
8039+
}
8040+
"""
8041+
80188042
QUERY_USAGE_GQL = """
80198043
query QueryUsage($startDate: DateTime!, $duration: TimeSpan!, $names: [String!], $offset: Int, $limit: Int) {
80208044
usage(
@@ -8470,6 +8494,16 @@
84708494
probability
84718495
chunkTokenLimit
84728496
}
8497+
bedrock {
8498+
tokenLimit
8499+
completionTokenLimit
8500+
model
8501+
accessKey
8502+
secretAccessKey
8503+
modelName
8504+
temperature
8505+
probability
8506+
}
84738507
groq {
84748508
tokenLimit
84758509
completionTokenLimit
@@ -8818,6 +8852,16 @@
88188852
probability
88198853
chunkTokenLimit
88208854
}
8855+
bedrock {
8856+
tokenLimit
8857+
completionTokenLimit
8858+
model
8859+
accessKey
8860+
secretAccessKey
8861+
modelName
8862+
temperature
8863+
probability
8864+
}
88218865
groq {
88228866
tokenLimit
88238867
completionTokenLimit

0 commit comments

Comments
 (0)