athena-intelligence 0.1.124__py3-none-any.whl → 0.1.126__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (58) hide show
  1. athena/__init__.py +3 -0
  2. athena/agents/client.py +88 -36
  3. athena/agents/drive/client.py +80 -32
  4. athena/agents/general/client.py +222 -91
  5. athena/agents/research/client.py +80 -32
  6. athena/agents/sql/client.py +80 -32
  7. athena/base_client.py +13 -11
  8. athena/client.py +191 -80
  9. athena/core/__init__.py +21 -4
  10. athena/core/client_wrapper.py +9 -10
  11. athena/core/file.py +37 -8
  12. athena/core/http_client.py +97 -41
  13. athena/core/jsonable_encoder.py +33 -31
  14. athena/core/pydantic_utilities.py +272 -4
  15. athena/core/query_encoder.py +38 -13
  16. athena/core/request_options.py +5 -2
  17. athena/core/serialization.py +272 -0
  18. athena/errors/internal_server_error.py +2 -3
  19. athena/errors/unauthorized_error.py +2 -3
  20. athena/errors/unprocessable_entity_error.py +2 -3
  21. athena/query/client.py +208 -58
  22. athena/tools/calendar/client.py +82 -30
  23. athena/tools/client.py +576 -184
  24. athena/tools/email/client.py +117 -43
  25. athena/tools/structured_data_extractor/client.py +118 -67
  26. athena/tools/tasks/client.py +41 -17
  27. athena/types/asset_node.py +14 -24
  28. athena/types/asset_not_found_error.py +11 -21
  29. athena/types/chunk.py +11 -21
  30. athena/types/chunk_content_item.py +21 -41
  31. athena/types/chunk_result.py +13 -23
  32. athena/types/custom_agent_response.py +12 -22
  33. athena/types/data_frame_request_out.py +11 -21
  34. athena/types/data_frame_unknown_format_error.py +11 -21
  35. athena/types/document_chunk.py +12 -22
  36. athena/types/drive_agent_response.py +12 -22
  37. athena/types/file_chunk_request_out.py +11 -21
  38. athena/types/file_too_large_error.py +11 -21
  39. athena/types/folder_response.py +11 -21
  40. athena/types/general_agent_config.py +11 -21
  41. athena/types/general_agent_config_enabled_tools_item.py +0 -1
  42. athena/types/general_agent_request.py +13 -23
  43. athena/types/general_agent_response.py +12 -22
  44. athena/types/image_url_content.py +11 -21
  45. athena/types/parent_folder_error.py +11 -21
  46. athena/types/prompt_message.py +12 -22
  47. athena/types/research_agent_response.py +12 -22
  48. athena/types/save_asset_request_out.py +11 -21
  49. athena/types/sql_agent_response.py +13 -23
  50. athena/types/structured_data_extractor_response.py +15 -25
  51. athena/types/text_content.py +11 -21
  52. athena/types/tool.py +1 -13
  53. athena/types/type.py +1 -21
  54. athena/version.py +0 -1
  55. {athena_intelligence-0.1.124.dist-info → athena_intelligence-0.1.126.dist-info}/METADATA +12 -4
  56. athena_intelligence-0.1.126.dist-info/RECORD +87 -0
  57. {athena_intelligence-0.1.124.dist-info → athena_intelligence-0.1.126.dist-info}/WHEEL +1 -1
  58. athena_intelligence-0.1.124.dist-info/RECORD +0 -86
@@ -1,15 +1,16 @@
1
1
  # This file was auto-generated by Fern from our API Definition.
2
2
 
3
3
  import typing
4
- from json.decoder import JSONDecodeError
5
-
6
- from ...core.api_error import ApiError
7
- from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
8
- from ...core.pydantic_utilities import pydantic_v1
9
- from ...core.request_options import RequestOptions
10
- from ...errors.unprocessable_entity_error import UnprocessableEntityError
4
+ from ...core.client_wrapper import SyncClientWrapper
11
5
  from ...types.general_agent_request import GeneralAgentRequest
