coco-framework 0.1.0__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 (400) hide show
  1. coco/__init__.py +3 -0
  2. coco/agents/__init__.py +9 -0
  3. coco/agents/_subagent_transformer.py +307 -0
  4. coco/agents/factory.py +2164 -0
  5. coco/agents/middleware/__init__.py +100 -0
  6. coco/agents/middleware/_execution.py +385 -0
  7. coco/agents/middleware/_redaction.py +454 -0
  8. coco/agents/middleware/_retry.py +134 -0
  9. coco/agents/middleware/_trace_policy.py +61 -0
  10. coco/agents/middleware/context_editing.py +306 -0
  11. coco/agents/middleware/file_search.py +433 -0
  12. coco/agents/middleware/human_in_the_loop.py +680 -0
  13. coco/agents/middleware/internal_call_transformer.py +108 -0
  14. coco/agents/middleware/model_call_limit.py +267 -0
  15. coco/agents/middleware/model_fallback.py +417 -0
  16. coco/agents/middleware/model_retry.py +327 -0
  17. coco/agents/middleware/pii.py +878 -0
  18. coco/agents/middleware/provider_tool_search.py +307 -0
  19. coco/agents/middleware/shell_tool.py +949 -0
  20. coco/agents/middleware/summarization.py +910 -0
  21. coco/agents/middleware/todo.py +357 -0
  22. coco/agents/middleware/tool_call_limit.py +495 -0
  23. coco/agents/middleware/tool_emulator.py +251 -0
  24. coco/agents/middleware/tool_error.py +209 -0
  25. coco/agents/middleware/tool_retry.py +421 -0
  26. coco/agents/middleware/tool_selection.py +515 -0
  27. coco/agents/middleware/types.py +2201 -0
  28. coco/agents/structured_output.py +472 -0
  29. coco/chat_models/__init__.py +7 -0
  30. coco/chat_models/base.py +1082 -0
  31. coco/embeddings/__init__.py +18 -0
  32. coco/embeddings/base.py +279 -0
  33. coco/mcp/__init__.py +36 -0
  34. coco/mcp/adapter.py +245 -0
  35. coco/mcp/elicitation.py +394 -0
  36. coco/mcp/tools.py +307 -0
  37. coco/messages/__init__.py +73 -0
  38. coco/py.typed +0 -0
  39. coco/rate_limiters/__init__.py +13 -0
  40. coco/tools/__init__.py +22 -0
  41. coco/tools/tool_node.py +20 -0
  42. coco_anthropic/__init__.py +15 -0
  43. coco_anthropic/_client_utils.py +81 -0
  44. coco_anthropic/_compat.py +228 -0
  45. coco_anthropic/_sdk_compat.py +124 -0
  46. coco_anthropic/_version.py +3 -0
  47. coco_anthropic/chat_models.py +3349 -0
  48. coco_anthropic/data/__init__.py +1 -0
  49. coco_anthropic/data/_profiles.py +491 -0
  50. coco_anthropic/data/profile_augmentations.toml +69 -0
  51. coco_anthropic/experimental.py +140 -0
  52. coco_anthropic/llms.py +428 -0
  53. coco_anthropic/middleware/__init__.py +25 -0
  54. coco_anthropic/middleware/anthropic_tools.py +1197 -0
  55. coco_anthropic/middleware/bash.py +84 -0
  56. coco_anthropic/middleware/file_search.py +360 -0
  57. coco_anthropic/middleware/prompt_caching.py +276 -0
  58. coco_anthropic/output_parsers.py +93 -0
  59. coco_anthropic/py.typed +0 -0
  60. coco_core/__init__.py +20 -0
  61. coco_core/_api/__init__.py +87 -0
  62. coco_core/_api/beta_decorator.py +266 -0
  63. coco_core/_api/deprecation.py +637 -0
  64. coco_core/_api/internal.py +23 -0
  65. coco_core/_api/path.py +50 -0
  66. coco_core/_import_utils.py +41 -0
  67. coco_core/_protocol.py +625 -0
  68. coco_core/_security/__init__.py +36 -0
  69. coco_core/_security/_exceptions.py +9 -0
  70. coco_core/_security/_policy.py +313 -0
  71. coco_core/_security/_ssrf_protection.py +155 -0
  72. coco_core/_security/_transport.py +254 -0
  73. coco_core/agents.py +258 -0
  74. coco_core/caches.py +272 -0
  75. coco_core/callbacks/__init__.py +132 -0
  76. coco_core/callbacks/base.py +1229 -0
  77. coco_core/callbacks/file.py +266 -0
  78. coco_core/callbacks/manager.py +2826 -0
  79. coco_core/callbacks/stdout.py +123 -0
  80. coco_core/callbacks/streaming_stdout.py +154 -0
  81. coco_core/callbacks/usage.py +124 -0
  82. coco_core/chat_history.py +267 -0
  83. coco_core/chat_loaders.py +26 -0
  84. coco_core/chat_sessions.py +19 -0
  85. coco_core/cross_encoders.py +18 -0
  86. coco_core/document_loaders/__init__.py +36 -0
  87. coco_core/document_loaders/base.py +155 -0
  88. coco_core/document_loaders/blob_loaders.py +38 -0
  89. coco_core/documents/__init__.py +55 -0
  90. coco_core/documents/base.py +347 -0
  91. coco_core/documents/compressor.py +74 -0
  92. coco_core/documents/transformers.py +79 -0
  93. coco_core/embeddings/__init__.py +31 -0
  94. coco_core/embeddings/embeddings.py +78 -0
  95. coco_core/embeddings/fake.py +129 -0
  96. coco_core/env.py +22 -0
  97. coco_core/example_selectors/__init__.py +47 -0
  98. coco_core/example_selectors/base.py +60 -0
  99. coco_core/example_selectors/length_based.py +131 -0
  100. coco_core/example_selectors/semantic_similarity.py +366 -0
  101. coco_core/exceptions.py +166 -0
  102. coco_core/globals.py +72 -0
  103. coco_core/indexing/__init__.py +53 -0
  104. coco_core/indexing/api.py +954 -0
  105. coco_core/indexing/base.py +661 -0
  106. coco_core/indexing/in_memory.py +104 -0
  107. coco_core/language_models/__init__.py +116 -0
  108. coco_core/language_models/_compat_bridge.py +844 -0
  109. coco_core/language_models/_utils.py +343 -0
  110. coco_core/language_models/base.py +500 -0
  111. coco_core/language_models/chat_model_stream.py +1525 -0
  112. coco_core/language_models/chat_models.py +2748 -0
  113. coco_core/language_models/fake.py +137 -0
  114. coco_core/language_models/fake_chat_models.py +396 -0
  115. coco_core/language_models/llms.py +1569 -0
  116. coco_core/language_models/model_profile.py +183 -0
  117. coco_core/load/__init__.py +44 -0
  118. coco_core/load/_validation.py +191 -0
  119. coco_core/load/dump.py +120 -0
  120. coco_core/load/load.py +828 -0
  121. coco_core/load/mapping.py +1097 -0
  122. coco_core/load/serializable.py +406 -0
  123. coco_core/messages/__init__.py +198 -0
  124. coco_core/messages/ai.py +854 -0
  125. coco_core/messages/base.py +518 -0
  126. coco_core/messages/block_translators/__init__.py +112 -0
  127. coco_core/messages/block_translators/anthropic.py +534 -0
  128. coco_core/messages/block_translators/bedrock.py +108 -0
  129. coco_core/messages/block_translators/bedrock_converse.py +329 -0
  130. coco_core/messages/block_translators/coco_v0.py +303 -0
  131. coco_core/messages/block_translators/google_genai.py +610 -0
  132. coco_core/messages/block_translators/google_vertexai.py +21 -0
  133. coco_core/messages/block_translators/groq.py +157 -0
  134. coco_core/messages/block_translators/openai.py +1089 -0
  135. coco_core/messages/chat.py +64 -0
  136. coco_core/messages/content.py +1488 -0
  137. coco_core/messages/function.py +62 -0
  138. coco_core/messages/human.py +70 -0
  139. coco_core/messages/modifier.py +33 -0
  140. coco_core/messages/system.py +70 -0
  141. coco_core/messages/tool.py +418 -0
  142. coco_core/messages/utils.py +2442 -0
  143. coco_core/output_parsers/__init__.py +101 -0
  144. coco_core/output_parsers/base.py +361 -0
  145. coco_core/output_parsers/format_instructions.py +16 -0
  146. coco_core/output_parsers/json.py +139 -0
  147. coco_core/output_parsers/list.py +249 -0
  148. coco_core/output_parsers/openai_functions.py +324 -0
  149. coco_core/output_parsers/openai_tools.py +383 -0
  150. coco_core/output_parsers/pydantic.py +141 -0
  151. coco_core/output_parsers/string.py +63 -0
  152. coco_core/output_parsers/transform.py +175 -0
  153. coco_core/output_parsers/xml.py +301 -0
  154. coco_core/outputs/__init__.py +66 -0
  155. coco_core/outputs/chat_generation.py +147 -0
  156. coco_core/outputs/chat_result.py +39 -0
  157. coco_core/outputs/generation.py +80 -0
  158. coco_core/outputs/llm_result.py +112 -0
  159. coco_core/outputs/run_info.py +22 -0
  160. coco_core/prompt_values.py +161 -0
  161. coco_core/prompts/__init__.py +101 -0
  162. coco_core/prompts/base.py +508 -0
  163. coco_core/prompts/chat.py +1495 -0
  164. coco_core/prompts/dict.py +177 -0
  165. coco_core/prompts/few_shot.py +483 -0
  166. coco_core/prompts/few_shot_with_templates.py +239 -0
  167. coco_core/prompts/image.py +177 -0
  168. coco_core/prompts/loading.py +296 -0
  169. coco_core/prompts/message.py +97 -0
  170. coco_core/prompts/prompt.py +312 -0
  171. coco_core/prompts/string.py +399 -0
  172. coco_core/prompts/structured.py +214 -0
  173. coco_core/py.typed +0 -0
  174. coco_core/rate_limiters.py +254 -0
  175. coco_core/retrievers.py +328 -0
  176. coco_core/runnables/__init__.py +136 -0
  177. coco_core/runnables/base.py +6725 -0
  178. coco_core/runnables/branch.py +457 -0
  179. coco_core/runnables/config.py +712 -0
  180. coco_core/runnables/configurable.py +710 -0
  181. coco_core/runnables/fallbacks.py +663 -0
  182. coco_core/runnables/graph.py +741 -0
  183. coco_core/runnables/graph_ascii.py +366 -0
  184. coco_core/runnables/graph_mermaid.py +503 -0
  185. coco_core/runnables/graph_png.py +215 -0
  186. coco_core/runnables/history.py +631 -0
  187. coco_core/runnables/passthrough.py +849 -0
  188. coco_core/runnables/retry.py +379 -0
  189. coco_core/runnables/router.py +239 -0
  190. coco_core/runnables/schema.py +188 -0
  191. coco_core/runnables/utils.py +798 -0
  192. coco_core/stores.py +291 -0
  193. coco_core/structured_query.py +203 -0
  194. coco_core/sys_info.py +154 -0
  195. coco_core/tools/__init__.py +95 -0
  196. coco_core/tools/base.py +1948 -0
  197. coco_core/tools/convert.py +497 -0
  198. coco_core/tools/render.py +67 -0
  199. coco_core/tools/retriever.py +94 -0
  200. coco_core/tools/simple.py +208 -0
  201. coco_core/tools/structured.py +322 -0
  202. coco_core/tracers/__init__.py +47 -0
  203. coco_core/tracers/_compat.py +95 -0
  204. coco_core/tracers/_context.py +337 -0
  205. coco_core/tracers/_streaming.py +53 -0
  206. coco_core/tracers/base.py +955 -0
  207. coco_core/tracers/coco.py +279 -0
  208. coco_core/tracers/context.py +205 -0
  209. coco_core/tracers/core.py +766 -0
  210. coco_core/tracers/event_stream.py +1101 -0
  211. coco_core/tracers/log_stream.py +765 -0
  212. coco_core/tracers/memory_stream.py +200 -0
  213. coco_core/tracers/root_listeners.py +130 -0
  214. coco_core/tracers/run_collector.py +40 -0
  215. coco_core/tracers/schemas.py +14 -0
  216. coco_core/tracers/stdout.py +205 -0
  217. coco_core/utils/__init__.py +111 -0
  218. coco_core/utils/_gateway.py +274 -0
  219. coco_core/utils/_merge.py +218 -0
  220. coco_core/utils/aiter.py +357 -0
  221. coco_core/utils/env.py +86 -0
  222. coco_core/utils/formatting.py +80 -0
  223. coco_core/utils/function_calling.py +869 -0
  224. coco_core/utils/html.py +132 -0
  225. coco_core/utils/image.py +15 -0
  226. coco_core/utils/input.py +82 -0
  227. coco_core/utils/interactive_env.py +12 -0
  228. coco_core/utils/iter.py +229 -0
  229. coco_core/utils/json.py +230 -0
  230. coco_core/utils/json_schema.py +273 -0
  231. coco_core/utils/mustache.py +706 -0
  232. coco_core/utils/pydantic.py +631 -0
  233. coco_core/utils/strings.py +69 -0
  234. coco_core/utils/usage.py +61 -0
  235. coco_core/utils/utils.py +528 -0
  236. coco_core/utils/uuid.py +57 -0
  237. coco_core/vectorstores/__init__.py +53 -0
  238. coco_core/vectorstores/base.py +1111 -0
  239. coco_core/vectorstores/in_memory.py +546 -0
  240. coco_core/vectorstores/utils.py +163 -0
  241. coco_core/version.py +3 -0
  242. coco_framework-0.1.0.dist-info/METADATA +184 -0
  243. coco_framework-0.1.0.dist-info/RECORD +400 -0
  244. coco_framework-0.1.0.dist-info/WHEEL +4 -0
  245. coco_framework-0.1.0.dist-info/licenses/LICENSE +24 -0
  246. coco_framework-0.1.0.dist-info/licenses/NOTICE +15 -0
  247. coco_graph/_internal/__init__.py +4 -0
  248. coco_graph/_internal/_cache.py +31 -0
  249. coco_graph/_internal/_config.py +474 -0
  250. coco_graph/_internal/_constants.py +140 -0
  251. coco_graph/_internal/_fields.py +213 -0
  252. coco_graph/_internal/_future.py +220 -0
  253. coco_graph/_internal/_pydantic.py +275 -0
  254. coco_graph/_internal/_queue.py +124 -0
  255. coco_graph/_internal/_replay.py +90 -0
  256. coco_graph/_internal/_retry.py +29 -0
  257. coco_graph/_internal/_runnable.py +969 -0
  258. coco_graph/_internal/_scratchpad.py +19 -0
  259. coco_graph/_internal/_serde.py +253 -0
  260. coco_graph/_internal/_timeout.py +25 -0
  261. coco_graph/_internal/_typing.py +54 -0
  262. coco_graph/cache/base/__init__.py +48 -0
  263. coco_graph/cache/base/py.typed +0 -0
  264. coco_graph/cache/memory/__init__.py +73 -0
  265. coco_graph/cache/redis/__init__.py +144 -0
  266. coco_graph/cache/sqlite/__init__.py +120 -0
  267. coco_graph/callbacks.py +394 -0
  268. coco_graph/channels/__init__.py +29 -0
  269. coco_graph/channels/any_value.py +72 -0
  270. coco_graph/channels/base.py +121 -0
  271. coco_graph/channels/binop.py +155 -0
  272. coco_graph/channels/delta.py +202 -0
  273. coco_graph/channels/ephemeral_value.py +79 -0
  274. coco_graph/channels/last_value.py +151 -0
  275. coco_graph/channels/named_barrier_value.py +167 -0
  276. coco_graph/channels/topic.py +94 -0
  277. coco_graph/channels/untracked_value.py +73 -0
  278. coco_graph/checkpoint/base/__init__.py +861 -0
  279. coco_graph/checkpoint/base/id.py +109 -0
  280. coco_graph/checkpoint/base/py.typed +0 -0
  281. coco_graph/checkpoint/memory/__init__.py +698 -0
  282. coco_graph/checkpoint/memory/py.typed +0 -0
  283. coco_graph/checkpoint/postgres/__init__.py +602 -0
  284. coco_graph/checkpoint/postgres/_ainternal.py +23 -0
  285. coco_graph/checkpoint/postgres/_internal.py +21 -0
  286. coco_graph/checkpoint/postgres/aio.py +691 -0
  287. coco_graph/checkpoint/postgres/base.py +666 -0
  288. coco_graph/checkpoint/postgres/py.typed +0 -0
  289. coco_graph/checkpoint/postgres/shallow.py +967 -0
  290. coco_graph/checkpoint/serde/__init__.py +0 -0
  291. coco_graph/checkpoint/serde/_msgpack.py +97 -0
  292. coco_graph/checkpoint/serde/base.py +64 -0
  293. coco_graph/checkpoint/serde/encrypted.py +80 -0
  294. coco_graph/checkpoint/serde/event_hooks.py +52 -0
  295. coco_graph/checkpoint/serde/jsonplus.py +883 -0
  296. coco_graph/checkpoint/serde/py.typed +0 -0
  297. coco_graph/checkpoint/serde/types.py +68 -0
  298. coco_graph/checkpoint/sqlite/__init__.py +645 -0
  299. coco_graph/checkpoint/sqlite/_delta.py +172 -0
  300. coco_graph/checkpoint/sqlite/aio.py +738 -0
  301. coco_graph/checkpoint/sqlite/py.typed +0 -0
  302. coco_graph/checkpoint/sqlite/utils.py +116 -0
  303. coco_graph/config.py +196 -0
  304. coco_graph/constants.py +64 -0
  305. coco_graph/errors.py +241 -0
  306. coco_graph/func/__init__.py +620 -0
  307. coco_graph/graph/__init__.py +12 -0
  308. coco_graph/graph/_branch.py +225 -0
  309. coco_graph/graph/_node.py +103 -0
  310. coco_graph/graph/message.py +437 -0
  311. coco_graph/graph/state.py +1978 -0
  312. coco_graph/graph/ui.py +227 -0
  313. coco_graph/managed/__init__.py +3 -0
  314. coco_graph/managed/base.py +31 -0
  315. coco_graph/managed/is_last_step.py +24 -0
  316. coco_graph/prebuilt/__init__.py +23 -0
  317. coco_graph/prebuilt/_tool_call_stream.py +117 -0
  318. coco_graph/prebuilt/_tool_call_transformer.py +165 -0
  319. coco_graph/prebuilt/chat_agent_executor.py +1015 -0
  320. coco_graph/prebuilt/interrupt.py +105 -0
  321. coco_graph/prebuilt/py.typed +0 -0
  322. coco_graph/prebuilt/tool_node.py +2030 -0
  323. coco_graph/prebuilt/tool_validator.py +221 -0
  324. coco_graph/pregel/__init__.py +3 -0
  325. coco_graph/pregel/_algo.py +1460 -0
  326. coco_graph/pregel/_call.py +298 -0
  327. coco_graph/pregel/_checkpoint.py +331 -0
  328. coco_graph/pregel/_config.py +0 -0
  329. coco_graph/pregel/_draw.py +294 -0
  330. coco_graph/pregel/_executor.py +223 -0
  331. coco_graph/pregel/_io.py +174 -0
  332. coco_graph/pregel/_log.py +3 -0
  333. coco_graph/pregel/_loop.py +1988 -0
  334. coco_graph/pregel/_messages.py +461 -0
  335. coco_graph/pregel/_read.py +312 -0
  336. coco_graph/pregel/_retry.py +854 -0
  337. coco_graph/pregel/_runner.py +941 -0
  338. coco_graph/pregel/_tools.py +268 -0
  339. coco_graph/pregel/_utils.py +221 -0
  340. coco_graph/pregel/_validate.py +120 -0
  341. coco_graph/pregel/_write.py +192 -0
  342. coco_graph/pregel/debug.py +302 -0
  343. coco_graph/pregel/main.py +4364 -0
  344. coco_graph/pregel/protocol.py +288 -0
  345. coco_graph/pregel/types.py +38 -0
  346. coco_graph/py.typed +0 -0
  347. coco_graph/runtime.py +350 -0
  348. coco_graph/store/base/__init__.py +1322 -0
  349. coco_graph/store/base/batch.py +371 -0
  350. coco_graph/store/base/embed.py +436 -0
  351. coco_graph/store/base/py.typed +0 -0
  352. coco_graph/store/memory/__init__.py +592 -0
  353. coco_graph/store/memory/py.typed +0 -0
  354. coco_graph/store/postgres/__init__.py +4 -0
  355. coco_graph/store/postgres/aio.py +592 -0
  356. coco_graph/store/postgres/base.py +1480 -0
  357. coco_graph/store/postgres/py.typed +0 -0
  358. coco_graph/store/sqlite/__init__.py +4 -0
  359. coco_graph/store/sqlite/aio.py +632 -0
  360. coco_graph/store/sqlite/base.py +1514 -0
  361. coco_graph/stream/__init__.py +45 -0
  362. coco_graph/stream/_convert.py +32 -0
  363. coco_graph/stream/_mux.py +523 -0
  364. coco_graph/stream/_types.py +330 -0
  365. coco_graph/stream/run_stream.py +735 -0
  366. coco_graph/stream/stream_channel.py +341 -0
  367. coco_graph/stream/transformers.py +1039 -0
  368. coco_graph/types.py +1079 -0
  369. coco_graph/typing.py +48 -0
  370. coco_graph/utils/__init__.py +1 -0
  371. coco_graph/utils/config.py +4 -0
  372. coco_graph/utils/runnable.py +3 -0
  373. coco_graph/version.py +13 -0
  374. coco_graph/warnings.py +69 -0
  375. coco_openai/__init__.py +20 -0
  376. coco_openai/_compat.py +45 -0
  377. coco_openai/_version.py +3 -0
  378. coco_openai/chat_models/__init__.py +6 -0
  379. coco_openai/chat_models/_client_utils.py +684 -0
  380. coco_openai/chat_models/_compat.py +525 -0
  381. coco_openai/chat_models/azure.py +1227 -0
  382. coco_openai/chat_models/base.py +5654 -0
  383. coco_openai/chat_models/codex.py +557 -0
  384. coco_openai/chatgpt_oauth.py +1063 -0
  385. coco_openai/data/__init__.py +1 -0
  386. coco_openai/data/_profiles.py +1655 -0
  387. coco_openai/data/profile_augmentations.toml +127 -0
  388. coco_openai/embeddings/__init__.py +6 -0
  389. coco_openai/embeddings/azure.py +240 -0
  390. coco_openai/embeddings/base.py +820 -0
  391. coco_openai/llms/__init__.py +6 -0
  392. coco_openai/llms/azure.py +249 -0
  393. coco_openai/llms/base.py +919 -0
  394. coco_openai/middleware/__init__.py +11 -0
  395. coco_openai/middleware/openai_moderation.py +484 -0
  396. coco_openai/output_parsers/__init__.py +9 -0
  397. coco_openai/output_parsers/tools.py +9 -0
  398. coco_openai/py.typed +0 -0
  399. coco_openai/tools/__init__.py +5 -0
  400. coco_openai/tools/custom_tool.py +111 -0