6
+ from ...core.request_options import RequestOptions
12
7
  from ...types.general_agent_response import GeneralAgentResponse
8
+ from ...core.serialization import convert_and_respect_annotation_metadata
9
+ from ...core.pydantic_utilities import parse_obj_as
10
+ from ...errors.unprocessable_entity_error import UnprocessableEntityError
11
+ from json.decoder import JSONDecodeError
12
+ from ...core.api_error import ApiError
13
+ from ...core.client_wrapper import AsyncClientWrapper
13
14
 
14
15
  # this is used as the default value for optional parameters
15
16
  OMIT = typing.cast(typing.Any, ...)
@@ -39,8 +40,7 @@ class GeneralClient:
39
40
 
40
41
  Examples
41
42
  --------
42
- from athena import GeneralAgentConfig, GeneralAgentRequest, Tool
43
- from athena.client import Athena
43
+ from athena import Athena, GeneralAgentConfig, GeneralAgentRequest
44
44
 
45
45
  client = Athena(
46
46
  api_key="YOUR_API_KEY",
@@ -49,7 +49,7 @@ class GeneralClient:
49
49
  request=[
50
50
  GeneralAgentRequest(
51
51
  config=GeneralAgentConfig(
52
- enabled_tools=[Tool.SEARCH],
52
+ enabled_tools=["search"],
53
53
  ),
54
54
  messages=[
55
55
  {
@@ -62,13 +62,33 @@ class GeneralClient:
62
62
  )
63
63
  """
64
64
  _response = self._client_wrapper.httpx_client.request(
65
- "api/v0/agents/general/batch", method="POST", json=request, request_options=request_options, omit=OMIT
65
+ "api/v0/agents/general/batch",
66
+ method="POST",
67
+ json=convert_and_respect_annotation_metadata(
68
+ object_=request, annotation=typing.Sequence[GeneralAgentRequest], direction="write"
69
+ ),
70
+ request_options=request_options,
71
+ omit=OMIT,
66
72
  )
67
- if 200 <= _response.status_code < 300:
68
- return pydantic_v1.parse_obj_as(typing.List[GeneralAgentResponse], _response.json()) # type: ignore
69
- if _response.status_code == 422:
70
- raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json())) # type: ignore
71
73
  try:
74
+ if 200 <= _response.status_code < 300:
75
+ return typing.cast(
76
+ typing.List[GeneralAgentResponse],
77
+ parse_obj_as(
78
+ type_=typing.List[GeneralAgentResponse], # type: ignore
79
+ object_=_response.json(),
80
+ ),
81
+ )
82
+ if _response.status_code == 422:
83
+ raise UnprocessableEntityError(
84
+ typing.cast(
85
+ typing.Optional[typing.Any],
86
+ parse_obj_as(
87
+ type_=typing.Optional[typing.Any], # type: ignore
88
+ object_=_response.json(),
89
+ ),
90
+ )
91
+ )
72
92
  _response_json = _response.json()
73
93
  except JSONDecodeError:
74
94
  raise ApiError(status_code=_response.status_code, body=_response.text)
@@ -97,8 +117,7 @@ class GeneralClient:
97
117
 
98
118
  Examples
99
119
  --------
100
- from athena import GeneralAgentConfig, GeneralAgentRequest, Tool
101
- from athena.client import Athena
120
+ from athena import Athena, GeneralAgentConfig, GeneralAgentRequest
102
121
 
103
122
  client = Athena(
104
123
  api_key="YOUR_API_KEY",
@@ -106,7 +125,7 @@ class GeneralClient:
106
125
  client.agents.general.invoke(
107
126
  request=GeneralAgentRequest(
108
127
  config=GeneralAgentConfig(
109
- enabled_tools=[Tool.SEARCH],
128
+ enabled_tools=["search"],
110
129
  ),
111
130
  messages=[
112
131
  {
@@ -118,13 +137,33 @@ class GeneralClient:
118
137
  )
119
138
  """
120
139
  _response = self._client_wrapper.httpx_client.request(
121
- "api/v0/agents/general/invoke", method="POST", json=request, request_options=request_options, omit=OMIT
140
+ "api/v0/agents/general/invoke",
141
+ method="POST",
142
+ json=convert_and_respect_annotation_metadata(
143
+ object_=request, annotation=GeneralAgentRequest, direction="write"
144
+ ),
145
+ request_options=request_options,
146
+ omit=OMIT,
122
147
  )
123
- if 200 <= _response.status_code < 300:
124
- return pydantic_v1.parse_obj_as(GeneralAgentResponse, _response.json()) # type: ignore
125
- if _response.status_code == 422:
126
- raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json())) # type: ignore
127
148
  try:
149
+ if 200 <= _response.status_code < 300:
150
+ return typing.cast(
151
+ GeneralAgentResponse,
152
+ parse_obj_as(
153
+ type_=GeneralAgentResponse, # type: ignore
154
+ object_=_response.json(),
155
+ ),
156
+ )
157
+ if _response.status_code == 422:
158
+ raise UnprocessableEntityError(
159
+ typing.cast(
160
+ typing.Optional[typing.Any],
161
+ parse_obj_as(
162
+ type_=typing.Optional[typing.Any], # type: ignore
163
+ object_=_response.json(),
164
+ ),
165
+ )
166
+ )
128
167
  _response_json = _response.json()
129
168
  except JSONDecodeError:
130
169
  raise ApiError(status_code=_response.status_code, body=_response.text)
@@ -150,8 +189,7 @@ class GeneralClient:
150
189
 
151
190
  Examples
152
191
  --------
153
- from athena import GeneralAgentConfig, GeneralAgentRequest, Tool
154
- from athena.client import Athena
192
+ from athena import Athena, GeneralAgentConfig, GeneralAgentRequest
155
193
 
156
194
  client = Athena(
157
195
  api_key="YOUR_API_KEY",
@@ -159,7 +197,7 @@ class GeneralClient:
159
197
  client.agents.general.stream_events(
160
198
  request=GeneralAgentRequest(
161
199
  config=GeneralAgentConfig(
162
- enabled_tools=[Tool.SEARCH],
200
+ enabled_tools=["search"],
163
201
  ),
164
202
  messages=[
165
203
  {
@@ -173,15 +211,31 @@ class GeneralClient:
173
211
  _response = self._client_wrapper.httpx_client.request(
174
212
  "api/v0/agents/general/stream_events",
175
213
  method="POST",
176
- json=request,
214
+ json=convert_and_respect_annotation_metadata(
215
+ object_=request, annotation=GeneralAgentRequest, direction="write"
216
+ ),
177
217
  request_options=request_options,
178
218
  omit=OMIT,
179
219
  )
180
- if 200 <= _response.status_code < 300:
181
- return pydantic_v1.parse_obj_as(GeneralAgentResponse, _response.json()) # type: ignore
182
- if _response.status_code == 422:
183
- raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json())) # type: ignore
184
220
  try:
221
+ if 200 <= _response.status_code < 300:
222
+ return typing.cast(
223
+ GeneralAgentResponse,
224
+ parse_obj_as(
225
+ type_=GeneralAgentResponse, # type: ignore
226
+ object_=_response.json(),
227
+ ),
228
+ )
229
+ if _response.status_code == 422:
230
+ raise UnprocessableEntityError(
231
+ typing.cast(
232
+ typing.Optional[typing.Any],
233
+ parse_obj_as(
234
+ type_=typing.Optional[typing.Any], # type: ignore
235
+ object_=_response.json(),
236
+ ),
237
+ )
238
+ )
185
239
  _response_json = _response.json()
186
240
  except JSONDecodeError:
187
241
  raise ApiError(status_code=_response.status_code, body=_response.text)
@@ -212,36 +266,63 @@ class AsyncGeneralClient:
212
266
 
213
267
  Examples
214
268
  --------
215
- from athena import GeneralAgentConfig, GeneralAgentRequest, Tool
216
- from athena.client import AsyncAthena
269
+ import asyncio
270
+
271
+ from athena import AsyncAthena, GeneralAgentConfig, GeneralAgentRequest
217
272
 
218
273
  client = AsyncAthena(
219
274
  api_key="YOUR_API_KEY",
220
275
  )
221
- await client.agents.general.batch(
222
- request=[
223
- GeneralAgentRequest(
224
- config=GeneralAgentConfig(
225
- enabled_tools=[Tool.SEARCH],
226
- ),
227
- messages=[
228
- {
229
- "content": "Please call the search tool for AAPL news.",
230
- "type": "user",
231
- }
232
- ],
233
- )
234
- ],
235
- )
276
+
277
+
278
+ async def main() -> None:
279
+ await client.agents.general.batch(
280
+ request=[
281
+ GeneralAgentRequest(
282
+ config=GeneralAgentConfig(
283
+ enabled_tools=["search"],
284
+ ),
285
+ messages=[
286
+ {
287
+ "content": "Please call the search tool for AAPL news.",
288
+ "type": "user",
289
+ }
290
+ ],
291
+ )
292
+ ],
293
+ )
294
+
295
+
296
+ asyncio.run(main())
236
297
  """
237
298
  _response = await self._client_wrapper.httpx_client.request(
238
- "api/v0/agents/general/batch", method="POST", json=request, request_options=request_options, omit=OMIT
299
+ "api/v0/agents/general/batch",
300
+ method="POST",
301
+ json=convert_and_respect_annotation_metadata(
302
+ object_=request, annotation=typing.Sequence[GeneralAgentRequest], direction="write"
303
+ ),
304
+ request_options=request_options,
305
+ omit=OMIT,
239
306
  )
240
- if 200 <= _response.status_code < 300:
241
- return pydantic_v1.parse_obj_as(typing.List[GeneralAgentResponse], _response.json()) # type: ignore
242
- if _response.status_code == 422:
243
- raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json())) # type: ignore
244
307
  try:
308
+ if 200 <= _response.status_code < 300:
309
+ return typing.cast(
310
+ typing.List[GeneralAgentResponse],
311
+ parse_obj_as(
312
+ type_=typing.List[GeneralAgentResponse], # type: ignore
313
+ object_=_response.json(),
314
+ ),
315
+ )
316
+ if _response.status_code == 422:
317
+ raise UnprocessableEntityError(
318
+ typing.cast(
319
+ typing.Optional[typing.Any],
320
+ parse_obj_as(
321
+ type_=typing.Optional[typing.Any], # type: ignore
322
+ object_=_response.json(),
323
+ ),
324
+ )
325
+ )
245
326
  _response_json = _response.json()
246
327
  except JSONDecodeError:
247
328
  raise ApiError(status_code=_response.status_code, body=_response.text)
@@ -270,34 +351,61 @@ class AsyncGeneralClient:
270
351
 
271
352
  Examples
272
353
  --------
273
- from athena import GeneralAgentConfig, GeneralAgentRequest, Tool
274
- from athena.client import AsyncAthena
354
+ import asyncio
355
+
356
+ from athena import AsyncAthena, GeneralAgentConfig, GeneralAgentRequest
275
357
 
276
358
  client = AsyncAthena(
277
359
  api_key="YOUR_API_KEY",
278
360
  )
279
- await client.agents.general.invoke(
280
- request=GeneralAgentRequest(
281
- config=GeneralAgentConfig(
282
- enabled_tools=[Tool.SEARCH],
361
+
362
+
363
+ async def main() -> None:
364
+ await client.agents.general.invoke(
365
+ request=GeneralAgentRequest(
366
+ config=GeneralAgentConfig(
367
+ enabled_tools=["search"],
368
+ ),
369
+ messages=[
370
+ {
371
+ "content": "Please call the search tool for AAPL news.",
372
+ "type": "user",
373
+ }
374
+ ],
283
375
  ),
284
- messages=[
285
- {
286
- "content": "Please call the search tool for AAPL news.",
287
- "type": "user",
288
- }
289
- ],
290
- ),
291
- )
376
+ )
377
+
378
+
379
+ asyncio.run(main())
292
380
  """
293
381
  _response = await self._client_wrapper.httpx_client.request(
294
- "api/v0/agents/general/invoke", method="POST", json=request, request_options=request_options, omit=OMIT
382
+ "api/v0/agents/general/invoke",
383
+ method="POST",
384
+ json=convert_and_respect_annotation_metadata(
385
+ object_=request, annotation=GeneralAgentRequest, direction="write"
386
+ ),
387
+ request_options=request_options,
388
+ omit=OMIT,
295
389
  )
296
- if 200 <= _response.status_code < 300:
297
- return pydantic_v1.parse_obj_as(GeneralAgentResponse, _response.json()) # type: ignore
298
- if _response.status_code == 422:
299
- raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json())) # type: ignore
300
390
  try:
391
+ if 200 <= _response.status_code < 300:
392
+ return typing.cast(
393
+ GeneralAgentResponse,
394
+ parse_obj_as(
395
+ type_=GeneralAgentResponse, # type: ignore
396
+ object_=_response.json(),
397
+ ),
398
+ )
399
+ if _response.status_code == 422:
400
+ raise UnprocessableEntityError(
401
+ typing.cast(
402
+ typing.Optional[typing.Any],
403
+ parse_obj_as(
404
+ type_=typing.Optional[typing.Any], # type: ignore
405
+ object_=_response.json(),
406
+ ),
407
+ )
408
+ )
301
409
  _response_json = _response.json()
302
410
  except JSONDecodeError:
303
411
  raise ApiError(status_code=_response.status_code, body=_response.text)
@@ -323,38 +431,61 @@ class AsyncGeneralClient:
323
431
 
324
432
  Examples
325
433
  --------
326
- from athena import GeneralAgentConfig, GeneralAgentRequest, Tool
327
- from athena.client import AsyncAthena
434
+ import asyncio
435
+
436
+ from athena import AsyncAthena, GeneralAgentConfig, GeneralAgentRequest
328
437
 
329
438
  client = AsyncAthena(
330
439
  api_key="YOUR_API_KEY",
331
440
  )
332
- await client.agents.general.stream_events(
333
- request=GeneralAgentRequest(
334
- config=GeneralAgentConfig(
335
- enabled_tools=[Tool.SEARCH],
441
+
442
+
443
+ async def main() -> None:
444
+ await client.agents.general.stream_events(
445
+ request=GeneralAgentRequest(
446
+ config=GeneralAgentConfig(
447
+ enabled_tools=["search"],
448
+ ),
449
+ messages=[
450
+ {
451
+ "content": "Please call the search tool for AAPL news.",
452
+ "type": "user",
453
+ }
454
+ ],
336
455
  ),
337
- messages=[
338
- {
339
- "content": "Please call the search tool for AAPL news.",
340
- "type": "user",
341
- }
342
- ],
343
- ),
344
- )
456
+ )
457
+
458
+
459
+ asyncio.run(main())
345
460
  """
346
461
  _response = await self._client_wrapper.httpx_client.request(
347
462
  "api/v0/agents/general/stream_events",
348
463
  method="POST",
349
- json=request,
464
+ json=convert_and_respect_annotation_metadata(
465
+ object_=request, annotation=GeneralAgentRequest, direction="write"
466
+ ),
350
467
  request_options=request_options,
351
468
  omit=OMIT,
352
469
  )
353
- if 200 <= _response.status_code < 300:
354
- return pydantic_v1.parse_obj_as(GeneralAgentResponse, _response.json()) # type: ignore
355
- if _response.status_code == 422:
356
- raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json())) # type: ignore
357
470
  try:
471
+ if 200 <= _response.status_code < 300:
472
+ return typing.cast(
473
+ GeneralAgentResponse,
474
+ parse_obj_as(
475
+ type_=GeneralAgentResponse, # type: ignore
476
+ object_=_response.json(),
477
+ ),
478
+ )
479
+ if _response.status_code == 422:
480
+ raise UnprocessableEntityError(
481
+ typing.cast(
482
+ typing.Optional[typing.Any],
483
+ parse_obj_as(
484
+ type_=typing.Optional[typing.Any], # type: ignore
485
+ object_=_response.json(),
486
+ ),
487
+ )
488
+ )
358
489
  _response_json = _response.json()
359
490
  except JSONDecodeError:
360
491
  raise ApiError(status_code=_response.status_code, body=_response.text)
@@ -1,14 +1,14 @@
1
1
  # This file was auto-generated by Fern from our API Definition.
2
2
 
3
3
  import typing
4
- from json.decoder import JSONDecodeError
5
-
6
- from ...core.api_error import ApiError
7
- from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
8
- from ...core.pydantic_utilities import pydantic_v1
4
+ from ...core.client_wrapper import SyncClientWrapper
9
5
  from ...core.request_options import RequestOptions
10
- from ...errors.unprocessable_entity_error import UnprocessableEntityError
11
6
  from ...types.research_agent_response import ResearchAgentResponse
7
+ from ...core.pydantic_utilities import parse_obj_as
8
+ from ...errors.unprocessable_entity_error import UnprocessableEntityError
9
+ from json.decoder import JSONDecodeError
10
+ from ...core.api_error import ApiError
11
+ from ...core.client_wrapper import AsyncClientWrapper
12
12
 
13
13
  # this is used as the default value for optional parameters
14
14
  OMIT = typing.cast(typing.Any, ...)
@@ -21,19 +21,19 @@ class ResearchClient:
21
21
  def invoke(
22
22
  self,
23
23
  *,
24
- config: typing.Dict[str, typing.Any],
25
- messages: typing.Sequence[typing.Dict[str, typing.Any]],
26
- request_options: typing.Optional[RequestOptions] = None
24
+ config: typing.Dict[str, typing.Optional[typing.Any]],
25
+ messages: typing.Sequence[typing.Dict[str, typing.Optional[typing.Any]]],
26
+ request_options: typing.Optional[RequestOptions] = None,
27
27
  ) -> ResearchAgentResponse:
28
28
  """
29
29
  Coming soon! Conduct research using web and other sources.
30
30
 
31
31
  Parameters
32
32
  ----------
33
- config : typing.Dict[str, typing.Any]
33
+ config : typing.Dict[str, typing.Optional[typing.Any]]
34
34
  Configuration for the research agent including search parameters and sources
35
35
 
36
- messages : typing.Sequence[typing.Dict[str, typing.Any]]
36
+ messages : typing.Sequence[typing.Dict[str, typing.Optional[typing.Any]]]
37
37
  The messages to send to the research agent
38
38
 
39
39
  request_options : typing.Optional[RequestOptions]
@@ -46,7 +46,7 @@ class ResearchClient:
46
46
 
47
47
  Examples
48
48
  --------
49
- from athena.client import Athena
49
+ from athena import Athena
50
50
 
51
51
  client = Athena(
52
52
  api_key="YOUR_API_KEY",
@@ -59,15 +59,35 @@ class ResearchClient:
59
59
  _response = self._client_wrapper.httpx_client.request(
60
60
  "api/v0/agents/research/invoke",
61
61
  method="POST",
62
- json={"config": config, "messages": messages},
62
+ json={
63
+ "config": config,
64
+ "messages": messages,
65
+ },
66
+ headers={
67
+ "content-type": "application/json",
68
+ },
63
69
  request_options=request_options,
64
70
  omit=OMIT,
65
71
  )
66
- if 200 <= _response.status_code < 300:
67
- return pydantic_v1.parse_obj_as(ResearchAgentResponse, _response.json()) # type: ignore
68
- if _response.status_code == 422:
69
- raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json())) # type: ignore
70
72
  try:
73
+ if 200 <= _response.status_code < 300:
74
+ return typing.cast(
75
+ ResearchAgentResponse,
76
+ parse_obj_as(
77
+ type_=ResearchAgentResponse, # type: ignore
78
+ object_=_response.json(),
79
+ ),
80
+ )
81
+ if _response.status_code == 422:
82
+ raise UnprocessableEntityError(
83
+ typing.cast(
84
+ typing.Optional[typing.Any],
85
+ parse_obj_as(
86
+ type_=typing.Optional[typing.Any], # type: ignore
87
+ object_=_response.json(),
88
+ ),
89
+ )
90
+ )
71
91
  _response_json = _response.json()
72
92
  except JSONDecodeError:
73
93
  raise ApiError(status_code=_response.status_code, body=_response.text)
@@ -81,19 +101,19 @@ class AsyncResearchClient:
81
101
  async def invoke(
82
102
  self,
83
103
  *,
84
- config: typing.Dict[str, typing.Any],
85
- messages: typing.Sequence[typing.Dict[str, typing.Any]],
86
- request_options: typing.Optional[RequestOptions] = None
104
+ config: typing.Dict[str, typing.Optional[typing.Any]],
105
+ messages: typing.Sequence[typing.Dict[str, typing.Optional[typing.Any]]],
106
+ request_options: typing.Optional[RequestOptions] = None,
87
107
  ) -> ResearchAgentResponse:
88
108
  """
89
109
  Coming soon! Conduct research using web and other sources.
90
110
 
91
111
  Parameters
92
112
  ----------
93
- config : typing.Dict[str, typing.Any]
113
+ config : typing.Dict[str, typing.Optional[typing.Any]]
94
114
  Configuration for the research agent including search parameters and sources
95
115
 
96
- messages : typing.Sequence[typing.Dict[str, typing.Any]]
116
+ messages : typing.Sequence[typing.Dict[str, typing.Optional[typing.Any]]]
97
117
  The messages to send to the research agent
98
118
 
99
119
  request_options : typing.Optional[RequestOptions]
@@ -106,28 +126,56 @@ class AsyncResearchClient:
106
126
 
107
127
  Examples
108
128
  --------
109
- from athena.client import AsyncAthena
129
+ import asyncio
130
+
131
+ from athena import AsyncAthena
110
132
 
111
133
  client = AsyncAthena(
112
134
  api_key="YOUR_API_KEY",
113
135
  )
114
- await client.agents.research.invoke(
115
- config={"key": "value"},
116
- messages=[{"key": "value"}],
117
- )
136
+
137
+
138
+ async def main() -> None:
139
+ await client.agents.research.invoke(
140
+ config={"key": "value"},
141
+ messages=[{"key": "value"}],
142
+ )
143
+
144
+
145
+ asyncio.run(main())
118
146
  """
119
147
  _response = await self._client_wrapper.httpx_client.request(
120
148
  "api/v0/agents/research/invoke",
121
149
  method="POST",
122
- json={"config": config, "messages": messages},
150
+ json={
151
+ "config": config,
152
+ "messages": messages,
153
+ },
154
+ headers={
155
+ "content-type": "application/json",
156
+ },
123
157
  request_options=request_options,
124
158
  omit=OMIT,
125
159
  )
126
- if 200 <= _response.status_code < 300:
127
- return pydantic_v1.parse_obj_as(ResearchAgentResponse, _response.json()) # type: ignore
128
- if _response.status_code == 422:
129
- raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json())) # type: ignore
130
160
  try:
161
+ if 200 <= _response.status_code < 300:
162
+ return typing.cast(
163
+ ResearchAgentResponse,
164
+ parse_obj_as(
165
+ type_=ResearchAgentResponse, # type: ignore
166
+ object_=_response.json(),
167
+ ),
168
+ )
169
+ if _response.status_code == 422:
170
+ raise UnprocessableEntityError(
171
+ typing.cast(
172
+ typing.Optional[typing.Any],
173
+ parse_obj_as(
174
+ type_=typing.Optional[typing.Any], # type: ignore
175
+ object_=_response.json(),
176
+ ),
177
+ )
178
+ )
131
179
  _response_json = _response.json()
132
180
  except JSONDecodeError:
133
181
  raise ApiError(status_code=_response.status_code, body=_response.text)