@@ -0,0 +1,127 @@
1
+ provider = "openai"
2
+
3
+ [overrides]
4
+ image_url_inputs = true
5
+ pdf_inputs = true
6
+ pdf_tool_message = true
7
+ image_tool_message = true
8
+ tool_choice = true
9
+ tool_call_streaming = true
10
+
11
+ [overrides."gpt-3.5-turbo"]
12
+ image_url_inputs = false
13
+ pdf_inputs = false
14
+ pdf_tool_message = false
15
+ image_tool_message = false
16
+
17
+ [overrides."gpt-5.1-codex"]
18
+ max_input_tokens = 272000
19
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
20
+
21
+ [overrides."gpt-5.2-pro"]
22
+ max_input_tokens = 272000
23
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
24
+
25
+ [overrides."gpt-5.1-codex-mini"]
26
+ max_input_tokens = 272000
27
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
28
+
29
+ [overrides."gpt-5.2-chat-latest"]
30
+ max_input_tokens = 272000
31
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
32
+
33
+ [overrides."gpt-5.1"]
34
+ max_input_tokens = 272000
35
+ reasoning_effort_levels = ["none", "low", "medium", "high"]
36
+
37
+ [overrides."gpt-5-nano"]
38
+ max_input_tokens = 272000
39
+ reasoning_effort_levels = ["minimal", "low", "medium", "high"]
40
+
41
+ [overrides."gpt-5-codex"]
42
+ max_input_tokens = 272000
43
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
44
+
45
+ [overrides."gpt-5-mini"]
46
+ max_input_tokens = 272000
47
+ reasoning_effort_levels = ["minimal", "low", "medium", "high"]
48
+
49
+ [overrides."gpt-5.1-codex-max"]
50
+ max_input_tokens = 272000
51
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
52
+
53
+ [overrides."gpt-5-chat-latest"]
54
+ max_input_tokens = 272000
55
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
56
+
57
+ [overrides."gpt-5"]
58
+ max_input_tokens = 272000
59
+ reasoning_effort_levels = ["minimal", "low", "medium", "high"]
60
+
61
+ [overrides."gpt-5-pro"]
62
+ max_input_tokens = 272000
63
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
64
+
65
+ [overrides."gpt-5.2"]
66
+ max_input_tokens = 272000
67
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
68
+
69
+ [overrides."gpt-5.1-chat-latest"]
70
+ max_input_tokens = 272000
71
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
72
+
73
+ [overrides."gpt-5.2-codex"]
74
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
75
+
76
+ [overrides."gpt-5.3-codex"]
77
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
78
+
79
+ [overrides."gpt-5.3-codex-spark"]
80
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
81
+
82
+ [overrides."gpt-5.4"]
83
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
84
+
85
+ [overrides."gpt-5.4-mini"]
86
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
87
+
88
+ [overrides."gpt-5.4-nano"]
89
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
90
+
91
+ [overrides."gpt-5.4-pro"]
92
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
93
+
94
+ [overrides."gpt-5.5"]
95
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
96
+ reasoning_effort_default = "medium"
97
+
98
+ [overrides."gpt-5.5-pro"]
99
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh"]
100
+ reasoning_effort_default = "medium"
101
+
102
+ [overrides."gpt-5.6"]
103
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh", "max"]
104
+ reasoning_effort_default = "medium"
105
+
106
+ [overrides."gpt-5.6-luna"]
107
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh", "max"]
108
+ reasoning_effort_default = "medium"
109
+
110
+ [overrides."gpt-5.6-sol"]
111
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh", "max"]
112
+ reasoning_effort_default = "medium"
113
+
114
+ [overrides."gpt-5.6-terra"]
115
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh", "max"]
116
+ reasoning_effort_default = "medium"
117
+
118
+ [overrides."gpt-6-astra"]
119
+ reasoning_effort_levels = ["low", "medium", "high", "xhigh", "max"]
120
+
121
+ [overrides."gpt-6-sol"]
122
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh", "max"]
123
+ reasoning_effort_default = "medium"
124
+
125
+ [overrides."gpt-6-luna"]
126
+ reasoning_effort_levels = ["none", "low", "medium", "high", "xhigh", "max"]
127
+ reasoning_effort_default = "medium"
@@ -0,0 +1,6 @@
1
+ """Module for OpenAI embeddings."""
2
+
3
+ from coco_openai.embeddings.azure import AzureOpenAIEmbeddings
4
+ from coco_openai.embeddings.base import OpenAIEmbeddings
5
+
6
+ __all__ = ["AzureOpenAIEmbeddings", "OpenAIEmbeddings"]
@@ -0,0 +1,240 @@
1
+ """Azure OpenAI embeddings wrapper."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from collections.abc import Awaitable, Callable
6
+ from typing import cast
7
+
8
+ import openai
9
+ from coco_core.utils import from_env, secret_from_env
10
+ from pydantic import Field, SecretStr, model_validator
11
+ from typing_extensions import Self
12
+
13
+ from coco_openai.embeddings.base import OpenAIEmbeddings
14
+
15
+
16
+ class AzureOpenAIEmbeddings(OpenAIEmbeddings): # type: ignore[override]
17
+ """AzureOpenAI embedding model integration.
18
+
19
+ Setup:
20
+ To access AzureOpenAI embedding models you'll need to create an Azure account,
21
+ get an API key, and install the `trustcoco-openai` integration package.
22
+
23
+ You'll need to have an Azure OpenAI instance deployed.
24
+ You can deploy a version on Azure Portal following this
25
+ [guide](https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/create-resource?pivots=web-portal).
26
+
27
+ Once you have your instance running, make sure you have the name of your
28
+ instance and key. You can find the key in the Azure Portal,
29
+ under the “Keys and Endpoint” section of your instance.
30
+
31
+ ```bash
32
+ pip install -U coco_openai
33
+
34
+ # Set up your environment variables (or pass them directly to the model)
35
+ export AZURE_OPENAI_API_KEY="your-api-key"
36
+ export AZURE_OPENAI_ENDPOINT="https://<your-endpoint>.openai.azure.com/"
37
+ export AZURE_OPENAI_API_VERSION="2024-02-01"
38
+ ```
39
+
40
+ Key init args — completion params:
41
+ model:
42
+ Name of `AzureOpenAI` model to use.
43
+ dimensions:
44
+ Number of dimensions for the embeddings. Can be specified only if the
45
+ underlying model supports it.
46
+
47
+ See full list of supported init args and their descriptions in the params section.
48
+
49
+ Instantiate:
50
+ ```python
51
+ from coco_openai import AzureOpenAIEmbeddings
52
+
53
+ embeddings = AzureOpenAIEmbeddings(
54
+ model="text-embedding-3-large"
55
+ # dimensions: int | None = None, # Can specify dimensions with new text-embedding-3 models
56
+ # azure_endpoint="https://<your-endpoint>.openai.azure.com/", If not provided, will read env variable AZURE_OPENAI_ENDPOINT
57
+ # api_key=... # Can provide an API key directly. If missing read env variable AZURE_OPENAI_API_KEY
58
+ # openai_api_version=..., # If not provided, will read env variable AZURE_OPENAI_API_VERSION
59
+ )
60
+ ```
61
+
62
+ Embed single text:
63
+ ```python
64
+ input_text = "The meaning of life is 42"
65
+ vector = embed.embed_query(input_text)
66
+ print(vector[:3])
67
+ ```
68
+ ```python
69
+ [-0.024603435769677162, -0.007543657906353474, 0.0039630369283258915]
70
+ ```
71
+
72
+ Embed multiple texts:
73
+ ```python
74
+ input_texts = ["Document 1...", "Document 2..."]
75
+ vectors = embed.embed_documents(input_texts)
76
+ print(len(vectors))
77
+ # The first 3 coordinates for the first vector
78
+ print(vectors[0][:3])
79
+ ```
80
+ ```python
81
+ 2
82
+ [-0.024603435769677162, -0.007543657906353474, 0.0039630369283258915]
83
+ ```
84
+
85
+ Async:
86
+ ```python
87
+ vector = await embed.aembed_query(input_text)
88
+ print(vector[:3])
89
+
90
+ # multiple:
91
+ # await embed.aembed_documents(input_texts)
92
+ ```
93
+ ```python
94
+ [-0.009100092574954033, 0.005071679595857859, -0.0029193938244134188]
95
+ ```
96
+ """ # noqa: E501
97
+
98
+ azure_endpoint: str | None = Field(
99
+ default_factory=from_env("AZURE_OPENAI_ENDPOINT", default=None)
100
+ )
101
+ """Your Azure endpoint, including the resource.
102
+
103
+ Automatically inferred from env var `AZURE_OPENAI_ENDPOINT` if not provided.
104
+
105
+ Example: `https://example-resource.azure.openai.com/`
106
+ """
107
+ deployment: str | None = Field(default=None, alias="azure_deployment")
108
+ """A model deployment.
109
+
110
+ If given sets the base client URL to include `/deployments/{azure_deployment}`.
111
+
112
+ !!! note
113
+ This means you won't be able to use non-deployment endpoints.
114
+
115
+ """
116
+ # Check OPENAI_KEY for backwards compatibility.
117
+ # TODO: Remove OPENAI_API_KEY support to avoid possible conflict when using
118
+ # other forms of azure credentials.
119
+ openai_api_key: SecretStr | None = Field(
120
+ alias="api_key",
121
+ default_factory=secret_from_env(
122
+ ["AZURE_OPENAI_API_KEY", "OPENAI_API_KEY"], default=None
123
+ ),
124
+ )
125
+ """Automatically inferred from env var `AZURE_OPENAI_API_KEY` if not provided."""
126
+ openai_api_version: str | None = Field(
127
+ default_factory=from_env("OPENAI_API_VERSION", default="2023-05-15"),
128
+ alias="api_version",
129
+ )
130
+ """Automatically inferred from env var `OPENAI_API_VERSION` if not provided.
131
+
132
+ Set to `'2023-05-15'` by default if env variable `OPENAI_API_VERSION` is not
133
+ set.
134
+ """
135
+ azure_ad_token: SecretStr | None = Field(
136
+ default_factory=secret_from_env("AZURE_OPENAI_AD_TOKEN", default=None)
137
+ )
138
+ """Your Azure Active Directory token.
139
+
140
+ Automatically inferred from env var `AZURE_OPENAI_AD_TOKEN` if not provided.
141
+
142
+ [For more, see this page.](https://www.microsoft.com/en-us/security/business/identity-access/microsoft-entra-id)
143
+ """
144
+ azure_ad_token_provider: Callable[[], str] | None = None
145
+ """A function that returns an Azure Active Directory token.
146
+
147
+ Will be invoked on every sync request. For async requests,
148
+ will be invoked if `azure_ad_async_token_provider` is not provided.
149
+ """
150
+ azure_ad_async_token_provider: Callable[[], Awaitable[str]] | None = None
151
+ """A function that returns an Azure Active Directory token.
152
+
153
+ Will be invoked on every async request.
154
+ """
155
+ openai_api_type: str | None = Field(
156
+ default_factory=from_env("OPENAI_API_TYPE", default="azure")
157
+ )
158
+ validate_base_url: bool = True
159
+ chunk_size: int = 2048
160
+ """Maximum number of texts to embed in each batch"""
161
+
162
+ @model_validator(mode="after")
163
+ def validate_environment(self) -> Self:
164
+ """Validate that api key and python package exists in environment."""
165
+ # For backwards compatibility. Before openai v1, no distinction was made
166
+ # between azure_endpoint and base_url (openai_api_base).
167
+ openai_api_base = self.openai_api_base
168
+ if openai_api_base and self.validate_base_url:
169
+ # Only validate openai_api_base if azure_endpoint is not provided
170
+ if not self.azure_endpoint and "/openai" not in openai_api_base:
171
+ self.openai_api_base = cast(str, self.openai_api_base) + "/openai"
172
+ msg = (
173
+ "As of openai>=1.0.0, Azure endpoints should be specified via "
174
+ "the `azure_endpoint` param not `openai_api_base` "
175
+ "(or alias `base_url`). "
176
+ )
177
+ raise ValueError(msg)
178
+ if self.deployment:
179
+ msg = (
180
+ "As of openai>=1.0.0, if `deployment` (or alias "
181
+ "`azure_deployment`) is specified then "
182
+ "`openai_api_base` (or alias `base_url`) should not be. "
183
+ "Instead use `deployment` (or alias `azure_deployment`) "
184
+ "and `azure_endpoint`."
185
+ )
186
+ raise ValueError(msg)
187
+ client_params: dict = {
188
+ "api_version": self.openai_api_version,
189
+ "azure_endpoint": self.azure_endpoint,
190
+ "azure_deployment": self.deployment,
191
+ "api_key": (
192
+ self.openai_api_key.get_secret_value() if self.openai_api_key else None
193
+ ),
194
+ "azure_ad_token": (
195
+ self.azure_ad_token.get_secret_value() if self.azure_ad_token else None
196
+ ),
197
+ "azure_ad_token_provider": self.azure_ad_token_provider,
198
+ "organization": self.openai_organization,
199
+ "base_url": self.openai_api_base,
200
+ "timeout": self.request_timeout,
201
+ "max_retries": self.max_retries,
202
+ "default_headers": {
203
+ "User-Agent": "trustcoco-partner-python-azure-openai",
204
+ **(self.default_headers or {}),
205
+ },
206
+ "default_query": self.default_query,
207
+ }
208
+ if (
209
+ self.azure_ad_token
210
+ or self.azure_ad_token_provider
211
+ or self.azure_ad_async_token_provider
212
+ ):
213
+ client_params["api_key"] = None
214
+ if not self.client and (
215
+ not self.azure_ad_async_token_provider
216
+ or self.azure_ad_token
217
+ or self.azure_ad_token_provider
218
+ ):
219
+ sync_specific: dict = {"http_client": self.http_client}
220
+ self.client = openai.AzureOpenAI(
221
+ **client_params, # type: ignore[arg-type]
222
+ **sync_specific,
223
+ ).embeddings
224
+ if not self.async_client:
225
+ async_specific: dict = {"http_client": self.http_async_client}
226
+
227
+ if self.azure_ad_async_token_provider:
228
+ client_params["azure_ad_token_provider"] = (
229
+ self.azure_ad_async_token_provider
230
+ )
231
+
232
+ self.async_client = openai.AsyncAzureOpenAI(
233
+ **client_params, # type: ignore[arg-type]
234
+ **async_specific,
235
+ ).embeddings
236
+ return self
237
+
238
+ @property
239
+ def _llm_type(self) -> str:
240
+ return "azure-openai-chat"