[ { "path": ".env-example", "content": "MODEL_PROVIDER=stepfun\nMODEL_BASE=https://api.stepfun.com\nSTEP_MODEL_API_KEY=your-model-api-key\nSTEP_SEARCH_API_BASE=https://api.stepfun.com\nSTEP_SEARCH_API_KEY=your-search-api-key\n" }, { "path": ".gitignore", "content": "# common file types\n*.log\n*.ipynb\n#*.json\n# Allow json files in frontend directory\n!agent_cortex/adapter/frontend/**/*.json\n*.csv\n*.jsonl\n*.png\n*.jpg\n\n# local dev files/directories\n**/_version.py\ndist/\n.env\n.venv\n__pycache__/\ndownloads*\nlogs\ndeprecated/\n*_debug.yaml\nscripts/*_debug.py\ntests/*_debug.py\n*.png\n!assets/**\ntrash/\ntraces/\n.idea/\n.cursor/\n*.db\n\neval/\n" }, { "path": "LICENSE", "content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n" }, { "path": "README.md", "content": "# Step-DeepResearch\n\n
\n \n
\n
\n \"Homepage\"  \n \"Twitter  \n \"Discord\"\n
\n
\n  \n \n \"License\"\n
\n\n\n## News\n\n* Feb 2, 2026: 👋 We have released **Step 3.5 Flash**, achieving **65.27** on ResearchRubrics. Try it out by setting the environment variable `MODEL_NAME=step-3.5-flash`. [Details](https://static.stepfun.com/blog/step-3.5-flash/)\n\n* Dec 25, 2025: 👋 You can join our group chat to get updates on your beta API application status and the latest project developments.\n
\n \"WeChat\n \"Feishu\n
\n\n* Dec 24, 2025: 👋 We have made our technical report available. [Read](https://arxiv.org/pdf/2512.20491)\n\n\n## Introduction\n### Model Summary\n **Step-DeepResearch** is a cost-effective, end-to-end deep research agent model designed for autonomous information exploration and professional report generation in open-ended research scenarios.\n - **Atomic Capability Integration**: By decomposing complex research tasks into trainable atomic capabilities—including planning, information seeking, reflection and cross-validation, and professional report generation—and achieving deep internalization at the model level, the system ensures closed-loop reflection and dynamic correction within a single inference pass.\n - **Progressive Training Pipeline**: We establish a complete optimization path from Agentic Mid-Training to Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL), reshaping the training objective from \"predicting the next token\" to \"deciding the next atomic action.\" This approach effectively enhances the model's adaptive capabilities and generalization performance in complex environments.\n - **Strong Performance Across Model Scales**: With only 32B parameters, Step-DeepResearch achieves 61.4% on Scale AI ResearchRubrics, matching OpenAI Deep Research and Gemini Deep Research. In expert human evaluations on ADR-Bench, its Elo score significantly outperforms larger models including DeepSeek-v3.2 and GLM-4.6, and rivals top-tier closed-source models.\n - **New model (Step 3.5 Flash)**: **Step 3.5 Flash** reaches **65.27** on ResearchRubrics, delivering research quality that competes with OpenAI and Gemini Deep Research while maintaining significantly higher inference efficiency. [Details](https://static.stepfun.com/blog/step-3.5-flash/)\n - **Superior Cost-Effectiveness**: With extremely low deployment and inference costs while maintaining expert-level research capabilities, Step-DeepResearch stands as the most cost-effective deep research agent solution currently available in the industry.\n - **Access**: Available via StepFun Open Platform API, free for the first month.\n\n
\n \"Cost-Performance
\n (left) Cost-Efficiency on ResearchRubrics: Step-DeepResearch achieves near-top performance (61.42) while significantly reducing inference costs (RMB), positioned at the high-efficiency frontier. \n (right) Expert Evaluation on ADR-Bench: Step-DeepResearch consistently leads in Elo ratings across all dimensions, rivaling top-tier closed-source models.\n
\n\n
\n \"Performance
\n Performance on ResearchRubrics: Step 3.5 Flash achieves 65.27 on the benchmark.\n
\n \n### System Architecture\n Step-DeepResearch adopts a single-agent architecture based on the ReAct paradigm, enabling autonomous deep research through a dynamic cycle of reasoning, action, and reflection.\n\n Building upon our internal proprietary toolset, we provide a streamlined toolset with local implementation in the framework, including `batch_web_surfer` (batch web search and browsing), `file` (file reading, writing, and editing), `todo` (task state management), and `shell` (interactive command execution), to support the complete research workflow. For a comprehensive introduction to our full toolset, please refer to Section 5.2 of the technical report.\n\n
\n \"results\"\n

Step-DeepResearch System Architecture: The agent operates within a ReAct loop, utilizing a specialized toolset (e.g., batch_web_surfer, todo, shell) for planning, execution, and reflection to generate comprehensive research reports.

\n
\n\n## Case\n\nWe provide a sample case below, including the query and the model-generated report.\n\n### Query\n\n```text\n我是个 AI 产品经理，正在考虑给产品增加个性化记忆功能。想系统了解一下现在 Agent 里的 Memory 模块都在往哪些方向发展，包括短期和长期记忆是怎么做的、各种技术路线的差异和取舍。你帮我写个调研报告，从工程落地和未来演进的角度分析一下 哪些方向更值得投入\n```\n\n### Model Output (Report)\n\n
\nClick to expand the report\n\n\n# AI Agent记忆模块深度调研报告：技术演进与产品化路径\n\n## 一、引言与核心问题界定\n\n### 1.1 记忆能力为何成为Agent时代的分水岭\n\n在Agentic AI加速渗透的2025年，智能体已从\"能对话、会调用工具\"的基础形态，向\"具备长期记忆、可自主进化\"的高阶阶段迈进[大模型进阶之路：AI Agent记忆能力构建技术详解（值得收藏）](https://blog.csdn.net/xxue345678/article/details/150983939)。正如OpenAI CEO山姆·奥尔特曼在2025年12月的Big Technology访谈中所言：\"这是整个系统里我个人最期待的部分之一。AI的下一个重大飞跃并非是更敏锐的推理能力，而是更为根本的记忆\"[超级Agent重要拼图？奥尔特曼点名“AI记忆” 存储环节迎来新叙事](https://finance.sina.com.cn/stock/t/2025-12-22/doc-inhcsmfc4526733.shtml)。他进一步指出，\"人类本身是有局限的：即使你拥有世界上最好的私人助理，他们也不可能记住你说过的每一句话，不可能读过你的每一封邮件，不可能看过你写的每一份文件\"——而这正是AI能够做到的[超级Agent重要拼图？奥尔特曼点名“AI记忆” 存储环节迎来新叙事](http://m.cls.cn/detail/2236511)。\n\n当前，大多数人以为通过更大的上下文窗口或巧妙的提示词工程，AI就拥有了\"记忆\"，但真相是，大多数AI Agent仍是无状态的，无法从过去的交互中学习，也无法随时间适应用户需求[探寻AI Agent 中隐秘的角落：记忆(Memory) - 定义、价值与实践](https://developer.volcengine.com/articles/7540134113190412324)。要从一次性工具迈向真正智能的伙伴，我们需要为AI赋予真正的记忆，而非仅仅依赖更大的提示或更强的检索。\n\n### 1.2 记忆的本质定义\n\n在AI Agent中，记忆是指跨时间、任务和多次用户交互，保留并调用相关信息的能力[探寻AI Agent 中隐秘的角落：记忆(Memory) - 定义、价值与实践](https://developer.volcengine.com/articles/7540134113190412324)。它让AI能记住过去发生的事，并利用这些信息优化未来的行为。记忆不是简单地存储聊天记录，也不是把更多数据塞进提示框——它是一种持久的内部状态，随着每次交互不断进化，哪怕间隔数周或数月，依然能为AI提供连续的上下文[探寻AI Agent 中隐秘的角落：记忆(Memory) - 定义、价值与实践](https://developer.volcengine.com/articles/7540134113190412324)。\n\n记忆的三大支柱包括：**状态**（了解当前情境，掌握正在发生的事情）、**持久性**（跨会话保留知识，确保信息不因对话结束而丢失）、**选择性**（判断哪些信息值得记住，哪些可以忽略）[探寻AI Agent 中隐秘的角落：记忆(Memory) - 定义、价值与实践](https://developer.volcengine.com/articles/7540134113190412324)。这三者共同赋予AI一种前所未有的能力——连续性。\n\n### 1.3 报告研究范围与目标\n\n本报告旨在系统梳理Agent记忆模块的技术演进方向，深入分析短期与长期记忆的实现机制、各类技术路线的差异与取舍，并从业务价值、技术成熟度和工程可行性三个维度，为产品团队提供明确的技术选型建议和投入优先级排序。\n\n---\n\n## 二、Agent记忆模块的分类体系与核心架构\n\n### 2.1 从认知科学到工程实践的记忆分层\n\n人类记忆遵循从感觉记忆到短期记忆再到长期记忆的一般性进程[A survey on large language model based autonomous agents](https://link.springer.com/content/pdf/10.1007/s11704-024-40231-1.pdf)。当设计Agent记忆结构时，研究者从中汲取灵感：\n\n**短期记忆（Working Memory）**：指AI正在进行的对话、脑中即时活跃的上下文。它容量有限（通常为数千tokens），但访问速度极快[从理论到落地：分层记忆架构在AI Agent中的应用实践](https://blog.csdn.net/whitehat_zhou/article/details/150269603)。例如，在MemGPT架构中，短期工作上下文由系统指令、工作上下文和FIFO队列组成[9.4k Star！MemGPT：伯克利大学最新开源、将LLM作为操作系统、无限上下文记忆、服务化部署自定义Agent](http://www.wehelpwin.com/m_article/5363)。\n\n**中期记忆（Episodic Memory）**：指最近读完的书籍核心内容或近期发生的重大事件。它们比短期记忆更持久，但不如长期记忆根深蒂固[从理论到落地：分层记忆架构在AI Agent中的应用实践](https://blog.csdn.net/whitehat_zhou/article/details/150269603)。这类记忆通常通过向量数据库或结构化存储实现语义检索。\n\n**长期记忆（Long-term Memory）**：指用户的个人经历、学到的技能、世界观以及那些已掌握的知识。它容量近乎无限，但检索可能需要更长时间[从理论到落地：分层记忆架构在AI Agent中的应用实践](https://blog.csdn.net/whitehat_zhou/article/details/150269603)。长期记忆是Agent个性化和持续学习的基础。\n\n### 2.2 三种记忆形式的技术实现\n\n根据2025年最新综述研究，Agent记忆可从形式（Forms）、功能（Functions）和动态（Dynamics）三个正交维度进行重构[2025年Memory最全综述！AI Agent记忆统一分类体系](https://zhuanlan.zhihu.com/p/1985435669187825983)：\n\n**Token级记忆（Token-level Memory）**：信息被处理在模型的上下文窗口内。这是最直接的记忆方式，包括滑动窗口注意力（Sliding Window Attention）和分块处理（Chunking）等技术[EPISODIC MEMORIES GENERATION AND EVALUATION BENCHMARK FOR LARGE LANGUAGE MODELS](https://arxiv.org/pdf/2501.13121v1.pdf)。优点是实现简单、无额外存储开销；缺点是在长序列上存在\"上下文腐烂\"问题[AI Agent 性能优化：核心策略与实战技巧（超详细）从零基础入门到精通！](https://m.blog.csdn.net/xiaoganbuaiuk/article/details/154012804)。\n\n**参数级记忆（Parameter-level Memory）**：通过微调将知识直接编码到模型参数中。这种方式使信息成为模型\"先天具有的知识\"，可在任意上下文中激活。优点是检索速度快、无需额外存储；缺点是训练成本高昂、难以增量更新[EPISODIC MEMORIES GENERATION AND EVALUATION BENCHMARK FOR LARGE LANGUAGE MODELS](https://arxiv.org/pdf/2501.13121v1.pdf)。\n\n**潜在级记忆（Latent-level Memory）**：通过检索增强生成（RAG）等外部存储机制实现。信息以向量或结构化形式存储在外存，需要时通过语义检索召回[EPISODIC MEMORIES GENERATION AND EVALUATION BENCHMARK FOR LARGE LANGUAGE MODELS](https://arxiv.org/pdf/2501.13121v1.pdf)。这是当前商业化产品中最广泛采用的方式。\n\n### 2.3 记忆模块在Agent架构中的位置\n\n典型的AI Agent包括LLM用于推理和生成答案、策略或规划模块（如ReAct或AutoGPT风格）、工具或API访问权限以及检索器用于获取文档或历史数据[探寻AI Agent 中隐秘的角落：记忆(Memory) - 定义、价值与实践](https://developer.volcengine.com/articles/7540134113190412324)。然而，这些组件都无法记住昨天发生了什么——它们没有内部状态，也没有随时间进化的理解力。\n\n加入记忆后，AI Agent的架构会发生质的改变：记忆层成为串联其他模块的\"数据中枢\"[大模型进阶之路：AI Agent记忆能力构建技术详解（值得收藏）](https://blog.csdn.net/xxue345678/article/details/150983939)。Agent架构的四大核心模块（感知、决策、记忆、行动）逐渐清晰，其中记忆模块（Brain-Memory & Knowledge）存储Agent的知识、用户偏好和历史交互记录，是决策与行动的\"数据基础\"[大模型进阶之路：AI Agent记忆能力构建技术详解（值得收藏）](https://blog.csdn.net/xxue345678/article/details/150983939)。\n\n---\n\n## 三、短期记忆的技术实现与优化策略\n\n### 3.1 上下文窗口的根本限制\n\n大语言模型在单次推理中能够接收和处理的Token数量上限被称为上下文窗口（Context Window），它直接决定了模型的\"短期记忆容量\"[LLM面试50问精解：从基础到进阶，掌握大语言模型核心知识，助你面试一战成名！大模型面试](https://m.blog.csdn.net/qkh1234567/article/details/155600143)。当前主流模型的上下文长度从4K token（早期GPT系列）发展到128K token（GPT-4o）甚至更长。然而，上下文窗口的扩大存在明显的\"性能代价\"：\n\n由于Transformer的自注意力机制需要计算序列中任意两个位置之间的相关性，其计算复杂度为O(n²)，内存复杂度同样为O(n²)[Hi，我是 OceanBase PowerMem，了解一下？](http://cdn.modb.pro/db/1994227212866904064)。这意味着当上下文窗口从4K增加到128K（32倍）时，理论上的计算量和内存需求的增长不是32倍，而是1024倍[Hi，我是 OceanBase PowerMem，了解一下？](http://cdn.modb.pro/db/1994227212866904064)。\n\n更重要的是，研究发现有效上下文长度远低于最大支持长度[LLM的长期记忆系统-成为你的个性化聊天陪伴](https://zhuanlan.zhihu.com/p/1917652340346946388)。已有研究表明，LLM的有效上下文长度在达到约2048 tokens后就开始显著下降，模型精度会随着上下文长度的增加而显著下降[LLM的长期记忆系统-成为你的个性化聊天陪伴](https://zhuanlan.zhihu.com/p/1917652340346946388)。\n\n### 3.2 KV Cache优化技术\n\nKV Cache（Key-Value Cache）的有效利用对提高LLM的运行效率至关重要[AI Agent优化技术深度解析：从Prompt到架构的全面指南（珍藏版）](https://m.blog.csdn.net/2401_85375186/article/details/155227353)。当前KV Cache优化可分为三个层次：\n\n**Token级策略**：包括键值缓存选择、预算分配、合并、量化和低秩分解[【文献阅读】A Survey on Large Language Model Acceleration based on KV Cache Management](https://m.blog.csdn.net/Toky_min/article/details/146019523)。代表性技术包括：\n- **SnapKV**：无需微调即可有效最小化KV缓存大小。在处理16K令牌输入时，实现一致的解码速度，与基线相比生成速度提高3.6倍，内存效率提高8.2倍[SnapKV: LLM在生成内容之前就知道您在寻找什么](https://blog.csdn.net/qq_36931982/article/details/139118015)\n- **LM-Infinite**和**StreamingLLM**：保留初始token和近期token以实现无限长度上下文\n- **H₂O**：基于注意力分数选择重要token进行缓存\n\n**模型级优化**：通过架构创新和注意力机制增强键值重用[【文献阅读】A Survey on Large Language Model Acceleration based on KV Cache Management](https://m.blog.csdn.net/Toky_min/article/details/146019523)。例如Layer-Condensed KV Cache通过减少KV Cache层数而非保留所有层来降低内存消耗[Layer-Condensed KV Cache for Efficient Inference of Large Language Models](https://arxiv.org/pdf/2405.10637v2.pdf)。\n\n**系统级方法**：解决内存管理、调度和硬件感知设计等问题[【文献阅读】A Survey on Large Language Model Acceleration based on KV Cache Management](https://m.blog.csdn.net/Toky_min/article/details/146019523)。NVIDIA Dynamo框架就采用了LLM-aware路由器来跟踪GPU集群中KV Cache的位置，当请求到达时计算新请求与缓存KV块之间的重叠度，将流量导向能最大化缓存复用的GPU[NVIDIA Dynamo Addresses Multi-Node LLM Inference Challenges](https://www.infoq.com/news/2025/12/nvidia-dynamo-kubernetes/?utm_campaign=infoq_content&utm_source=infoq&utm_medium=feed&utm_term=global)。\n\n### 3.3 短期记忆的工程实现方案\n\n**LangChain的记忆实现**提供了多种短期记忆方案\\cite{web_bd759108}：\n- **ConversationBufferMemory**：存储所有历史消息\n- **ConversationBufferWindowMemory**：仅保留最近N轮对话（如k=3表示保留最近3轮）\n- **ConversationSummaryMemory**：对对话历史生成摘要以节省Token\n\n这种分级设计体现了\"选择性\"原则——判断哪些信息值得记住，哪些可以忽略[探寻AI Agent 中隐秘的角落：记忆(Memory) - 定义、价值与实践](https://developer.volcengine.com/articles/7540134113190412324)。\n\n---\n\n## 四、长期记忆的技术路线全景\n\n### 4.1 RAG与Memory的区别与联系\n\nRAG（检索增强生成）是OpenAI、Meta等公司提出的一种框架，用来增强语言模型的知识能力[让大模型“记住”更多：RAG与长期记忆](https://m.blog.csdn.net/2401_82452722/article/details/148744627)。然而，需要明确区分RAG与真正的Agent Memory：\n\nRAG主要用于知识问答场景，在这些场景下知识点之间相对独立，一个用户问题往往只需涉及一两个知识点[LLM的长期记忆系统-成为你的个性化聊天陪伴](https://zhuanlan.zhihu.com/p/1917652340346946388)。企业级RAG系统通常会通过人工干预优化chunk划分，因此能够更好地支持知识类检索。\n\n但在\"个人记忆\"这一场景下存在独特挑战：首先很难依赖人工完成chunk划分；其次用户的提问往往需要LLM进行推理、联想才能找到真正相关的记忆[LLM的长期记忆系统-成为你的个性化聊天陪伴](https://zhuanlan.zhihu.com/p/1917652340346946388)。当前RAG技术将文本按chunk进行切分，对每个chunk独立生成特征向量，在检索阶段query也是独立与每个chunk进行相似度计算——这种做法不仅容易造成信息\"割裂\"，也难以保证检索到的是上下文完整的语义单元[LLM的长期记忆系统-成为你的个性化聊天陪伴](https://zhuanlan.zhihu.com/p/1917652340346946388)。\n\n### 4.2 主流长期记忆框架对比\n\n#### 4.2.1 MemGPT：操作系统灵感的虚拟内存管理\n\nMemGPT是由UC Berkeley研究团队开发的记忆增强LLM框架，其核心创新是借鉴操作系统内存管理的概念[Meet MemGPT: The Future of Memory-Augmented AI - Medium](https://medium.com/@harshavasana1/meet-memgpt-the-future-of-memory-augmented-ai-9b547bbd3879)：\n\n**主上下文（Main Context）**：类比RAM的短期内存，用于即时处理。包含系统指令、工作上下文和FIFO队列[9.4k Star！MemGPT：伯克利大学最新开源、将LLM作为操作系统、无限上下文记忆、服务化部署自定义Agent](http://www.wehelpwin.com/m_article/5363)。\n\n**外部上下文（External Context）**：类比硬盘的长期存储，用于存放不太常用的信息。MemGPT引入两个外部记忆模块——Recall Memory和Archival Memory，分别对应短期交互记录与长期知识存储[LLM的长期记忆系统-成为你的个性化聊天陪伴](https://zhuanlan.zhihu.com/p/1917652340346946388)。\n\nMemGPT通过函数在主上下文和外部上下文之间移动数据，LLM可以通过在其输出中生成特殊关键字参数来请求立即后续推理，以将函数调用链接在一起；函数链允许MemGPT执行多步骤检索来回答用户查询[9.4k Star！MemGPT：伯克利大学最新开源、将LLM作为操作系统、无限上下文记忆、服务化部署自定义Agent](http://www.wehelpwin.com/m_article/5363)。这种架构不仅允许处理更大的数据集和更长的对话，而且还提高了模型在扩展交互中保持一致性能力[9.4k Star！MemGPT：伯克利大学最新开源、将LLM作为操作系统、无限上下文记忆、服务化部署自定义Agent](http://www.wehelpwin.com/m_article/5363)。\n\nMemGPT的核心优势在于它证明了LLM可以被教会管理自己的内存[9.4k Star！MemGPT：伯克利大学最新开源、将LLM作为操作系统、无限上下文记忆、服务化部署自定义Agent](http://www.wehelpwin.com/m_article/5363)。其后商业化的Letta公司已获得1000万美元种子轮融资，在7000万美元估值下由Felicis Ventures领投[Letta, one of UC Berkeley's most anticipated AI startups, has just ...](https://techcrunch.com/2024/09/23/letta-one-of-uc-berkeleys-most-anticipated-ai-startups-has-just-come-out-of-stealth/)。\n\n#### 4.2.2 Mem0：轻量级实用主义路线\n\nMem0被官方定义为\"The Memory Layer for Personalized AI\"（个性化AI的记忆层）[Mem0 获2400万美元融资，开发AI内存层技术-搜索](https://cn.bing.com/search?q=Mem0%20%E8%8E%B72400%E4%B8%87%E7%BE%8E%E5%85%83%E8%9E%8D%E8%B5%84%EF%BC%8C%E5%BC%80%E5%8F%91AI%E5%86%85%E5%AD%98%E5%B1%82%E6%8A%80%E6%9C%AF&rdr=1&rdrig=7UV8W6XI2C2HFZFXE23T7R5222E2IX6F&first=1)。其核心设计理念包括：\n\n- **记忆是可搜索和可管理的**：通过自然语言索引+向量化混合检索\n- **以用为先**：轻量级实现适合实际部署\n- **结构化存储**：使用slot存储结构化长期信息（如角色设定、兴趣偏好）\n\nMem0在LOCOmo Benchmark中达到了业界开源Memory解决方案的SOTA水平[Hi，我是 OceanBase PowerMem，了解一下？](http://cdn.modb.pro/db/1994227212866904064)。其设计哲学是\"足够好而非完美\"——通过可恢复的压缩策略（如仅保留URL而非完整内容）在不永久丢失信息的情况下减少上下文开销[AI Agent 性能优化：核心策略与实战技巧（超详细）从零基础入门到精通！](https://m.blog.csdn.net/xiaoganbuaiuk/article/details/154012804)。\n\n#### 4.2.3 OpenMemory：分层记忆分解架构\n\nOpenMemory是一个自托管、模块化的AI记忆引擎，采用Hierarchical Memory Decomposition（HMD）架构[CaviraOSS/OpenMemory: Add long-term memory to any AI ...](https://github.com/CaviraOSS/OpenMemory)：\n\n- **单一规范节点每记忆**（无数据复制）\n- **多扇区嵌入**（情节、语义、程序性、情感、反思性）\n- **单一点位链接**（稀疏、生物启发式图）\n- **复合相似性检索**（扇区融合+激活传播）\n\n性能对比数据显示[CaviraOSS/OpenMemory: Add long-term memory to any AI ...](https://github.com/CaviraOSS/OpenMemory)：\n| 指标 | OpenMemory（自托管） | Zep（云服务） | Supermemory（SaaS） | Mem0 | 向量DB（平均） |\n|------|---------------------|---------------|--------------------|-------|----------------|\n| 成本 | 自托管 | 按API调用付费 | 按API调用付费 | 按API调用付费 | 按存储+查询付费 |\n\n#### 4.2.4 A-MEM：卡片盒笔记法驱动的记忆系统\n\nA-MEM是一种新型面向LLM Agent的智能记忆系统，遵循卡片盒笔记法（Zettelkasten）的基本原理[A-MEM智能记忆系统，让你的大模型拥有“学习”能力（收藏版）](https://blog.csdn.net/xiaoganbuaiuk/article/details/151215108)：\n\n当添加新记忆时，系统会生成包含多个结构化属性的综合笔记，包括上下文描述、关键词和标签。随后分析历史记忆以识别相关联系，并在存在有意义的相似性时建立链接——这一过程还支持记忆进化——随着新记忆的整合，它们会触发对现有历史记忆的上下文表示和属性更新[A-MEM智能记忆系统，让你的大模型拥有“学习”能力（收藏版）](https://blog.csdn.net/xiaoganbuaiuk/article/details/151215108)。\n\nA-MEM的成本效益分析显示：通过选择性top-k检索机制，每次记忆操作约需1200个tokens，与基线方法（LoCoMo和MemGPT为16900个tokens）相比，Token使用量减少了85%-93%[A-MEM智能记忆系统，让你的大模型拥有“学习”能力（收藏版）](https://blog.csdn.net/xiaoganbuaiuk/article/details/151215108)。\n\n### 4.3 多模态记忆的发展前沿\n\n随着多模态LLM的发展，记忆系统也开始向多模态方向演进：\n\n**M3-Agent**是字节跳动Seed团队提出的新型多模态代理框架，配备长期记忆能力。与人类类似，M3-Agent可以处理实时视觉和听觉输入来构建和更新其长期记忆[Research](https://seed.bytedance.com/en/public_papers)。\n\n**VideoAgent**探讨如何将基础模型（大型语言模型和视觉语言模型）与统一记忆机制协调起来解决视频理解问题。其构建结构化内存来存储视频的通用时间事件描述和以对象为中心的跟踪状态，在NExTQA上平均提升6.6%，在EgoSchema上平均提升26.0%[VideoAgent: A Memory-augmented Multimodal Agent for Video Understanding（翻译）](https://m.blog.csdn.net/weixin_56764022/article/details/139862630)。\n\n---\n\n## 五、OpenAI ChatGPT记忆功能解析：行业标杆案例\n\n### 5.1 发展历程与产品演进\n\nOpenAI的记忆功能经历了三个重要阶段：\n\n**第一阶段（2024年2月）**：小范围测试ChatGPT的记忆功能——记住用户在聊天中讨论过的事情，并避免重复信息。用户可以要求它记住特定的内容或让它自行获取详细信息。用得越多，ChatGPT的记忆力就会越好[ChatGPT要有记忆力了，OpenAI宣布小范围测试“记忆”功能！](http://m.nbd.com.cn/articles/2024-02-14/3245164.html)。\n\n**第二阶段（2025年4月）**：向所有ChatGPT Plus和Pro用户推出记忆提升功能。这次更新后，ChatGPT的记忆功能能够参考用户过去的所有聊天以提供更个性化的回复[OpenAI：ChatGPT的记忆提升功能今天起对所有Plus和Pro用户推出](https://finance.jrj.com.cn/2025/04/11072749477951.shtml)。例如，如果用户曾经提到喜欢泰国菜，下次问\"中午应该吃什么\"时，ChatGPT可能会考虑到这一点[刚刚！OpenAI发布ChatGPT记忆功能，秒变私人助理！](https://m.blog.csdn.net/Leinwin/article/details/148428926)。\n\n**第三阶段（2025年6月）**：为免费版ChatGPT推出轻量级记忆功能。该功能根据用户过往一段时间内的对话习惯、写作风格、提问方式进行个性化回答[OpenAI发布ChatGPT记忆功能](https://c.m.163.com/news/a/K16NIQ1D0511A6N9.html)。\n\n### 5.2 记忆功能的核心机制\n\nChatGPT的记忆功能通过两个设置来控制[刚刚！OpenAI发布ChatGPT记忆功能，秒变私人助理！](https://m.blog.csdn.net/Leinwin/article/details/148428926)：\n\n- **\"引用已保存的记忆\"**：用户明确要求ChatGPT记住的细节（如名字、喜欢的颜色或饮食偏好）\n- **\"引用聊天历史\"**：ChatGPT可以利用过去的对话信息来使未来的对话更有帮助\n\nOpenAI称用户可以随时关闭新功能，在设置菜单中找到个性化选项即可将记忆设置为关闭。用户也可以在设置中找到\"管理内存\"选项，在其中查看和删除特定的记忆内容或清除所有记忆[ChatGPT要有记忆力了，OpenAI宣布小范围测试“记忆”功能！](http://m.nbd.com.cn/articles/2024-02-14/3245164.html)。\n\n### 5.3 记忆功能的战略意义\n\n投资人Allie K. Miller在X平台上感叹：\"这相当于ChatGPT全天候在'偷听'——不管你有没有叫它记住，它都在默默收集\"[OpenAI最近推出的ChatGPT更新简直像给AI打了“记忆芯片”](https://blog.csdn.net/2301_79342058/article/details/147156180)。她还表示，在平台功能越来越趋同的今天，真正拉开差距的关键就是\"记忆+个性化\"——AI的记忆就是平台的护城河[OpenAI最近推出的ChatGPT更新简直像给AI打了“记忆芯片”](https://blog.csdn.net/2301_79342058/article/details/147156180)。\n\nOpenAI研究还表明，ChatGPT记住什么或不记住什么与人们如何体验ChatGPT的人格密切相关——许多Plus和Pro订阅者告诉OpenAI，更好的记忆是体验中最有价值的部分之一[超越“一刀切”：ChatGPT如何为8亿用户定制个性化体验| 图文 ...](https://www.zhihu.com/pin/1972634281496057035)。\n\n---\n\n## 六、评测基准与性能分析\n\n### 6.1 主要评测基准概述\n\n当前用于评估Agent Memory的主要数据集包括[Agent 又又失忆了！我来做一次记忆体检](https://juejin.cn/post/7564589641258139648)[Survey on Evaluation of LLM-based Agents](http://arxiv.org/pdf/2503.16416v1)：\n\n**LoCoMo Benchmark**：专门测试长上下文任务的记忆能力，在该基准中Mem0和Letta曾出现过评测分值分歧[2025 AI 记忆系统大横评：从插件到操作系统，谁在定义下一代Agent Infra？](https://m.toutiao.com/a7578151050135044643/)。\n\n**LongMemEval**：UCLA团队开发的系统性基准，更像是一个\"记忆体检表\"——从信息提取、跨会话推理、知识更新到拒答未知，一共五项指标[Agent 又又失忆了！我来做一次记忆体检](https://juejin.cn/post/7564589641258139648)。\n\n**Reflection-Bench**：上海人工智能实验室开发的认知心理学导向评测平台，围绕七个认知维度设计了354个任务——预测能力与决策能力、感知能力与记忆能力、反事实思维、信念更新、元反思能力等[AI大模型代理评测全攻略：从入门到精通，一篇就够了！](https://m.blog.csdn.net/m0_56255097/article/details/154832865)。\n\n这些评测指标超越了传统的自然语言处理性能指标——聚焦事实信息的存储与利用，通过准确性指标（基于历史信息生成响应的正确性）和召回率@5指标（前5条检索结果中相关信息的占比）来衡量[上下文工程最新综述A Survey of Context Engineering for ...](https://zhuanlan.zhihu.com/p/1956121336968647834)。\n\n### 6.2 性能数据对比\n\n根据MemOS横向评测数据，在实际任务中的表现差异显著[2025 AI 记忆系统大横评：从插件到操作系统，谁在定义下一代Agent Infra？](https://m.toutiao.com/a7578151050135044643/)：\n\n| 技术路线 | 上下文扩展能力 | 检索准确性 | 响应延迟 | 开发复杂度 | 成本效益 |\n|---------|--------------|-----------|---------|----------|---------|\n| 基础RAG | 中等 | 较低 | 较低 | 低 | 高 |\n| MemGPT | 高 | 中等 | 较高 | 高 | 中等 |\n| Mem0 | 高 | 较高 | 中等 | 中等 | 高 |\n| 向量数据库+自定义 | 可定制 | 高 | 中等 | 高 | 中等 |\n\n---\n\n## 七、工程实践案例与落地挑战\n\n### 7.1 成功案例分析\n\n来自硅谷一线AI创业者的数据显示：只有5%的AI Agent成功部署到生产环境[为什么只有5%的AI Agent落地成功？](https://wwww.huxiu.com/comment/4792610.html)。这5%的成功案例都有一个共同点——都采用了\"human-in-the-loop\"的设计：\n- 将AI定位为辅助工具而非自主决策者\n- 构建反馈循环让系统能从人类修正中学习\n- 让人类可以轻松验证和否决AI的输出\n\n这些成功案例在记忆设计上的共同特点是：精细调整模型的需求其实非常少见——一个设计完善的RAG系统通常就能满足需求。但大多数RAG系统的设计都太过初级——将所有内容编入索引导致检索信息过量反而迷惑模型；编入索引的内容过少导致模型缺乏有效信号；不加区分地混合结构化与非结构化数据会破坏嵌入向量的语义[为什么只有5%的AI Agent落地成功？](https://wwww.huxiu.com/comment/4792610.html)。\n\n### 7.2 记忆功能的产品化架构\n\n亚马逊云科技发布的Amazon Bedrock AgentCore为开发者打通了AI Agents从概念验证到生产部署的关键环节[跨越“演示”到“生产”鸿沟，亚马逊云科技开启AI Agents新纪元](http://app.myzaker.com/news/article.php?pk=6878d44a8e9f091282757f81)。其核心能力包括：\n\n- **会话管理**：处理多用户并发场景下的状态隔离\n- **身份权限控制**：确保AI访问敏感数据时的身份与权限可控\n- **记忆系统**：支持分级存储和版本控制\n- **可观测性机制**：监控Agent行为并进行调试\n\n该平台的一个关键洞察是：大多数创始人以为自己在打造AI产品，实际上构建的是上下文选择系统[为什么只有5%的AI Agent落地成功？](https://wwww.huxiu.com/comment/4792610.html)。真正的工程工作应该得到应有的重视——精细调整模型的需求其实非常少见。\n\n### 7.3 记忆层级的设计原则\n\n产品化设计中应该考虑多层次的记忆架构[为什么只有5%的AI Agent落地成功？](https://wwww.huxiu.com/comment/4792610.html)：\n\n**用户级**：个人偏好设置（如图表类型、风格、写作语气）\n**团队级**：高频查询、仪表盘、标准操作手册（runbooks）\n**企业级**：知识库、政策文档、历史决策记录\n\n这种分层设计既能满足个性化需求又能保护敏感信息不被泄露给其他用户。\n\n---\n\n## 八、技术路线对比与产品化建议\n\n### 8.1 不同技术路线的核心差异矩阵\n\n基于技术成熟度、落地难度、成本效益和业务价值四个维度进行综合评估：\n\n| 技术路线 | 技术成熟度 | 落地难度 | 成本效益 | 业务价值 |\n|---------|----------|---------|---------|---------|\n| Token级短期记忆 | ★★★★★ | ★★☆☆☆ | ★★★★★ | ★★★☆☆ |\n| 基础RAG方案 | ★★★★☆ | ★★★☆☆ | ★★★★☆ | ★★★★☆ |\n| MemGPT虚拟内存 | ★★★☆☆ | ★★★★★ | ★★☆☆☆ | ★★★★★ |\n| Mem0轻量级方案 | ★★★★☆ | ★★★☆☆ | ★★★★★ | ★★★★☆ |\n| A-MEM卡片盒方案 | ★★★☆☆ | ★★★★☆ | ★★★★☆ | ★★★★★ |\n\n### 8.2 各类技术路线优劣势详解\n\n#### 基础RAG方案（推荐指数：⭐⭐⭐⭐⭐）\n\n**优势**：\n- 技术成熟度高，生态系统完善\n- 成本可控，可使用开源组件构建\n- 易于理解和调试\n- 可快速上线验证业务价值\n\n**劣势**：\n- 检索准确性受chunk划分影响大\n- 难以支持复杂推理关联\n- 对非结构化个人记忆支持有限\n\n**适用场景**：知识问答系统、文档检索助手、企业知识库查询\n\n#### Mem0轻量级方案（推荐指数：⭐⭐⭐⭐⭐）\n\n**优势**：\n- 实现简单，有现成SDK\n- 性能达到开源方案SOTA水平\n- 资源占用低适合中小规模部署\n- 社区活跃支持良好\n\n**劣势**：\n- 功能相对基础，缺乏深度定制能力\n- 对复杂多模态场景支持有限\n\n**适用场景**：个性化聊天助手、客服机器人升级、中小型企业AI应用\n\n#### MemGPT虚拟内存方案（推荐指数：⭐⭐⭐★☆）\n\n**优势**：\n- 理论架构先进，支持无限扩展\n- 内存利用率高\n- 可学习管理自身内存\n\n**劣势**：\n- 实现复杂度高\n- 运维难度大\n- 对开发团队要求高\n- 生产环境稳定性待验证\n\n**适用场景**：研发型团队探索前沿技术、学术研究、复杂任务型Agent\n\n#### 定制化向量数据库方案（推荐指数：⭐⭐⭐★☆）\n\n**优势**：\n- 可高度定制化\n- 支持复杂查询逻辑\n- 数据完全可控（隐私合规）\n\n**劣势**：\n- 开发周期长\n- 运维成本高\n- 需要专业的数据工程团队\n\n**适用场景**：大型企业私有化部署、高隐私要求行业（医疗、金融）\n\n### 8.3 投入优先级建议\n\n基于您的产品定位和团队能力，建议采取以下投入策略：\n\n#### 第一优先级（立即投入）\n\n**基础RAG+短期记忆组合方案**\n\n理由：\n1. 技术风险最低——有成熟的开源框架和丰富的社区支持\n2. 成本可控——可使用开源向量数据库（如FAISS、Milvus）\n3. 快速验证价值——几周内即可上线原型\n4. 底座稳固——为后续升级奠定基础\n\n实施方案建议：\n```\n短期记忆层 → LangChain/LLamaIndex内置memory模块\n长期记忆层 → FAISS/Pinecone向量数据库 + 自定义Chunk策略\n检索增强 → 多模态Embedding + Top-k筛选 + 人工反馈闭环\n```\n\n#### 第二优先级（3个月内）\n\n**引入Mem0或类似轻量级方案**\n\n理由：\n1. 在第一优先级方案验证价值后逐步升级\n2. 成本效益比高——开源方案可快速迭代\n3. 技术风险可控——有成熟产品参考\n\n#### 第三优先级（6-12个月）\n\n**探索差异化方案**\n\n根据业务场景可选择：\n- **如果是知识密集型应用**：深入定制向量数据库方案\n- **如果是任务型Agent**：探索MemGPT虚拟内存方案\n- **如果是多模态应用**：关注M3-Agent等前沿框架演进\n\n---\n\n## 九、未来演进方向与趋势判断\n\n### 9.1 短期趋势（2025-2026）\n\n根据奥尔特曼的说法：\"这是2026年要考虑的事\"[超级Agent重要拼图？奥尔特曼点名“AI记忆” 存储环节迎来新叙事](https://finance.sina.com.cn/stock/t/2025-12-22/doc-inhcsmfc4526733.shtml)——真正成熟的长期记忆系统将在明年成为焦点。预计的发展方向包括：\n\n**KV Cache优化技术成熟化**\n目前KV Cache优化仍在快速演进中。SnapKV已在单个A100-80GB GPU上实现处理多达38万个上下文令牌的能力[SnapKV: LLM在生成内容之前就知道您在寻找什么](https://blog.csdn.net/qq_36931982/article/details/139118015)。随着技术成熟，这一能力将成为标准配置而非差异化优势。\n\n**多模态记忆标准化**\nM3-Agent等多模态记忆框架正在快速发展[Research](https://seed.bytedance.com/en/public_papers)。预计到2026年将出现标准化的多模态记忆接口和评估基准。\n\n### 9.2 中期趋势（2026-2027）\n\n**操作系统级别集成**\nAgent的记忆系统将从独立组件演进为操作系统级别的原生能力——类似MemGPT的理念但更加成熟稳定。届时可能出现专门针对AI应用优化的操作系统发行版或云服务层。\n\n**个性化与隐私平衡**\n随着监管加强和技术进步，个性化记忆将需要更严格的隐私保护机制。联邦学习、差分隐私等技术将与记忆系统深度整合。\n\n### 9.3 长期展望（2027+）\n\n真正的类人记忆能力将成为区分优秀AI产品与普通工具的关键标志。正如行业共识所言：\"在AI都拥有相似模型和工具的未来，记忆将成为决定胜负的关键\"[探寻AI Agent 中隐秘的角落：记忆(Memory) - 定义、价值与实践](https://developer.volcengine.com/articles/7540134113190412324)。\n\n---\n\n## 十、结论与行动建议\n\n### 10.1 核心结论\n\n通过对当前Agent Memory模块的技术全景调研，可以得出以下核心结论：\n\n第一，记忆模块已从\"锦上添花\"转变为\"不可或缺\"——它是Agent从工具到伙伴的关键分水岭[大模型进阶之路：AI Agent记忆能力构建技术详解（值得收藏）](https://blog.csdn.net/xxue345678/article/details/150983939)[超级Agent重要拼图？奥尔特曼点名“AI记忆” 存储环节迎来新叙事](https://finance.sina.com.cn/stock/t/2025-12-22/doc-inhcsmfc4526733.shtml)。\n\n第二，在技术成熟度和工程可行性方面，基础RAG方案配合短期记忆是最务实的起点；Mem0等轻量级方案提供了良好的平衡点；而MemGPT等前沿方案更适合有充足技术储备的研发型团队。\n\n第三，从成本效益角度分析，A-MEM等优化方案可将Token消耗降低85%-93%[A-MEM智能记忆系统，让你的大模型拥有“学习”能力（收藏版）](https://blog.csdn.net/xiaoganbuaiuk/article/details/151215108)——这说明优化空间巨大但仍需谨慎选择合适的技术路线。\n\n第四，在OpenAI的带领下，个性化记忆已成为行业共识方向[ChatGPT要有记忆力了，OpenAI宣布小范围测试“记忆”功能！](http://m.nbd.com.cn/articles/2024-02-14/3245164.html)[超级Agent重要拼图？奥尔特曼点名“AI记忆” 存储环节迎来新叙事](http://m.cls.cn/detail/2236511)——不跟进将面临竞争劣势风险。\n\n### 10.2 行动建议清单\n\n针对您的产品规划，建议按以下步骤推进：\n\n✅ **Week 1-2**: 明确产品对记忆的具体需求——是知识检索还是个性化对话？需要短期还是长期？对准确性和延迟的要求是多少？\n\n✅ **Week 3**: 基于需求选择技术路线——优先评估LangChain/LlamaIndex+FAISS的基础组合方案\n\n✅ **Week 4**: 构建最小可行原型——用两周时间验证核心功能是否满足业务需求\n\n✅ **Week 5**: 设置评估指标——参考LoCoMo或LongMemEval建立内部评测体系\n\n✅ **Month 2**: 根据原型反馈决定是否引入Mem0或其他轻量级方案进行升级\n\n✅ **Month 3+**: 持续优化并关注行业演进——特别是OpenAI的记忆功能迭代和其他竞品动态\n\n最终建议采取\"敏捷验证、快速迭代\"策略：不要追求一步到位的理想方案，在验证业务价值的基础上逐步升级技术栈。正如一位资深AI产品经理所言：\"大多数创始人以为自己在打造AI产品，但实际上他们构建的是上下文选择系统\"[为什么只有5%的AI Agent落地成功？](https://wwww.huxiu.com/comment/4792610.html)——聚焦用户价值而非技术炫技才是成功的关键。\n
\n\n\n## Quick Start\n\nYou can get beta access to the model API through StepFun Open Platform. \n\n### Requirements\n\n- Python >= 3.10\n- Node.js >= 18 (for frontend)\n- npm or yarn\n\n### 1. Environment Setup\n\n**Install dependencies (backend & frontend)**\n\n```bash\n# Using uv (recommended)\nuv sync\nsource .venv/bin/activate\n\n# Or using pip\npip install -e .\n```\n\n```bash\ncd cortex-ui\n\n# Using npm\nnpm install\n\n# Or using yarn\nyarn install\n```\n\n**Configure environment variables**\n\nGet your StepFun API key(s) from [StepFun Open Platform](https://platform.stepfun.com/interface-key), the StepFun API key is for both model and search.\n\nCreate a `.env` file in the repo root:\n\n```bash\n# Model provider\nMODEL_PROVIDER=stepfun\n# Model name for Deep Research Agent (optional: step-dr-1, step-3.5-flash; default: step-dr-1)\nMODEL_NAME=step-dr-1\n# Model API base URL (StepFun)\nMODEL_BASE=https://api.stepfun.com\n# StepFun model API key\nSTEP_MODEL_API_KEY=your-model-api-key\n# Search API base URL (StepFun search service)\nSTEP_SEARCH_API_BASE=https://api.stepfun.com\n# StepFun search API key\nSTEP_SEARCH_API_KEY=your-search-api-key\n```\n\n**Recommended System Prompt**\n\nWe recommend using the system prompt in [prompt.py](scripts/configs/prompt.py) to ensure optimal performance.\n\n### 2. Run with Demo UI\n\n**Start backend service**\n\n```bash\npython -m demo.server [OPTIONS]\n```\n\nOptions:\n- `--port PORT` Server port (default: `8001`)\n\nThe service runs on `http://localhost:8001` by default.\n\n**Main endpoints:**\n- `GET /agents` - Get list of all available Agents\n- `WebSocket /ws` - Real-time communication endpoint\n\n**Start frontend service**\n\n```bash\ncd cortex-ui\n\n# Development mode\nnpm run dev\n# Or\nyarn dev\n```\n\nThe frontend runs on `http://localhost:3000` by default and automatically proxies API requests to the backend service.\n\n### 3. Run with Offline Runner\n\nUse `python -m scripts.runner` to run DeepResearchAgent on ad-hoc tasks without the UI.\n\n- **Direct args**: pass a single prompt or a tasks file.\n ```bash\n python -m scripts.runner \\\n --task \"列出最近的 AI 安全新闻\" \\\n --output-dir scripts/results\n ```\n- **Config file**: provide defaults via YAML .\n ```bash\n python -m scripts.runner --config scripts/configs/runner_example.yaml\n ```\n\nInputs:\n- `--task` / `--task-id`: one-off prompt and optional id.\n- `--tasks-file`: JSON task list (relative to project root if not absolute), e.g. `scripts/configs/tasks.example.json`.\n- Optional flags: `--output-dir` (default `scripts/results`), `--mode` (`multi` default), `--no-stream`, `--request-timeout`, `--context-upper-limit`, `--context-lower-limit`, `--overwrite`.\n\nOutputs:\n- One JSON trace per task in `output_dir`, containing the final answer, full message/tool events, metadata, and status.\n\n\n## Our paper\n\nWe have made our technical report available. [Read](https://arxiv.org/pdf/2512.20491)\n\n## Contact Us\n\nYou can join our group chat to get updates on your beta API application status and the latest project developments.\n
\n \"WeChat\n \"Feishu\n
\n\n## License\n\nThe code in the repository is licensed under [Apache 2.0](LICENSE) License.\n\n## Citation\n\n```\n@misc{hu2025stepdeepresearchtechnicalreport,\n title={Step-DeepResearch Technical Report}, \n author={Chen Hu and Haikuo Du and Heng Wang and Lin Lin and Mingrui Chen and Peng Liu and Ruihang Miao and Tianchi Yue and Wang You and Wei Ji and Wei Yuan and Wenjin Deng and Xiaojian Yuan and Xiaoyun Zhang and Xiangyu Liu and Xikai Liu and Yanming Xu and Yicheng Cao and Yifei Zhang and Yongyao Wang and Yubo Shu and Yurong Zhang and Yuxiang Zhang and Zheng Gong and Zhichao Chang and Binyan Li and Dan Ma and Furong Jia and Hongyuan Wang and Jiayu Liu and Jing Bai and Junlan Liu and Manjiao Liu and Na Wang and Qiuping Wu and Qinxin Du and Shiwei Li and Wen Sun and Yifeng Gong and Yonglin Chen and Yuling Zhao and Yuxuan Lin and Ziqi Ren and Zixuan Wang and Aihu Zhang and Brian Li and Buyun Ma and Kang An and Li Xie and Mingliang Li and Pan Li and Shidong Yang and Xi Chen and Xiaojia Liu and Yuchu Luo and Yuan Song and YuanHao Ding and Yuanwei Liang and Zexi Li and Zhaoning Zhang and Zixin Zhang and Binxing Jiao and Daxin Jiang and Jiansheng Chen and Jing Li and Xiangyu Zhang and Yibo Zhu},\n year={2025},\n eprint={2512.20491},\n archivePrefix={arXiv},\n primaryClass={cs.CL},\n url={https://arxiv.org/abs/2512.20491}, \n}\n```\n\n## Star History\n\n[![Star History Chart](https://api.star-history.com/svg?repos=stepfun-ai/StepDeepResearch&type=date&legend=top-left)](https://www.star-history.com/#stepfun-ai/StepDeepResearch&type=date&legend=top-left)\n" }, { "path": "agentkit/trace/__init__.py", "content": "from .builder import EventBuilder, FunctionSpanBuilder, HTTPSpanBuilder, SpanBuilder\nfrom .context import (\n SpanContext,\n create_span,\n get_current_context,\n get_default_app_name,\n get_default_tracer,\n record_event,\n set_default_app_name,\n set_default_tracer,\n start_trace,\n trace_function,\n)\nfrom .local_tracer import LocalStorageTracer\nfrom .remote_tracer import HybridTracer, RemoteTracer\nfrom .span import (\n DeltaEventPayload,\n Error,\n Event,\n EventType,\n FunctionSpanPayload,\n HTTPSpanPayload,\n LLMSpanPayload,\n OtherEventPayload,\n OtherSpanPayload,\n Span,\n SpanType,\n ToolSpanPayload,\n)\nfrom .tracer import Tracer\n\n__all__ = [\n # Core types\n \"Span\",\n \"Event\",\n \"SpanType\",\n \"EventType\",\n \"LLMSpanPayload\",\n \"ToolSpanPayload\",\n \"FunctionSpanPayload\",\n \"HTTPSpanPayload\",\n \"OtherSpanPayload\",\n \"DeltaEventPayload\",\n \"OtherEventPayload\",\n \"Error\",\n \"Tracer\",\n \"LocalStorageTracer\",\n \"RemoteTracer\",\n \"HybridTracer\",\n # Context management\n \"SpanContext\",\n \"get_current_context\",\n \"get_default_tracer\",\n \"get_default_app_name\",\n \"set_default_tracer\",\n \"set_default_app_name\",\n \"start_trace\",\n \"create_span\",\n \"record_event\",\n \"trace_function\",\n # Builders\n \"SpanBuilder\",\n \"FunctionSpanBuilder\",\n \"HTTPSpanBuilder\",\n \"EventBuilder\",\n]\n" }, { "path": "agentkit/trace/builder.py", "content": "\"\"\"\nBuilder pattern for constructing Span and Event objects.\n\"\"\"\n\nfrom datetime import datetime\nfrom typing import Any, Optional\n\nfrom ulid import ULID\n\nfrom .default import get_default\nfrom .span import (\n Event,\n EventType,\n FunctionSpanPayload,\n HTTPSpanPayload,\n OtherEventPayload,\n Span,\n SpanType,\n)\n\n\nclass SpanBuilder:\n \"\"\"\n Span builder providing a fluent API to construct complex Spans.\n\n Usage:\n span = (SpanBuilder()\n .with_name(\"my_operation\")\n .with_app_name(\"my_app\")\n .with_type(SpanType.FUNCTION)\n .with_tag(\"user\", \"alice\")\n .with_parent(parent_span)\n .build())\n \"\"\"\n\n def __init__(self, trace_id: Optional[str] = None, app_name: Optional[str] = None):\n self._trace_id = trace_id or str(ULID())\n self._app_name = app_name or get_default(\"app_name\")\n self._id = str(ULID())\n self._name = \"\"\n self._start_time = datetime.now()\n self._end_time = None\n self._tags: dict[str, str] = {}\n self._payload: Optional[Any] = None\n self._parent_id: Optional[str] = None\n\n def with_id(self, span_id: str) -> \"SpanBuilder\":\n \"\"\"Set span ID.\"\"\"\n self._id = span_id\n return self\n\n def with_name(self, name: str) -> \"SpanBuilder\":\n \"\"\"Set span name.\"\"\"\n self._name = name\n return self\n\n def with_trace_id(self, trace_id: str) -> \"SpanBuilder\":\n \"\"\"Set trace ID.\"\"\"\n self._trace_id = trace_id\n return self\n\n def with_start_time(self, start_time: datetime) -> \"SpanBuilder\":\n \"\"\"Set start time.\"\"\"\n self._start_time = start_time\n return self\n\n def with_end_time(self, end_time: datetime) -> \"SpanBuilder\":\n \"\"\"Set end time.\"\"\"\n self._end_time = end_time\n return self\n\n def with_tag(self, key: str, value: str) -> \"SpanBuilder\":\n \"\"\"Add a tag.\"\"\"\n self._tags[key] = value\n return self\n\n def with_tags(self, tags: dict[str, str]) -> \"SpanBuilder\":\n \"\"\"Add multiple tags.\"\"\"\n self._tags.update(tags)\n return self\n\n def with_payload(self, payload: Any) -> \"SpanBuilder\":\n \"\"\"Set payload.\"\"\"\n self._payload = payload\n return self\n\n def with_parent(self, parent: Span) -> \"SpanBuilder\":\n \"\"\"Set parent span.\"\"\"\n self._parent_id = parent.id\n # Automatically inherit trace_id from parent span\n if parent:\n self._trace_id = parent.trace_id\n return self\n\n def with_parent_id(self, parent_id: str) -> \"SpanBuilder\":\n \"\"\"Set parent span ID.\"\"\"\n self._parent_id = parent_id\n return self\n\n def with_app_name(self, app_name: str) -> \"SpanBuilder\":\n \"\"\"Set application name.\"\"\"\n self._app_name = app_name\n return self\n\n def build(self) -> Span:\n \"\"\"Build Span object.\"\"\"\n span = Span(\n id=self._id,\n name=self._name,\n trace_id=self._trace_id,\n app_name=self._app_name,\n start_time=self._start_time,\n end_time=self._end_time,\n tags=self._tags,\n payload=self._payload,\n parent_id=self._parent_id,\n )\n return span\n\n\nclass FunctionSpanBuilder(SpanBuilder):\n \"\"\"\n Specialized builder for Function Spans.\n\n Usage:\n span = (FunctionSpanBuilder()\n .with_name(\"calculate\")\n .with_function_name(\"calculate_total\")\n .with_arguments({\"x\": 1, \"y\": 2})\n .with_return_value(3)\n .build())\n \"\"\"\n\n def __init__(self, trace_id: Optional[str] = None, app_name: Optional[str] = None):\n super().__init__(trace_id, app_name)\n self._function_name = \"\"\n self._arguments: dict[str, Any] = {}\n self._ret: Any = None\n self._error = None\n\n def with_function_name(self, name: str) -> \"FunctionSpanBuilder\":\n \"\"\"Set function name.\"\"\"\n self._function_name = name\n if not self._name:\n self._name = name\n return self\n\n def with_arguments(self, arguments: dict[str, Any]) -> \"FunctionSpanBuilder\":\n \"\"\"Set function arguments.\"\"\"\n self._arguments = arguments\n return self\n\n def with_return_value(self, ret: Any) -> \"FunctionSpanBuilder\":\n \"\"\"Set return value.\"\"\"\n self._ret = ret\n return self\n\n def with_error(self, code: int, message: str) -> \"FunctionSpanBuilder\":\n \"\"\"Set error information.\"\"\"\n self._error = {\"code\": code, \"message\": message}\n return self\n\n def build(self) -> Span:\n \"\"\"Build Function Span.\"\"\"\n payload = FunctionSpanPayload(\n type=SpanType.FUNCTION,\n name=self._function_name,\n arguments=self._arguments,\n ret=self._ret,\n error=self._error,\n )\n self._payload = payload\n return super().build()\n\n\nclass HTTPSpanBuilder(SpanBuilder):\n \"\"\"\n Specialized builder for HTTP Spans.\n\n Usage:\n span = (HTTPSpanBuilder()\n .with_url(\"https://api.example.com/users\")\n .with_method(\"GET\")\n .with_header(\"Authorization\", \"Bearer token\")\n .with_response('{\"users\": []}')\n .build())\n \"\"\"\n\n def __init__(self, trace_id: Optional[str] = None, app_name: Optional[str] = None):\n super().__init__(trace_id, app_name)\n self._url = \"\"\n self._method = \"GET\"\n self._headers: dict[str, list[str]] = {}\n self._body: Optional[str | bytes] = None\n self._response: Optional[str | bytes] = None\n self._error = None\n\n def with_url(self, url: str) -> \"HTTPSpanBuilder\":\n \"\"\"Set URL.\"\"\"\n self._url = url\n if not self._name:\n self._name = f\"{self._method} {url}\"\n return self\n\n def with_method(self, method: str) -> \"HTTPSpanBuilder\":\n \"\"\"Set HTTP method.\"\"\"\n self._method = method\n if self._url:\n self._name = f\"{method} {self._url}\"\n return self\n\n def with_header(self, key: str, value: str | list[str]) -> \"HTTPSpanBuilder\":\n \"\"\"Add HTTP header.\"\"\"\n if isinstance(value, str):\n value = [value]\n self._headers[key] = value\n return self\n\n def with_headers(self, headers: dict[str, list[str]]) -> \"HTTPSpanBuilder\":\n \"\"\"Add multiple HTTP headers.\"\"\"\n self._headers.update(headers)\n return self\n\n def with_body(self, body: str | bytes) -> \"HTTPSpanBuilder\":\n \"\"\"Set request body.\"\"\"\n self._body = body\n return self\n\n def with_response(self, response: str | bytes) -> \"HTTPSpanBuilder\":\n \"\"\"Set response body.\"\"\"\n self._response = response\n return self\n\n def with_error(self, code: int, message: str) -> \"HTTPSpanBuilder\":\n \"\"\"Set error information.\"\"\"\n self._error = {\"code\": code, \"message\": message}\n return self\n\n def build(self) -> Span:\n \"\"\"Build HTTP Span.\"\"\"\n payload = HTTPSpanPayload(\n type=SpanType.HTTP,\n url=self._url,\n method=self._method,\n headers=self._headers,\n body=self._body,\n response=self._response,\n error=self._error,\n )\n self._payload = payload\n return super().build()\n\n\nclass EventBuilder:\n \"\"\"\n Event builder.\n\n Usage:\n event = (EventBuilder()\n .with_name(\"user_input\")\n .with_data({\"text\": \"hello\"})\n .with_parent(span)\n .build())\n \"\"\"\n\n def __init__(self, trace_id: Optional[str] = None, app_name: Optional[str] = None):\n self._trace_id = trace_id or str(ULID())\n self._app_name = app_name or get_default(\"app_name\")\n self._id = str(ULID())\n self._name = \"\"\n self._timestamp = datetime.now()\n self._tags: dict[str, str] = {}\n self._data: Any = None\n self._parent_id: Optional[str] = None\n\n def with_id(self, event_id: str) -> \"EventBuilder\":\n \"\"\"Set event ID.\"\"\"\n self._id = event_id\n return self\n\n def with_name(self, name: str) -> \"EventBuilder\":\n \"\"\"Set event name.\"\"\"\n self._name = name\n return self\n\n def with_trace_id(self, trace_id: str) -> \"EventBuilder\":\n \"\"\"Set trace ID.\"\"\"\n self._trace_id = trace_id\n return self\n\n def with_timestamp(self, timestamp: datetime) -> \"EventBuilder\":\n \"\"\"Set timestamp.\"\"\"\n self._timestamp = timestamp\n return self\n\n def with_tag(self, key: str, value: str) -> \"EventBuilder\":\n \"\"\"Add a tag.\"\"\"\n self._tags[key] = value\n return self\n\n def with_tags(self, tags: dict[str, str]) -> \"EventBuilder\":\n \"\"\"Add multiple tags.\"\"\"\n self._tags.update(tags)\n return self\n\n def with_data(self, data: Any) -> \"EventBuilder\":\n \"\"\"Set data.\"\"\"\n self._data = data\n return self\n\n def with_parent(self, parent: Span) -> \"EventBuilder\":\n \"\"\"Set parent span.\"\"\"\n self._parent_id = parent.id\n # Automatically inherit trace_id from parent span\n if parent:\n self._trace_id = parent.trace_id\n return self\n\n def with_parent_id(self, parent_id: str) -> \"EventBuilder\":\n \"\"\"Set parent span ID.\"\"\"\n self._parent_id = parent_id\n return self\n\n def build(self) -> Event:\n \"\"\"Build Event object.\"\"\"\n payload = OtherEventPayload(type=EventType.OTHER, data=self._data)\n\n event = Event(\n id=self._id,\n name=self._name,\n trace_id=self._trace_id,\n timestamp=self._timestamp,\n tags=self._tags,\n payload=payload,\n parent_id=self._parent_id,\n app_name=self._app_name,\n )\n return event\n" }, { "path": "agentkit/trace/context.py", "content": "\"\"\"\nContext management and utility functions for Span and Event.\n\"\"\"\n\nfrom contextlib import contextmanager\nfrom contextvars import ContextVar\nfrom datetime import datetime\nfrom functools import wraps\ntry: # Python 3.11+ has HTTPMethod in stdlib\n from http import HTTPMethod\nexcept ImportError: # pragma: no cover - fallback for Python 3.10\n from enum import Enum\n\n class HTTPMethod(str, Enum):\n GET = \"GET\"\n POST = \"POST\"\n PUT = \"PUT\"\n DELETE = \"DELETE\"\n PATCH = \"PATCH\"\n HEAD = \"HEAD\"\n OPTIONS = \"OPTIONS\"\n CONNECT = \"CONNECT\"\n TRACE = \"TRACE\"\nfrom typing import Any, Callable, Literal, Optional\n\nfrom ulid import ULID\n\nfrom .default import get_default, get_default_settings, set_default\nfrom .span import (\n Event,\n EventType,\n FunctionSpanPayload,\n HTTPSpanPayload,\n LLMSpanPayload,\n OtherEventPayload,\n OtherSpanPayload,\n Span,\n SpanType,\n ToolSpanPayload,\n)\nfrom .tracer import Tracer\nfrom .types import Error\n\n# Use contextvars to manage the current context\n_current_context: ContextVar[Optional[\"SpanContext\"]] = ContextVar(\n \"current_context\", default=None\n)\n\n\nclass SpanContext:\n \"\"\"\n Manages the current active span context, supporting nested parent-child relationships.\n\n Supports cross-service span reconstruction: trace_id and parent_id can be specified at creation time.\n \"\"\"\n\n def __init__(\n self,\n app_name: Optional[str] = None,\n tags: Optional[dict[str, str]] = None,\n trace_id: Optional[str] = None,\n parent_id: Optional[str] = None,\n tracer: Optional[Tracer] = None,\n ):\n if tracer is None:\n self.tracer = get_default(\"tracer\")\n else:\n self.tracer = tracer\n self.app_name = app_name or get_default(\"app_name\")\n self.tags = tags\n self._span_stack: list[Span] = []\n self._current_trace_id: Optional[str] = trace_id\n self._root_parent_id: Optional[str] = parent_id # Root parent node for cross-service reconstruction\n\n def get_current_span(self) -> Optional[Span]:\n \"\"\"Get the current active span.\"\"\"\n return self._span_stack[-1] if self._span_stack else None\n\n def get_current_trace_id(self) -> str:\n \"\"\"Get or create the current trace_id.\"\"\"\n if self._current_trace_id is None:\n self._current_trace_id = str(ULID())\n return self._current_trace_id\n\n def set_trace_id(self, trace_id: str):\n \"\"\"Set the current trace_id.\"\"\"\n self._current_trace_id = trace_id\n\n def merge_tags(self, tags: dict[str, str]) -> dict[str, str]:\n \"\"\"Merge tags.\"\"\"\n merged_tags = self.tags or {}\n if tags:\n merged_tags.update(tags)\n return merged_tags\n\n @contextmanager\n def span(\n self,\n name: str,\n tags: Optional[dict[str, str]] = None,\n payload: Optional[\n HTTPSpanPayload\n | LLMSpanPayload\n | ToolSpanPayload\n | FunctionSpanPayload\n | OtherSpanPayload\n ] = None,\n ):\n \"\"\"\n Context manager for creating a span.\n\n Usage:\n with ctx.span(\"my_operation\", SpanType.FUNCTION):\n # your code\n pass\n \"\"\"\n parent_span = self.get_current_span()\n trace_id = self.get_current_trace_id()\n\n # Determine parent_id: prioritize span in current stack, otherwise use root parent_id (cross-service scenario)\n parent_id = None\n if parent_span:\n parent_id = parent_span.id\n elif self._root_parent_id:\n parent_id = self._root_parent_id\n\n span = Span(\n name=name,\n trace_id=trace_id,\n app_name=self.app_name,\n tags=self.merge_tags(tags),\n payload=payload,\n parent_id=parent_id,\n )\n\n self._span_stack.append(span)\n token = _current_context.set(self)\n\n try:\n yield span\n except Exception as e:\n # Record error\n span.tags[\"error\"] = str(e)\n raise\n finally:\n # Set end time\n span.end_time = datetime.now()\n self._span_stack.pop()\n\n # Record span\n if self.tracer:\n self.tracer.record_span(span)\n _current_context.reset(token)\n\n def record_event(\n self,\n name: str,\n data: Any,\n tags: Optional[dict[str, str]] = None,\n ):\n \"\"\"\n Record an event to the current span.\n\n Usage:\n ctx.record_event(\"user_input\", {\"text\": \"hello\"})\n \"\"\"\n parent_span = self.get_current_span()\n trace_id = self.get_current_trace_id()\n\n payload = OtherEventPayload(type=EventType.OTHER, data=data)\n\n event = Event(\n name=name,\n trace_id=trace_id,\n tags=self.merge_tags(tags),\n payload=payload,\n parent_id=parent_span.id if parent_span else None,\n app_name=self.app_name,\n )\n\n if self.tracer:\n self.tracer.record_event(event)\n\n return event\n\n @contextmanager\n def function_span(\n self,\n name: str,\n arguments: dict[str, Any],\n tags: Optional[dict[str, str]] = None,\n ):\n \"\"\"\n Create a span for a function call.\n\n Usage:\n with ctx.function_span(\"calculate\", {\"x\": 1, \"y\": 2}) as span:\n result = calculate(1, 2)\n span.update_payload_data(return_value=result)\n \"\"\"\n payload = FunctionSpanPayload(\n type=SpanType.FUNCTION,\n name=name,\n arguments=arguments,\n return_value=None,\n )\n\n with self.span(name, tags, payload) as span:\n try:\n yield span\n except Exception as e:\n if isinstance(span.payload, FunctionSpanPayload):\n span.payload.error = Error(code=-1, message=str(e))\n raise\n\n @contextmanager\n def llm_span(\n self,\n name: str = \"llm_call\",\n request: Any = None,\n tags: Optional[dict[str, str]] = None,\n ):\n \"\"\"\n Create a span for an LLM call.\n\n Usage:\n with ctx.llm_span(\"openai_call\", request=messages) as span:\n response = client.chat.completions.create(...)\n span.update_payload_data(response=response)\n \"\"\"\n payload = LLMSpanPayload(\n type=SpanType.LLM,\n request=request,\n )\n\n with self.span(name, tags, payload) as span:\n try:\n yield span\n except Exception as e:\n if isinstance(span.payload, LLMSpanPayload):\n span.payload.error = Error(code=-1, message=str(e))\n raise\n\n @contextmanager\n def tool_span(\n self,\n name: str = \"tool_call\",\n request: Any = None,\n tags: Optional[dict[str, str]] = None,\n ):\n \"\"\"\n Create a span for a tool call.\n\n Usage:\n with ctx.tool_span(\"search_tool\", request=tool_call) as span:\n result = search(query)\n span.update_payload_data(response=result)\n \"\"\"\n payload = ToolSpanPayload(type=SpanType.TOOL, request=request)\n\n with self.span(name, tags, payload) as span:\n try:\n yield span\n except Exception as e:\n if isinstance(span.payload, ToolSpanPayload):\n span.payload.error = Error(code=-1, message=str(e))\n raise\n\n @contextmanager\n def http_span(\n self,\n url: str,\n method: Literal[\n HTTPMethod.GET,\n HTTPMethod.POST,\n HTTPMethod.PUT,\n HTTPMethod.DELETE,\n HTTPMethod.PATCH,\n HTTPMethod.HEAD,\n HTTPMethod.OPTIONS,\n HTTPMethod.CONNECT,\n HTTPMethod.TRACE,\n ],\n name: Optional[str] = None,\n headers: Optional[dict[str, list[str]]] = None,\n body: Optional[str | bytes] = None,\n tags: Optional[dict[str, str]] = None,\n ):\n \"\"\"\n Create a span for an HTTP request.\n\n Usage:\n with ctx.http_span(\"https://api.example.com\", \"POST\", body=data) as span:\n response = requests.post(url, data=data)\n span.update_payload_data(response=response.text)\n \"\"\"\n span_name = name or f\"{method} {url}\"\n payload = HTTPSpanPayload(\n type=SpanType.HTTP,\n url=url,\n method=method,\n headers=headers or {},\n body=body,\n )\n\n with self.span(span_name, tags, payload) as span:\n try:\n yield span\n except Exception as e:\n if isinstance(span.payload, HTTPSpanPayload):\n span.payload.error = Error(code=-1, message=str(e))\n raise\n\n\ndef get_current_context() -> SpanContext:\n \"\"\"Get the current SpanContext, creating a default one if it doesn't exist.\"\"\"\n context = _current_context.get()\n if context is None:\n context = SpanContext()\n _current_context.set(context)\n return context\n\n\ndef set_default_tracer(tracer: Tracer):\n \"\"\"Set the global default Tracer.\"\"\"\n set_default(tracer=tracer)\n\n\ndef set_default_app_name(app_name: str):\n \"\"\"Set the global default app_name.\"\"\"\n set_default(app_name=app_name)\n\n\ndef get_default_tracer() -> Tracer:\n \"\"\"Get the global default Tracer.\"\"\"\n return get_default_settings().tracer\n\n\ndef get_default_app_name() -> str:\n \"\"\"Get the global default app_name.\"\"\"\n return get_default_settings().app_name\n\n\n# Convenient global functions\ndef trace_function(\n name: Optional[str] = None,\n tags: Optional[dict[str, str]] = None,\n context: Optional[SpanContext] = None,\n):\n \"\"\"\n Decorator: automatically trace function calls.\n\n Usage:\n @trace_function(name=\"my_function\", tags={\"category\": \"business\"})\n def my_function(x, y):\n return x + y\n \"\"\"\n\n def decorator(func: Callable) -> Callable:\n func_name = name or func.__name__\n\n @wraps(func)\n def wrapper(*args, **kwargs):\n ctx = context or get_current_context()\n\n # Build arguments dictionary\n arguments = {\"args\": args, \"kwargs\": kwargs}\n\n with ctx.function_span(func_name, arguments, tags) as span:\n result = func(*args, **kwargs)\n\n # Record return value\n if isinstance(span.payload, FunctionSpanPayload):\n span.payload.return_value = result\n\n return result\n\n @wraps(func)\n async def async_wrapper(*args, **kwargs):\n ctx = context or get_current_context()\n\n # Build arguments dictionary\n arguments = {\"args\": args, \"kwargs\": kwargs}\n\n with ctx.function_span(func_name, arguments, tags) as span:\n result = await func(*args, **kwargs)\n\n # Record return value\n if isinstance(span.payload, FunctionSpanPayload):\n span.payload.return_value = result\n\n return result\n\n # Return corresponding wrapper based on function type\n import asyncio\n\n if asyncio.iscoroutinefunction(func):\n return async_wrapper\n else:\n return wrapper\n\n return decorator\n\n\ndef start_trace(trace_id: Optional[str] = None, context: Optional[SpanContext] = None):\n \"\"\"\n Start a new trace.\n\n Usage:\n trace_id = start_trace()\n \"\"\"\n ctx = context or get_current_context()\n if trace_id:\n ctx.set_trace_id(trace_id)\n else:\n trace_id = ctx.get_current_trace_id()\n return trace_id\n\n\ndef record_event(\n name: str,\n data: Any,\n tags: Optional[dict[str, str]] = None,\n context: Optional[SpanContext] = None,\n):\n \"\"\"\n Record an event.\n\n Usage:\n record_event(\"user_action\", {\"action\": \"click\", \"button\": \"submit\"})\n \"\"\"\n ctx = context or get_current_context()\n return ctx.record_event(name, data, tags=tags)\n\n\n@contextmanager\ndef create_span(\n name: str,\n tags: Optional[dict[str, str]] = None,\n context: Optional[SpanContext] = None,\n):\n \"\"\"\n Convenient function to create a span.\n\n Usage:\n with create_span(\"my_operation\"):\n # your code\n pass\n \"\"\"\n ctx = context or get_current_context()\n with ctx.span(name, tags) as span:\n yield span\n" }, { "path": "agentkit/trace/default.py", "content": "from dataclasses import dataclass\n\nfrom .local_tracer import LocalStorageTracer\nfrom .tracer import Tracer\n\n\n@dataclass\nclass DefaultSettings:\n app_name: str = \"default\"\n tracer: Tracer = LocalStorageTracer(storage_dir=\"./traces\")\n\n\n_settings = DefaultSettings()\n\n\ndef set_default(**kwargs):\n for key, value in kwargs.items():\n if hasattr(_settings, key):\n setattr(_settings, key, value)\n else:\n raise ValueError(f\"Unknown setting: {key}\")\n\n\ndef get_default_settings() -> DefaultSettings:\n return _settings\n\n\ndef get_default(key: str):\n \"\"\"Get a single default value.\"\"\"\n return getattr(_settings, key)\n" }, { "path": "agentkit/trace/local_tracer.py", "content": "import json\nfrom datetime import datetime\nfrom pathlib import Path\nfrom typing import Optional\n\nfrom .span import Event, Span\nfrom .tracer import Tracer\n\n\nclass LocalStorageTracer(Tracer):\n def __init__(self, storage_dir: str = \"./traces\"):\n self.storage_dir = Path(storage_dir)\n self.storage_dir.mkdir(parents=True, exist_ok=True)\n self.spans_dir = self.storage_dir / \"spans\"\n self.events_dir = self.storage_dir / \"events\"\n self.spans_dir.mkdir(exist_ok=True)\n self.events_dir.mkdir(exist_ok=True)\n\n def _get_trace_spans_file(self, trace_id: str) -> Path:\n return self.spans_dir / f\"{trace_id}.jsonl\"\n\n def _get_trace_events_file(self, trace_id: str) -> Path:\n return self.events_dir / f\"{trace_id}.jsonl\"\n\n def record_span(self, span: Span) -> None:\n spans_file = self._get_trace_spans_file(span.trace_id)\n with open(spans_file, \"a\", encoding=\"utf-8\") as f:\n span_data = span.model_dump_json(exclude_none=True, ensure_ascii=False)\n f.write(span_data + \"\\n\")\n\n def record_event(self, event: Event) -> None:\n events_file = self._get_trace_events_file(event.trace_id)\n with open(events_file, \"a\", encoding=\"utf-8\") as f:\n event_data = event.model_dump_json(exclude_none=True, ensure_ascii=False)\n f.write(event_data + \"\\n\")\n\n def get_spans(self, trace_id: str) -> list[Span]:\n spans_file = self._get_trace_spans_file(trace_id)\n if not spans_file.exists():\n return []\n\n spans = []\n with open(spans_file, \"r\", encoding=\"utf-8\") as f:\n for line_num, line in enumerate(f, 1):\n if line.strip():\n try:\n span_data = json.loads(line)\n spans.append(Span(**span_data))\n except Exception as e:\n # Log error but continue processing other lines\n print(\n f\"Warning: Failed to parse span at line {line_num} in {spans_file}: {e}\"\n )\n continue\n return spans\n\n def get_events(self, trace_id: str) -> list[Event]:\n events_file = self._get_trace_events_file(trace_id)\n if not events_file.exists():\n return []\n\n events = []\n with open(events_file, \"r\", encoding=\"utf-8\") as f:\n for line_num, line in enumerate(f, 1):\n if line.strip():\n try:\n event_data = json.loads(line)\n events.append(Event(**event_data))\n except Exception as e:\n # Log error but continue processing other lines\n print(\n f\"Warning: Failed to parse event at line {line_num} in {events_file}: {e}\"\n )\n continue\n return events\n\n def get_trace(self, trace_id: str) -> Optional[dict]:\n spans = self.get_spans(trace_id)\n events = self.get_events(trace_id)\n\n if not spans and not events:\n return None\n\n return {\n \"trace_id\": trace_id,\n \"spans\": [span.model_dump(mode=\"json\") for span in spans],\n \"events\": [event.model_dump(mode=\"json\") for event in events],\n \"span_count\": len(spans),\n \"event_count\": len(events),\n }\n\n def get_trace_raw(self, trace_id: str) -> Optional[dict]:\n \"\"\"\n Get raw trace data (without Pydantic model validation).\n Used for frontend display to avoid serialization/deserialization issues.\n \"\"\"\n spans_file = self._get_trace_spans_file(trace_id)\n events_file = self._get_trace_events_file(trace_id)\n\n if not spans_file.exists() and not events_file.exists():\n return None\n\n spans = []\n events = []\n\n # Read spans (raw JSON)\n if spans_file.exists():\n with open(spans_file, \"r\", encoding=\"utf-8\") as f:\n for line in f:\n if line.strip():\n try:\n span_data = json.loads(line)\n spans.append(span_data)\n except json.JSONDecodeError as e:\n # Skip invalid line\n print(f\"Warning: Failed to parse span line: {e}\")\n continue\n\n # Read events (raw JSON)\n if events_file.exists():\n with open(events_file, \"r\", encoding=\"utf-8\") as f:\n for line in f:\n if line.strip():\n try:\n event_data = json.loads(line)\n events.append(event_data)\n except json.JSONDecodeError as e:\n # Skip invalid line\n print(f\"Warning: Failed to parse event line: {e}\")\n continue\n\n return {\n \"trace_id\": trace_id,\n \"spans\": spans,\n \"events\": events,\n \"span_count\": len(spans),\n \"event_count\": len(events),\n }\n\n def list_traces(self, limit: int = 100, offset: int = 0) -> list[dict]:\n trace_files = {}\n\n for spans_file in self.spans_dir.glob(\"*.jsonl\"):\n trace_id = spans_file.stem\n trace_files[trace_id] = {\n \"trace_id\": trace_id,\n \"spans_file\": spans_file,\n \"mtime\": spans_file.stat().st_mtime,\n }\n\n for events_file in self.events_dir.glob(\"*.jsonl\"):\n trace_id = events_file.stem\n if trace_id not in trace_files:\n trace_files[trace_id] = {\n \"trace_id\": trace_id,\n \"events_file\": events_file,\n \"mtime\": events_file.stat().st_mtime,\n }\n else:\n trace_files[trace_id][\"events_file\"] = events_file\n trace_files[trace_id][\"mtime\"] = max(\n trace_files[trace_id][\"mtime\"], events_file.stat().st_mtime\n )\n\n sorted_traces = sorted(\n trace_files.values(), key=lambda x: x[\"mtime\"], reverse=True\n )\n\n traces = []\n for trace_info in sorted_traces[offset : offset + limit]:\n trace_id = trace_info[\"trace_id\"]\n spans = self.get_spans(trace_id)\n events = self.get_events(trace_id)\n\n traces.append(\n {\n \"trace_id\": trace_id,\n \"span_count\": len(spans),\n \"event_count\": len(events),\n \"last_modified\": datetime.fromtimestamp(trace_info[\"mtime\"]),\n }\n )\n\n return traces\n" }, { "path": "agentkit/trace/remote_tracer.py", "content": "import logging\nfrom datetime import datetime\nfrom typing import Optional\n\nimport httpx\nfrom pydantic import BaseModel\n\nfrom .span import DataType, Event, Span\nfrom .tracer import Tracer\n\nlogger = logging.getLogger(__name__)\n\n\nclass RemoteEvent(BaseModel):\n id: str\n data_type: DataType\n timestamp: datetime\n app_name: str\n data: Span | Event\n\n\nclass RemoteTracer(Tracer):\n \"\"\"\n A tracer that sends trace data to a remote API.\n\n According to API docs: /trace/agent/event\n \"\"\"\n\n def __init__(\n self,\n base_url: str = \"\",\n timeout: float = 10.0,\n enable_span: bool = True,\n enable_event: bool = True,\n ):\n \"\"\"\n Initialize RemoteTracer.\n\n Args:\n base_url: API base URL.\n timeout: Request timeout in seconds.\n enable_span: Whether to enable span sending (default True).\n enable_event: Whether to enable event sending (default True).\n \"\"\"\n self.base_url = base_url.rstrip(\"/\")\n self.event_endpoint = f\"{self.base_url}/trace/agent/event\"\n self.timeout = timeout\n self.enable_span = enable_span\n self.enable_event = enable_event\n self.client = httpx.Client(timeout=timeout)\n\n def __del__(self):\n \"\"\"Clean up resources.\"\"\"\n try:\n self.client.close()\n except Exception:\n pass\n\n def _send_to_api(self, data: str) -> bool:\n \"\"\"\n Send data to remote API.\n\n Args:\n data: Data to send.\n\n Returns:\n Whether the send was successful.\n \"\"\"\n try:\n response = self.client.post(\n self.event_endpoint,\n data=data,\n headers={\"Content-Type\": \"application/json\"},\n )\n\n if response.status_code == 200:\n result = response.json()\n if result.get(\"code\") == 0:\n logger.debug(f\"Successfully sent trace data: {data}\")\n return True\n else:\n logger.error(\n f\"API returned error code {result.get('code')}: {result.get('msg')}\"\n )\n return False\n else:\n logger.error(f\"HTTP error {response.status_code}: {response.text}\")\n return False\n\n except httpx.TimeoutException:\n logger.error(f\"Request timeout when sending trace data: {data}\")\n return False\n except Exception as e:\n logger.error(f\"Failed to send trace data: {e}\")\n return False\n\n def record_span(self, span: Span) -> None:\n \"\"\"Record a span to the remote service.\"\"\"\n if not self.enable_span:\n return\n\n remote_event = RemoteEvent(\n id=span.id,\n data_type=DataType.SPAN,\n timestamp=span.start_time,\n app_name=span.app_name,\n data=span,\n )\n\n self._send_to_api(remote_event.model_dump_json(exclude_none=True))\n\n def record_event(self, event: Event) -> None:\n \"\"\"Record an event to the remote service.\"\"\"\n if not self.enable_event:\n return\n\n remote_event = RemoteEvent(\n id=event.id,\n data_type=DataType.EVENT,\n timestamp=event.timestamp,\n app_name=event.app_name,\n data=event,\n )\n\n self._send_to_api(remote_event.model_dump_json(exclude_none=True))\n\n def get_spans(self, trace_id: str) -> list[Span]:\n \"\"\"\n RemoteTracer does not support read operations.\n\n Note: The current remote API only provides write interface, query is not supported.\n \"\"\"\n logger.warning(\"RemoteTracer does not support reading spans\")\n return []\n\n def get_events(self, trace_id: str) -> list[Event]:\n \"\"\"\n RemoteTracer does not support read operations.\n\n Note: The current remote API only provides write interface, query is not supported.\n \"\"\"\n logger.warning(\"RemoteTracer does not support reading events\")\n return []\n\n def get_trace(self, trace_id: str) -> Optional[dict]:\n \"\"\"\n RemoteTracer does not support read operations.\n\n Note: The current remote API only provides write interface, query is not supported.\n \"\"\"\n logger.warning(\"RemoteTracer does not support reading traces\")\n return None\n\n def list_traces(self, limit: int = 100, offset: int = 0) -> list[dict]:\n \"\"\"\n RemoteTracer does not support read operations.\n\n Note: The current remote API only provides write interface, query is not supported.\n \"\"\"\n logger.warning(\"RemoteTracer does not support listing traces\")\n return []\n\n\nclass HybridTracer(Tracer):\n \"\"\"\n Hybrid tracer: sends data to both remote service and local storage.\n\n Usage:\n from agentkit.trace.local_tracer import LocalStorageTracer\n from agentkit.trace.remote_tracer import RemoteTracer, HybridTracer\n\n local_tracer = LocalStorageTracer(\"./traces\")\n remote_tracer = RemoteTracer(\"xxx\")\n hybrid_tracer = HybridTracer(local_tracer, remote_tracer)\n \"\"\"\n\n def __init__(self, local_tracer: Tracer, remote_tracer: RemoteTracer):\n \"\"\"\n Initialize hybrid tracer.\n\n Args:\n local_tracer: Local tracer (for reading and local storage).\n remote_tracer: Remote tracer (for remote reporting).\n \"\"\"\n self.local_tracer = local_tracer\n self.remote_tracer = remote_tracer\n\n def record_span(self, span: Span) -> None:\n \"\"\"Record to both local and remote.\"\"\"\n self.local_tracer.record_span(span)\n self.remote_tracer.record_span(span)\n\n def record_event(self, event: Event) -> None:\n \"\"\"Record to both local and remote.\"\"\"\n self.local_tracer.record_event(event)\n self.remote_tracer.record_event(event)\n\n def get_spans(self, trace_id: str) -> list[Span]:\n \"\"\"Read spans from local.\"\"\"\n return self.local_tracer.get_spans(trace_id)\n\n def get_events(self, trace_id: str) -> list[Event]:\n \"\"\"Read events from local.\"\"\"\n return self.local_tracer.get_events(trace_id)\n\n def get_trace(self, trace_id: str) -> Optional[dict]:\n \"\"\"Read trace from local.\"\"\"\n return self.local_tracer.get_trace(trace_id)\n\n def get_trace_raw(self, trace_id: str) -> Optional[dict]:\n \"\"\"Read raw trace from local.\"\"\"\n return self.local_tracer.get_trace_raw(trace_id)\n\n def list_traces(self, limit: int = 100, offset: int = 0) -> list[dict]:\n \"\"\"List traces from local.\"\"\"\n return self.local_tracer.list_traces(limit, offset)\n" }, { "path": "agentkit/trace/span.py", "content": "from datetime import datetime\nfrom enum import Enum\ntry: # Python 3.11+ has HTTPMethod in stdlib\n from http import HTTPMethod\nexcept ImportError: # pragma: no cover - fallback for Python 3.10\n class HTTPMethod(str, Enum):\n GET = \"GET\"\n POST = \"POST\"\n PUT = \"PUT\"\n DELETE = \"DELETE\"\n PATCH = \"PATCH\"\n HEAD = \"HEAD\"\n OPTIONS = \"OPTIONS\"\n CONNECT = \"CONNECT\"\n TRACE = \"TRACE\"\nfrom typing import Any, Literal, Optional\n\nfrom pydantic import BaseModel, Field\nfrom ulid import ULID\n\nfrom .types import Error\n\n\nclass DataType(str, Enum):\n SPAN = \"span\"\n EVENT = \"event\"\n\n\nclass SpanType(str, Enum):\n LLM = \"llm_span\"\n TOOL = \"tool_span\"\n HTTP = \"http_span\"\n FUNCTION = \"function_span\"\n OTHER = \"other_span\"\n\n\nclass LLMSpanPayload(BaseModel):\n type: Literal[SpanType.LLM]\n request: Any = None\n response: Any = None\n error: Optional[Error] = None\n\n\nclass ToolSpanPayload(BaseModel):\n type: Literal[SpanType.TOOL]\n request: Any = None\n response: Any = None\n error: Optional[Error] = None\n\n\nclass FunctionSpanPayload(BaseModel):\n type: Literal[SpanType.FUNCTION]\n name: str = \"\"\n arguments: dict[str, Any] = Field(default_factory=dict)\n return_value: Any = None\n error: Optional[Error] = None\n\n\nclass HTTPSpanPayload(BaseModel):\n type: Literal[SpanType.HTTP]\n url: str\n method: Literal[\n HTTPMethod.GET,\n HTTPMethod.POST,\n HTTPMethod.PUT,\n HTTPMethod.DELETE,\n HTTPMethod.PATCH,\n HTTPMethod.HEAD,\n HTTPMethod.OPTIONS,\n HTTPMethod.CONNECT,\n HTTPMethod.TRACE,\n ]\n headers: dict[str, list[str]] = Field(default_factory=dict)\n body: Optional[str | bytes] = None\n response: Optional[str | bytes] = None\n error: Optional[Error] = None\n\n\nclass OtherSpanPayload(BaseModel):\n type: Literal[SpanType.OTHER]\n data: Any\n\n\nclass Span(BaseModel):\n # use ulid\n id: str = Field(default_factory=lambda: str(ULID()))\n name: str = Field(default=\"\")\n data_type: Literal[DataType.SPAN] = DataType.SPAN\n start_time: datetime = Field(default_factory=datetime.now)\n end_time: Optional[datetime] = Field(default=None)\n tags: dict[str, str] = Field(default_factory=dict)\n payload: (\n LLMSpanPayload\n | ToolSpanPayload\n | FunctionSpanPayload\n | OtherSpanPayload\n | HTTPSpanPayload\n | None\n ) = Field(default=None)\n parent_id: Optional[str] = Field(default=None) # Parent span ID\n trace_id: str\n app_name: str\n\n def update_payload(\n self,\n payload: (\n LLMSpanPayload\n | ToolSpanPayload\n | FunctionSpanPayload\n | HTTPSpanPayload\n | OtherSpanPayload\n | None\n ),\n ) -> \"Span\":\n \"\"\"Update span payload.\"\"\"\n self.payload = payload\n return self\n\n def update_payload_data(self, **kwargs) -> \"Span\":\n \"\"\"\n Update specific fields in payload.\n\n Usage:\n span.update_payload_data(ret=result, error=None)\n \"\"\"\n if self.payload is not None and hasattr(self.payload, \"model_copy\"):\n # Use pydantic's model_copy to update fields\n self.payload = self.payload.model_copy(update=kwargs)\n return self\n\n def add_tag(self, key: str, value: str) -> \"Span\":\n \"\"\"Add or update a tag.\"\"\"\n self.tags[key] = value\n return self\n\n def add_tags(self, tags: dict[str, str]) -> \"Span\":\n \"\"\"Add or update multiple tags.\"\"\"\n self.tags.update(tags)\n return self\n\n\nclass EventType(str, Enum):\n DELTA = \"delta_event\"\n OTHER = \"other_event\"\n\n\nclass DeltaEventPayload(BaseModel):\n type: Literal[EventType.DELTA]\n delta: Any\n\n\nclass OtherEventPayload(BaseModel):\n type: Literal[EventType.OTHER]\n data: Any\n\n\nclass Event(BaseModel):\n id: str = Field(default_factory=lambda: str(ULID()))\n name: str = Field(default=\"\")\n data_type: Literal[DataType.EVENT] = DataType.EVENT\n timestamp: datetime = Field(default_factory=datetime.now)\n tags: dict[str, str] = Field(default_factory=dict)\n payload: DeltaEventPayload | OtherEventPayload\n parent_id: Optional[str] = Field(default=None) # Parent span ID\n trace_id: str\n app_name: str\n" }, { "path": "agentkit/trace/tracer.py", "content": "from abc import ABC, abstractmethod\nfrom typing import Optional\n\nfrom .span import Event, Span\n\n\nclass Tracer(ABC):\n @abstractmethod\n def record_span(self, span: Span) -> None:\n pass\n\n @abstractmethod\n def record_event(self, event: Event) -> None:\n pass\n\n @abstractmethod\n def get_spans(self, trace_id: str) -> list[Span]:\n pass\n\n @abstractmethod\n def get_events(self, trace_id: str) -> list[Event]:\n pass\n\n @abstractmethod\n def get_trace(self, trace_id: str) -> Optional[dict]:\n pass\n\n @abstractmethod\n def list_traces(self, limit: int = 100, offset: int = 0) -> list[dict]:\n pass\n\n def get_trace_raw(self, trace_id: str) -> Optional[dict]:\n \"\"\"\n Get raw trace data (optional implementation).\n Default behavior is to call get_trace().\n \"\"\"\n return self.get_trace(trace_id)\n" }, { "path": "agentkit/trace/types.py", "content": "from pydantic import BaseModel\n\n\nclass Error(BaseModel):\n code: int\n message: str\n\n\n# class ModelParams(BaseModel):\n# name: str\n# response_format: Optional[Any] = None\n# toolcall_parser_version: Optional[str] = None\n# parallel_tool_calls: bool = True\n# infer_kwargs: dict[str, Any] = {}\n#\n#\n# class Function(BaseModel):\n# arguments: str | None = None\n# \"\"\"\n# The arguments to call the function with, as generated by the model in JSON\n# format. Note that the model does not always generate valid JSON, and may\n# hallucinate parameters not defined by your function schema. Validate the\n# arguments in your code before calling your function.\n# \"\"\"\n#\n# name: str | None = None\n# \"\"\"The name of the function to call.\"\"\"\n#\n#\n# class ChatToolCall(BaseModel):\n# index: int | None = None\n# \"\"\"The index of the tool call.\"\"\"\n#\n# id: str | None = None\n# \"\"\"The ID of the tool call.\"\"\"\n#\n# function: Function\n# \"\"\"The function that the model called.\"\"\"\n#\n# type: str | None = \"function\"\n# \"\"\"The type of the tool. Currently, only `function` is supported.\"\"\"\n#\n#\n# class ChatMessage(BaseModel):\n# id: Optional[str] = None\n# role: Optional[str] = None\n# content: str | list[dict] | None = None\n# tool_call_id: Optional[str] = None\n# tool_calls: Optional[list[ChatToolCall]] = None\n#\n# # for train use\n# extra_info: Optional[dict] = None\n#\n#\n# class LLMRequest(BaseModel):\n# messages: list[ChatMessage]\n# model: ModelParams\n" }, { "path": "config.yaml", "content": "# Global runtime defaults for StepDeepResearch.\n#\n# Used as the default context budget for the force-final-answer workflow in BaseStepAgent.\n# `scripts.runner` can override these via its own `--config` YAML or CLI flags.\ncontext_upper_limit: 80000\ncontext_lower_limit: 60000\n" }, { "path": "cortex/__init__.py", "content": "from .env import load_env as _load_env\n\n_load_env()\n" }, { "path": "cortex/agents/__init__.py", "content": "\"\"\"Agent components module.\"\"\"\n\nfrom cortex.agents.base_agent import BaseAgent\nfrom cortex.agents.base_step_agent import BaseStepAgent\nfrom cortex.agents.react_agent import ReActAgent\nfrom cortex.agents.types import (\n AgentConfig,\n AgentMessageType,\n AgentResponse,\n AgentRunningStatus,\n)\n\n__all__ = [\n \"BaseAgent\",\n \"BaseStepAgent\",\n \"AgentConfig\",\n \"AgentResponse\",\n \"AgentRunningStatus\",\n \"AgentMessageType\",\n \"ReActAgent\",\n]\n" }, { "path": "cortex/agents/agent_factory.py", "content": "\"\"\"\nAgentFactory is the factory class for Agent, responsible for creating and managing Agents\n\"\"\"\n\nfrom typing import Awaitable, Callable\n\nfrom cortex.agents.base_agent import BaseAgent\nfrom cortex.agents.types import AgentConfig\n\n\nclass AgentFactory:\n \"\"\"\n AgentFactory is the factory class for Agent, responsible for creating and managing Agents\n \"\"\"\n\n agent_make_func: dict[str, Callable[[AgentConfig, str], Awaitable[BaseAgent]]] = {}\n default_agent_configs: dict[str, AgentConfig] = {}\n\n def list_agents(self) -> list[AgentConfig]:\n \"\"\"\n Return all registered Agent configurations\n \"\"\"\n return list(self.default_agent_configs.values())\n\n def get_default_agent_config(self, name: str) -> AgentConfig:\n \"\"\"\n Get Agent configuration\n \"\"\"\n config = self.default_agent_configs.get(name)\n if config is None:\n raise ValueError(\n f\"AgentConfig not provided, and no default configuration set for '{name}' in factory\"\n )\n return config\n\n def register_agent(\n self,\n name: str,\n make_agent_func: Callable[[AgentConfig, str], Awaitable[BaseAgent]],\n default_config: AgentConfig | None = None,\n ) -> None:\n \"\"\"\n Register Agent\n \"\"\"\n self.agent_make_func[name] = make_agent_func\n if default_config is not None:\n self.default_agent_configs[name] = default_config\n\n async def make_agent(\n self, name: str, context_id: str, agent_config: AgentConfig | None\n ) -> BaseAgent:\n \"\"\"\n Create Agent\n \"\"\"\n if name not in self.agent_make_func:\n raise ValueError(f\"Agent {name} not found\")\n config = agent_config or self.default_agent_configs.get(name)\n if config is None:\n raise ValueError(\n f\"AgentConfig not provided, and no default configuration set for '{name}' in factory\"\n )\n make_agent_func = self.agent_make_func[name]\n return await make_agent_func(config, context_id)\n" }, { "path": "cortex/agents/base_agent.py", "content": "\"\"\"Base Agent class, provides run() interface as the base class for all Agents.\"\"\"\n\nimport asyncio\nimport copy\nimport logging\nfrom abc import abstractmethod\nfrom typing import Any, AsyncGenerator\n\nfrom cortex.model.definition import ChatMessage, ChatToolCall, ContentBlockType\n\nfrom cortex.agents.input.input import InputChannel\nfrom cortex.agents.types import AgentConfig, AgentResponse\nfrom cortex.model import ModelAPI\nfrom cortex.model.provider import ModelProvider\nfrom cortex.model.stepfun_provider import StepFunModelProvider\nfrom cortex.tools.toolset import ToolSet\n\nlogger = logging.getLogger(__name__)\n\n\nclass BaseAgent:\n \"\"\"Base Agent class, provides run() interface.\"\"\"\n\n model: Any\n name: str | None = None\n description: str | None = None\n system_prompt: str | None = None\n max_steps: int | None = 5\n _input_channel: InputChannel[ChatMessage] | None = None\n provider: ModelAPI\n\n def update_from_config(self):\n \"\"\"Update agent properties from config.\"\"\"\n # Iterate config attributes and update self if attribute exists\n for key, _ in self.config.model_dump().items():\n setattr(self, key, copy.deepcopy(getattr(self.config, key)))\n\n self.name = self.name or self.__class__.__name__\n self.description = (\n self.description\n or f\"Call the {self.name} sub-agent to handle specific tasks\"\n )\n\n def __init__(\n self,\n provider: ModelProvider | None = None,\n config: AgentConfig | None = None,\n toolset: ToolSet | None = None,\n ):\n self.config = config\n self.update_from_config()\n\n self._toolset = toolset\n\n if provider is None:\n provider = StepFunModelProvider(model_params=self.model)\n self.provider = ModelAPI(provider)\n\n async def __aenter__(self):\n return self\n\n async def __aexit__(self, exc_type, exc_value, traceback):\n \"\"\"Async context manager exit.\"\"\"\n\n async def run(\n self,\n messages: list[ChatMessage] | InputChannel[ChatMessage],\n additional_kwargs: dict | None = None,\n ) -> AsyncGenerator[AgentResponse, None]:\n \"\"\"Run agent, returns AgentResponse generator.\"\"\"\n async for response in self._run(messages, additional_kwargs):\n if response.agent_name is None:\n response.agent_name = self.name\n yield response\n\n @abstractmethod\n async def _run(\n self,\n messages: list[ChatMessage] | InputChannel[ChatMessage],\n additional_kwargs: dict | None = None,\n ) -> AsyncGenerator[AgentResponse, None]:\n \"\"\"\n Run agent, subclasses must implement this method.\n\n Args:\n messages: Input message list or input channel\n additional_kwargs: Additional parameters\n\n Yields:\n AgentResponse: Agent response object\n \"\"\"\n raise NotImplementedError(\"Subclasses must implement this method\")\n\n def model_api(self) -> ModelAPI:\n return self.provider\n\n def toolset(self) -> ToolSet | None:\n \"\"\"\n Get toolset.\n\n Returns:\n ToolSet | None: Toolset\n \"\"\"\n return self._toolset\n\n def as_tool(self, timeout: float | None = None) -> dict[str, Any]:\n \"\"\"\n Convert Agent to parameters needed for Tool creation.\n\n Returns a dictionary containing parameters for creating AgentTool:\n - name: Tool name (uses agent's name)\n - description: Tool description (uses agent's description)\n - agent_name: Agent name (used to specify which agent to call, as metadata)\n - timeout: Timeout (optional, if provided)\n\n Note: When creating AgentTool, channel parameter is also required, which is not included in this method.\n\n Args:\n timeout: Timeout in seconds, if None it won't be included in the returned dictionary\n\n Returns:\n dict: Dictionary containing parameters for creating AgentTool:\n {\n \"name\": str, # Tool name\n \"description\": str, # Tool description\n \"agent_name\": str, # Agent name (metadata)\n \"timeout\": float, # Timeout (optional)\n }\n\n Example:\n >>> agent = MathAgent(config)\n >>> tool_params = agent.as_tool(timeout=60.0)\n >>> # Channel is also required when creating AgentTool\n >>> tool = AgentTool(**tool_params, channel=channel)\n >>> toolset.register(tool)\n \"\"\"\n agent_name = self.name or self.__class__.__name__\n\n tool_params: dict[str, Any] = {\n \"name\": agent_name,\n \"description\": self.description\n or f\"Call {agent_name} Agent to handle specific tasks\",\n \"agent_name\": agent_name, # As metadata, caller knows which agent this tool corresponds to\n }\n\n if timeout is not None:\n tool_params[\"timeout\"] = timeout\n\n logger.debug(\n \"BaseAgent.as_tool returns tool params: name=%s, agent_name=%s, has_timeout=%s\",\n tool_params[\"name\"],\n tool_params[\"agent_name\"],\n \"timeout\" in tool_params,\n )\n\n return tool_params\n\n @staticmethod\n def has_tool_call(message: ChatMessage) -> bool:\n \"\"\"\n Check if message contains tool calls.\n\n Args:\n message: ChatMessage object\n\n Returns:\n bool: True if message contains tool calls, False otherwise\n \"\"\"\n if not message:\n return False\n tool_calls = getattr(message, \"tool_calls\", None)\n return (\n tool_calls is not None and len(tool_calls) > 0\n if isinstance(tool_calls, (list, tuple))\n else tool_calls is not None\n )\n\n async def _execute_single_tool(self, tool_call: ChatToolCall) -> ChatMessage | None:\n \"\"\"\n Execute a single tool call.\n\n Args:\n tool_call: Tool call object containing function.name, function.arguments, id, etc.\n\n Returns:\n ChatMessage: Tool call result message with role \"tool\"\n \"\"\"\n tool_name = tool_call.function.name\n tool_args = tool_call.function.arguments\n tool_call_id = tool_call.id\n\n try:\n # Execute tool call\n result = await self._toolset.call(\n tool_name=tool_name, parameters=tool_args, tool_call_id=tool_call_id\n )\n logger.info(f\"@{self.name} Tool {tool_name} result: {result}\")\n if result is None:\n return None\n\n return ChatMessage(\n role=\"tool\",\n content=result,\n tool_call_id=tool_call_id,\n )\n\n except Exception as e:\n error_msg = f\"Error calling tool {tool_name}: {str(e)}\"\n logger.error(f\"@{self.name} {error_msg}\")\n\n tool_result_content = [\n {\n \"type\": ContentBlockType.TEXT.value,\n ContentBlockType.TEXT.value: error_msg,\n }\n ]\n\n return ChatMessage(\n role=\"tool\",\n content=tool_result_content,\n tool_call_id=tool_call_id,\n )\n\n async def run_tool_call(self, message: ChatMessage) -> list[ChatMessage]:\n \"\"\"\n Extract tool calls from message and execute them, returning list of tool call result messages.\n\n Args:\n message: ChatMessage object\n\n Returns:\n list[ChatMessage]: List of tool call result messages, each with role \"tool\"\n \"\"\"\n if not message:\n return []\n\n if not self._toolset:\n logger.warning(f\"@{self.name} run_tool_call: toolset not initialized\")\n return []\n\n tool_calls = getattr(message, \"tool_calls\", None)\n if not tool_calls:\n return []\n\n # Ensure tool_calls is a list\n if isinstance(tool_calls, (list, tuple)):\n toolcalls_list = list(tool_calls)\n else:\n toolcalls_list = [tool_calls]\n\n # Execute all tool calls sequentially\n tool_result_messages = []\n for tool_call in toolcalls_list:\n result_message = await self._execute_single_tool(tool_call)\n if result_message is not None:\n tool_result_messages.append(result_message)\n\n return tool_result_messages\n\n async def run_tool_call_concurrency(\n self, message: ChatMessage\n ) -> list[ChatMessage]:\n \"\"\"\n Extract concurrent tool calls from message and execute them (when there are multiple tool calls).\n\n Args:\n message: ChatMessage object\n\n Returns:\n list[ChatMessage]: List of concurrent tool call result messages, only returned when tool call count > 1\n \"\"\"\n if not message:\n return []\n\n if not self._toolset:\n logger.warning(f\"@{self.name} run_tool_call_concurrency: toolset not initialized\")\n return []\n\n tool_calls = getattr(message, \"tool_calls\", None)\n if not tool_calls:\n return []\n\n # Ensure tool_calls is a list\n if isinstance(tool_calls, (list, tuple)):\n toolcalls_list = list(tool_calls)\n else:\n toolcalls_list = [tool_calls]\n\n # Only return when there are multiple tool calls (concurrent scenario)\n if len(toolcalls_list) <= 1:\n return []\n\n # Execute all tool calls concurrently\n tool_result_messages = await asyncio.gather(\n *[self._execute_single_tool(tc) for tc in toolcalls_list]\n )\n\n # Filter out None results\n return [msg for msg in tool_result_messages if msg is not None]\n" }, { "path": "cortex/agents/base_step_agent.py", "content": "import copy\nimport json\nimport logging\nimport re\nfrom abc import abstractmethod\nfrom typing import Any, AsyncGenerator, Callable\n\nfrom agentkit.trace import create_span\nfrom cortex.agents.base_agent import BaseAgent\nfrom cortex.agents.input.input import InputChannel\nfrom cortex.agents.types import (\n AgentConfig,\n AgentMessageType,\n AgentResponse,\n AgentRunningStatus,\n)\nfrom cortex.context import BaseContext\nfrom cortex.model.definition import ChatMessage\nfrom cortex.model.provider import ModelProvider\nfrom cortex.tools.toolset import ToolSet\nfrom cortex.runtime_config import get_context_limit_overrides\n\ntry:\n import tiktoken\nexcept Exception: # noqa: BLE001\n tiktoken = None\n\nlogger = logging.getLogger(__name__)\n\nDEFAULT_FORCE_FINAL_ANSWER_UPPER_LIMIT = 100_000\nDEFAULT_FORCE_FINAL_ANSWER_THRESHOLD = DEFAULT_FORCE_FINAL_ANSWER_UPPER_LIMIT\nDEFAULT_FORCE_FINAL_ANSWER_LOWER_LIMIT_RATIO = 0.9\nDEFAULT_FORCE_FINAL_ANSWER_LOWER_LIMIT = max(\n int(DEFAULT_FORCE_FINAL_ANSWER_UPPER_LIMIT * DEFAULT_FORCE_FINAL_ANSWER_LOWER_LIMIT_RATIO),\n 1,\n)\nDEFAULT_FORCE_FINAL_ANSWER_PROMPT = (\n \"你现在已经达到了你所能处理的最大上下文长度。你应该停止进行工具调用，\"\n \"并基于以上所有信息重新思考，然后按照以下格式提供你认为最可能的答案：\"\n \"你的最终思考\\n你的答案\"\n)\n_AVG_CHARS_PER_TOKEN = 3\n\n\ndef _get_encoding(model_name: str | None):\n if not tiktoken:\n return None\n if not model_name:\n try:\n return tiktoken.get_encoding(\"cl100k_base\")\n except Exception: # noqa: BLE001\n return None\n try:\n return tiktoken.encoding_for_model(model_name)\n except Exception: # noqa: BLE001\n try:\n return tiktoken.get_encoding(\"cl100k_base\")\n except Exception: # noqa: BLE001\n return None\n\n\ndef _estimate_token_length(messages: list[ChatMessage], model_name: str | None) -> int:\n \"\"\"Token estimator with tiktoken fallback.\"\"\"\n encoding = _get_encoding(model_name)\n total_tokens = 0\n for message in messages:\n try:\n payload = message.model_dump(exclude_none=True)\n except Exception:\n payload = {\n \"role\": getattr(message, \"role\", None),\n \"content\": getattr(message, \"content\", None),\n }\n serialized = json.dumps(payload, ensure_ascii=False)\n if encoding:\n try:\n total_tokens += len(encoding.encode(serialized))\n continue\n except Exception: # noqa: BLE001\n encoding = None\n total_tokens += len(serialized) // _AVG_CHARS_PER_TOKEN\n return total_tokens\n\n\ndef _compress_batch_search_result(content: str) -> str:\n \"\"\"Strip verbose content and mark compressed.\"\"\"\n compressed = re.sub(r\".*?\\s*\", \"\", content, flags=re.S)\n compressed = compressed.replace(\n \"\", \"\", 1\n )\n compressed = compressed.replace(\n \"\", \"\", 1\n )\n return compressed\n\n\nclass BaseStepAgent(BaseAgent):\n \"\"\"\n Base class for step-based Agent\n\n Implements run() method, executes tasks by calling step() in a loop\n step() method returns a tuple containing a flag indicating whether to stop\n \"\"\"\n\n def __init__(\n self,\n context: BaseContext,\n config: AgentConfig,\n provider: ModelProvider | None = None,\n toolset: ToolSet | None = None,\n ):\n super().__init__(config=config, toolset=toolset, provider=provider)\n self.current_round = 0\n self.context = context\n extra_cfg = config.extra_config if config and config.extra_config else {}\n self._force_final_answer_enabled = extra_cfg.get(\"force_final_answer\", False)\n threshold_override = extra_cfg.get(\"final_answer_context_threshold\")\n upper_override = extra_cfg.get(\"final_answer_context_upper_limit\")\n lower_override = extra_cfg.get(\"final_answer_context_lower_limit\")\n upper_from_extra = upper_override is not None\n lower_from_extra = lower_override is not None\n if not upper_from_extra or not lower_from_extra:\n runtime_upper, runtime_lower = get_context_limit_overrides()\n if not upper_from_extra and runtime_upper is not None:\n upper_override = runtime_upper\n # Only apply runtime lower default when upper isn't explicitly overridden;\n # otherwise keep the existing \"derive lower from upper\" behavior.\n if (\n not lower_from_extra\n and not upper_from_extra\n and runtime_lower is not None\n ):\n lower_override = runtime_lower\n\n def _normalize_limit(value: Any) -> int | None:\n if isinstance(value, (int, float)):\n return int(value)\n return None\n\n upper_limit = _normalize_limit(upper_override)\n if upper_limit is None:\n upper_limit = _normalize_limit(threshold_override)\n if upper_limit is None or upper_limit <= 0:\n upper_limit = DEFAULT_FORCE_FINAL_ANSWER_UPPER_LIMIT\n upper_limit = max(upper_limit, 2)\n\n lower_limit = _normalize_limit(lower_override)\n if lower_limit is None or lower_limit <= 0:\n derived = int(upper_limit * DEFAULT_FORCE_FINAL_ANSWER_LOWER_LIMIT_RATIO)\n lower_limit = derived if derived > 0 else DEFAULT_FORCE_FINAL_ANSWER_LOWER_LIMIT\n lower_limit = max(lower_limit, 1)\n if lower_limit >= upper_limit:\n lower_limit = max(upper_limit - 1, 1)\n\n self._force_final_answer_upper_limit = upper_limit\n self._force_final_answer_lower_limit = lower_limit\n self._force_final_answer_prompt = extra_cfg.get(\n \"final_answer_prompt\", DEFAULT_FORCE_FINAL_ANSWER_PROMPT\n )\n self._force_prompt_inserted = False\n self._model_name = getattr(config.model, \"name\", None) if config and config.model else None\n\n def _insert_final_prompt(self) -> None:\n \"\"\"Activate the force-final-answer prompt for subsequent model calls.\n\n Kept for backward compatibility; this does not mutate persisted context history.\n \"\"\"\n if not self._force_final_answer_enabled or self._force_prompt_inserted:\n return\n self._force_prompt_inserted = True\n logger.info(\"@%s Final answer prompt activated\", self.name)\n\n def _make_force_final_answer_message(self) -> ChatMessage:\n prompt = self._force_final_answer_prompt or DEFAULT_FORCE_FINAL_ANSWER_PROMPT\n return ChatMessage(role=\"system\", content=prompt)\n\n def _ensure_final_prompt(self, messages: list[ChatMessage]) -> None:\n \"\"\"Ensure the force-final-answer prompt is present in model input.\n\n This should not mutate the persisted context history.\n \"\"\"\n if not self._force_final_answer_enabled or not self._force_prompt_inserted:\n return\n prompt_message = self._make_force_final_answer_message()\n if messages:\n last = messages[-1]\n if (\n getattr(last, \"role\", None) == prompt_message.role\n and getattr(last, \"content\", None) == prompt_message.content\n ):\n return\n messages.append(prompt_message)\n\n @staticmethod\n def _copy_messages(messages: list[ChatMessage]) -> list[ChatMessage]:\n \"\"\"Deep-copy messages so we can mutate model input without touching stored history.\"\"\"\n copied: list[ChatMessage] = []\n for message in messages:\n try:\n copied.append(message.model_copy(deep=True))\n except Exception:\n copied.append(copy.deepcopy(message))\n return copied\n\n @classmethod\n def _compress_batch_search_in_block(cls, block: dict) -> bool:\n \"\"\"Compress batch_search XML in the first text block found (in-place).\"\"\"\n if block.get(\"type\") == \"text\":\n text_value = block.get(\"text\")\n if (\n isinstance(text_value, str)\n and \" tuple[Any, bool]:\n if isinstance(content, str):\n if (\n \" list[ChatMessage]:\n \"\"\"Build the message list sent to the model, without mutating stored history.\"\"\"\n raw_messages = self.context.get_all()\n messages: list[ChatMessage] = list(raw_messages)\n\n if not self._force_final_answer_enabled:\n return messages\n\n # If final-answer mode has been activated before, always include the prompt in model input.\n self._ensure_final_prompt(messages)\n\n token_estimate = _estimate_token_length(messages, self._model_name)\n if token_estimate < self._force_final_answer_upper_limit:\n return messages\n\n # We are going to mutate messages for context management: deep-copy first.\n messages = self._copy_messages(messages)\n\n self._handle_context_overflow(messages)\n self._ensure_final_prompt(messages)\n self._ensure_context_within_upper_limit(messages)\n return messages\n\n def _shrink_batch_search_results(self, messages: list[ChatMessage]) -> bool:\n \"\"\"Compress earliest uncompressed batch_search_result content (model-input only).\"\"\"\n for message in messages:\n content = getattr(message, \"content\", None)\n new_content, changed = self._compress_batch_search_in_content(content)\n if changed:\n message.content = new_content\n logger.info(\"@%s Compressed batch_search_results to save tokens\", self.name)\n return True\n return False\n\n @staticmethod\n def _parse_tool_call_arguments(raw_arguments: Any) -> dict[str, Any]:\n \"\"\"Safely parse tool call arguments JSON.\"\"\"\n if not isinstance(raw_arguments, str) or not raw_arguments.strip():\n return {}\n try:\n parsed = json.loads(raw_arguments)\n return parsed if isinstance(parsed, dict) else {}\n except Exception: # noqa: BLE001\n return {}\n\n def _is_search_tool_call(self, tool_name: str | None, tool_args: dict[str, Any]) -> bool:\n \"\"\"Identify whether a tool call is search-related.\"\"\"\n if not tool_name:\n return False\n lowered = tool_name.lower()\n if \"search\" in lowered:\n return True\n if lowered == \"batch_web_surfer\":\n action = tool_args.get(\"action\")\n if isinstance(action, str) and action.lower() == \"batch_search\":\n return True\n return False\n\n def _drop_oldest_tool_cycle(\n self,\n messages: list[ChatMessage],\n predicate: Callable[[str | None, dict[str, Any]], bool] | None = None,\n log_context: str = \"tool\",\n ) -> bool:\n \"\"\"Drop earliest tool call message plus all corresponding tool results.\"\"\"\n drop_indices: set[int] = set()\n for idx, message in enumerate(messages):\n tool_calls = getattr(message, \"tool_calls\", None)\n if not tool_calls:\n continue\n matched_any = False\n for tc in tool_calls:\n try:\n tool_name = tc.function.name\n raw_arguments = tc.function.arguments\n except Exception:\n tool_name = None\n raw_arguments = None\n parsed_args = self._parse_tool_call_arguments(raw_arguments)\n if predicate and not predicate(tool_name, parsed_args):\n continue\n matched_any = True\n\n if not matched_any:\n continue\n\n drop_indices.add(idx)\n\n # Find corresponding tool results for all tool calls in this message.\n for tc in tool_calls:\n tc_id = getattr(tc, \"id\", None)\n if not tc_id:\n continue\n for j in range(idx + 1, len(messages)):\n tool_msg = messages[j]\n if getattr(tool_msg, \"role\", None) != \"tool\":\n continue\n if getattr(tool_msg, \"tool_call_id\", None) == tc_id:\n drop_indices.add(j)\n break\n\n break\n\n if not drop_indices:\n return False\n\n for del_idx in sorted(drop_indices, reverse=True):\n messages.pop(del_idx)\n\n logger.warning(\n \"@%s Dropped earliest %s tool call/results to shrink context\", self.name, log_context\n )\n return True\n\n def _trim_oldest_messages(self, messages: list[ChatMessage]) -> bool:\n \"\"\"Drop oldest non-system messages until under threshold.\"\"\"\n if not messages or len(messages) <= 1:\n return False\n removed = False\n idx = 0\n while (\n idx < len(messages)\n and _estimate_token_length(messages, self._model_name)\n > self._force_final_answer_upper_limit\n ):\n if getattr(messages[idx], \"role\", None) == \"system\":\n idx += 1\n continue\n messages.pop(idx)\n removed = True\n if removed:\n logger.warning(\"@%s Trimmed oldest messages to satisfy context budget\", self.name)\n return removed\n\n def _ensure_context_within_upper_limit(self, messages: list[ChatMessage]) -> None:\n \"\"\"Ensure context is below the configured upper limit before forcing final answer.\"\"\"\n if not self._force_final_answer_enabled:\n return\n while True:\n token_estimate = _estimate_token_length(messages, self._model_name)\n if token_estimate <= self._force_final_answer_upper_limit:\n return\n if self._drop_oldest_tool_cycle(messages, log_context=\"any\"):\n continue\n if self._trim_oldest_messages(messages):\n continue\n break\n\n final_tokens = _estimate_token_length(messages, self._model_name)\n if final_tokens > self._force_final_answer_upper_limit:\n logger.warning(\n \"@%s Unable to trim context below upper limit (%s tokens remaining)\",\n self.name,\n final_tokens,\n )\n\n def _handle_context_overflow(self, messages: list[ChatMessage]) -> None:\n \"\"\"Enforce two-threshold hysteresis for search results before triggering final answer.\"\"\"\n if not self._force_final_answer_enabled:\n return\n\n token_estimate = _estimate_token_length(messages, self._model_name)\n if token_estimate < self._force_final_answer_upper_limit:\n return\n\n logger.warning(\n (\n \"@%s Context %s tokens reached upper limit %s; processing search tool payloads \"\n \"until below lower limit %s\"\n ),\n self.name,\n token_estimate,\n self._force_final_answer_upper_limit,\n self._force_final_answer_lower_limit,\n )\n\n while True:\n token_estimate = _estimate_token_length(messages, self._model_name)\n if token_estimate < self._force_final_answer_lower_limit:\n break\n\n if self._shrink_batch_search_results(messages):\n continue\n\n if self._drop_oldest_tool_cycle(messages, self._is_search_tool_call, \"search\"):\n continue\n\n break\n\n final_tokens = _estimate_token_length(messages, self._model_name)\n if final_tokens < self._force_final_answer_lower_limit:\n return\n\n logger.warning(\n (\n \"@%s Exhausted search tool cleanup but context still %s tokens (>= lower limit %s); \"\n \"forcing final answer workflow\"\n ),\n self.name,\n final_tokens,\n self._force_final_answer_lower_limit,\n )\n\n if not self._force_prompt_inserted:\n self._force_prompt_inserted = True\n logger.info(\"@%s Final answer prompt activated\", self.name)\n\n async def _run(\n self,\n messages: list[ChatMessage] | InputChannel[ChatMessage],\n additional_kwargs: dict | None = None,\n ) -> AsyncGenerator[AgentResponse, None]:\n \"\"\"\n Run agent, calls step() method in a loop\n\n Args:\n messages: Input message list or input channel\n additional_kwargs: Additional parameters\n\n Yields:\n AgentResponse: Agent response object\n \"\"\"\n if additional_kwargs is None:\n additional_kwargs = {}\n\n input_messages = []\n # Handle input messages\n if isinstance(messages, list):\n input_messages = messages\n elif isinstance(messages, InputChannel):\n input_messages = await messages.get()\n\n self.current_round = 0\n # Initialize history messages as member variable\n self.context.add(input_messages)\n should_stop = False\n\n # Loop calling step() until stop or reach max steps\n while self.current_round < self.max_steps:\n with create_span(\n name=f\"@{self.name} Round {self.current_round}/{self.max_steps}\"\n ):\n if should_stop and not self.config.unfinished_mode:\n break\n if should_stop:\n if isinstance(messages, InputChannel):\n input_messages = await messages.get()\n self.context.add(input_messages)\n should_stop = False\n else:\n if isinstance(messages, InputChannel):\n input_messages = await messages.get_no_wait()\n self.context.add(input_messages)\n\n self.current_round += 1\n logger.info(f\"@{self.name} Round {self.current_round}/{self.max_steps}\")\n\n try:\n # Call step() method (now an async generator)\n last_response = None\n model_messages = self._prepare_messages_for_model()\n async for response in self._step(model_messages, additional_kwargs):\n # Update history messages (using member variable)\n # Only add complete messages to history (has role field and not STREAM type)\n logger.info(f\"@{self.name} Response: {response}\")\n if response is None:\n continue\n if response.message:\n # STREAM type messages are incremental updates, should not be added to history\n if response.message_type == AgentMessageType.STREAM.value:\n # Skip streaming incremental messages, they will be handled in accumulated messages\n pass\n else:\n # For non-streaming messages, ensure role field exists\n message_to_add = None\n if isinstance(response.message, ChatMessage):\n if (\n hasattr(response.message, \"role\")\n and response.message.role\n ):\n message_to_add = response.message\n elif isinstance(response.message, dict):\n if response.message.get(\"role\"):\n message_to_add = ChatMessage(**response.message)\n\n # Only add message to history when it has a valid role\n if message_to_add and message_to_add.role:\n self.context.add([message_to_add])\n else:\n logger.warning(\n f\"@{self.name} Skipping message without role: {response.message_type}\"\n )\n\n # Set metadata (only includes round, history_messages uses member variable)\n if response.metadata is None:\n response.metadata = {}\n response.metadata[\"step_count\"] = self.current_round\n\n # Return response\n yield response\n last_response = response\n\n # Check if error occurred\n if response.status == AgentRunningStatus.ERROR.value:\n logger.error(\n f\"@{self.name} Error at round {self.current_round}: {response.error_msg}\"\n )\n should_stop = True\n break\n\n # Check last response status, decide whether to stop\n if last_response:\n if last_response.status == AgentRunningStatus.FINISHED.value:\n should_stop = True\n logger.info(\n f\"@{self.name} Finished at round {self.current_round}\"\n )\n elif last_response.status == AgentRunningStatus.ERROR.value:\n should_stop = True\n elif last_response.status == AgentRunningStatus.RUNNING.value:\n # If still running, may need to continue (e.g., has tool calls)\n # Here can decide whether to stop based on actual situation\n pass\n\n except Exception as e:\n err_text = str(e) or repr(e)\n logger.error(\n f\"@{self.name} Exception at round {self.current_round}: {err_text}\",\n exc_info=True,\n )\n error_response = AgentResponse(\n status=AgentRunningStatus.ERROR.value,\n error_msg=err_text,\n metadata={\n \"step_count\": self.current_round,\n },\n )\n yield error_response\n should_stop = True\n break\n\n # If max steps reached and not stopped yet\n if self.current_round >= self.max_steps and not should_stop:\n logger.warning(\n f\"@{self.name} Reached max steps ({self.max_steps}) without stopping\"\n )\n # Generate final response\n final_response = AgentResponse(\n status=AgentRunningStatus.STOPPED.value,\n metadata={\n \"step_count\": self.current_round,\n },\n )\n yield final_response\n\n @abstractmethod\n async def _step(\n self,\n messages: list[ChatMessage],\n additional_kwargs: dict | None = None,\n ) -> AsyncGenerator[AgentResponse, None]:\n \"\"\"\n Execute a single step, subclasses must implement this method.\n\n Args:\n messages: Current message history.\n additional_kwargs: Additional parameters.\n\n Yields:\n AgentResponse: Response for current step, can yield multiple responses.\n - Last response with status FINISHED indicates completion.\n - Status ERROR indicates an error occurred.\n - Status RUNNING indicates continue execution.\n \"\"\"\n raise NotImplementedError(\"Subclasses must implement this method\")\n" }, { "path": "cortex/agents/checkpoint_agent/checkpoint_agent.py", "content": "import logging\nfrom abc import abstractmethod\nfrom typing import AsyncGenerator\n\nfrom cortex.model.definition import ChatMessage, ContentBlockType\nfrom pydantic import BaseModel\n\nfrom cortex.agents.base_agent import BaseAgent\nfrom cortex.agents.checkpoint_agent.checkpointer import CheckPointer, CheckpointStorage\nfrom cortex.agents.types import (\n AgentConfig,\n AgentMessageType,\n AgentResponse,\n AgentRunningStatus,\n)\nfrom cortex.model.provider import ModelProvider\nfrom cortex.tools.toolset import ToolSet\n\nlogger = logging.getLogger(__name__)\n\n\nclass PendingToolCall(BaseModel):\n \"\"\"Pending tool call\"\"\"\n\n request: ChatMessage | None\n results: dict[str, ChatMessage] | None\n\n\nclass CheckpointState(BaseModel):\n \"\"\"Checkpoint state definition\"\"\"\n\n messages: list[ChatMessage] | None\n pending_tool_calls: list[PendingToolCall] | None\n tool_call_results: list[ChatMessage] | None\n config: AgentConfig | None\n current_step: int\n max_steps: int | None\n finished: bool\n error: str | None\n\n\nclass CheckpointAgent(BaseAgent):\n \"\"\"Checkpoint-based single-step execution Agent (without LangGraph)\"\"\"\n\n def __init__(\n self,\n config: AgentConfig,\n storage: CheckpointStorage,\n provider: ModelProvider | None = None,\n toolset: ToolSet | None = None,\n thread_id: str | None = None,\n ):\n super().__init__(config=config, toolset=toolset, provider=provider)\n self.thread_id = thread_id or f\"{self.name}_main\"\n self.storage = storage\n self._state: CheckpointState | None = None\n\n def _init_state(\n self,\n ) -> CheckpointState:\n \"\"\"Initialize state\"\"\"\n return CheckpointState(\n messages=[],\n pending_tool_calls=[],\n tool_call_results=[],\n config=self.config,\n current_step=0,\n max_steps=getattr(self.config, \"max_steps\", 10),\n finished=False,\n error=None,\n )\n\n async def _process_tool_call_results(\n self, state: CheckpointState\n ) -> AsyncGenerator[AgentResponse, None]:\n \"\"\"Process tool call results\"\"\"\n if not state.tool_call_results or len(state.tool_call_results) == 0:\n return\n\n logger.debug(\n \"@%s Processing %d tool call results\", self.name, len(state.tool_call_results)\n )\n async for response in self._tool_call_handler(\n state.messages, state.tool_call_results\n ):\n yield response\n if response.message:\n state.messages.append(response.message)\n\n # Clear processed results\n state.tool_call_results = []\n\n async def _execute_step(\n self, state: CheckpointState, additional_kwargs: dict | None = None\n ) -> AsyncGenerator[AgentResponse, None]:\n \"\"\"Execute single step\"\"\"\n step_responses = []\n try:\n async for response in self._step(state.messages, additional_kwargs):\n step_responses.append(response)\n yield response\n\n # Handle message history update\n if response.message:\n # STREAM type messages are incremental updates, should not be added to history\n if response.message_type == AgentMessageType.STREAM.value:\n pass\n else:\n # For non-streaming messages, ensure role field exists\n message_to_add = None\n if isinstance(response.message, ChatMessage):\n if (\n response.message.tool_calls\n and len(response.message.tool_calls) > 0\n ):\n # Execute tool calls\n results = await self.run_tool_call(response.message)\n # Convert result list to dict with tool_call_id as key\n state.pending_tool_calls.append(\n PendingToolCall(\n request=response.message,\n results={\n result.tool_call_id: result\n for result in results\n },\n )\n )\n message_to_add = response.message\n elif response.message.tool_call_id:\n # Tool call result\n state.tool_call_results.append(response.message)\n else:\n message_to_add = response.message\n\n elif isinstance(response.message, dict):\n if response.message.get(\"role\"):\n message_to_add = ChatMessage(**response.message)\n\n # Only add message to history when it has a valid role\n if message_to_add and message_to_add.role:\n state.messages.append(message_to_add)\n else:\n logger.warning(\n \"@%s Skipping message without role: %s\",\n self.name,\n response.message_type,\n )\n\n state.current_step += 1\n # Check if finished\n if step_responses:\n last_response = step_responses[-1]\n if last_response.status == AgentRunningStatus.FINISHED.value:\n state.finished = True\n elif last_response.status == AgentRunningStatus.ERROR.value:\n state.error = last_response.error_msg\n state.finished = True\n\n except Exception as e:\n logger.error(\"@%s Error during step execution: %s\", self.name, str(e))\n state.error = str(e)\n state.finished = True\n raise\n\n async def _update_client_tool_results(self, state: CheckpointState) -> bool:\n \"\"\"\n Update client tool call results.\n Returns True if there are still pending tool calls.\n \"\"\"\n if not state.pending_tool_calls or len(state.pending_tool_calls) == 0:\n return False\n\n new_pending_tool_calls = []\n for pending_item in state.pending_tool_calls:\n request = pending_item.request\n results = pending_item.results\n # Check if each tool_call has a corresponding result\n all_matched = True\n\n for tool_call in request.tool_calls:\n tool_call_id = tool_call.id\n # Check if result already exists\n if results.get(tool_call_id):\n continue\n\n result_content = self.toolset().get_client_tool_call_result(\n tool_call_id\n )\n\n if result_content is None:\n all_matched = False\n continue\n results[tool_call_id] = ChatMessage(\n role=\"tool\",\n content=[\n {\n \"type\": ContentBlockType.TEXT.value,\n ContentBlockType.TEXT.value: str(result_content),\n }\n ],\n tool_call_id=tool_call_id,\n )\n\n if all_matched:\n # All tool_calls have corresponding results, move results to tool_call_results\n state.tool_call_results.extend(results.values())\n else:\n # Still have incomplete tool calls, update results and keep in pending\n pending_item.results = results\n new_pending_tool_calls.append(pending_item)\n\n state.pending_tool_calls = new_pending_tool_calls\n return len(new_pending_tool_calls) > 0\n\n def _should_continue(self, state: CheckpointState) -> bool:\n \"\"\"Determine whether execution should continue\"\"\"\n\n if len(state.pending_tool_calls) > 0:\n return False # Has pending tool calls, need to wait\n if len(state.tool_call_results) > 0:\n return True\n if state.finished:\n return False # Already finished\n if state.error:\n return False # Has error\n if state.current_step >= state.max_steps:\n return False # Reached max steps\n return True\n\n async def _run(\n self,\n messages: list[ChatMessage],\n additional_kwargs: dict | None = None,\n ) -> AsyncGenerator[AgentResponse, None]:\n \"\"\"\n Checkpoint-based run method (without LangGraph)\n\n Args:\n messages: Input message list or input channel\n additional_kwargs: Additional parameters\n\n Yields:\n AgentResponse: Agent response object\n \"\"\"\n # Initialize or load state\n async with CheckPointer[CheckpointState](\n self.thread_id, self.storage, self._init_state(), CheckpointState\n ) as state:\n state.messages.extend(messages)\n\n # Main execution loop\n while True:\n # 1. Update client tool call results and check if there are pending tool calls\n has_pending = await self._update_client_tool_results(state)\n if has_pending:\n logger.info(\"@%s Has pending tool calls, waiting for client response\", self.name)\n return\n\n # 2. Process completed tool call results\n if len(state.tool_call_results) > 0:\n async for response in self._process_tool_call_results(state):\n yield response\n\n # 3. Check whether execution should continue\n if not self._should_continue(state):\n logger.info(\"@%s Execution completed\", self.name)\n break\n\n # 4. Execute next step\n logger.info(\"@%s Executing step %d\", self.name, state.current_step + 1)\n async for response in self._execute_step(state, additional_kwargs):\n yield response\n\n @abstractmethod\n async def _step(\n self,\n messages: list[ChatMessage],\n additional_kwargs: dict | None = None,\n ) -> AsyncGenerator[AgentResponse, None]:\n \"\"\"\n Execute a single step, subclasses must implement this method.\n\n Args:\n messages: Current message history.\n additional_kwargs: Additional parameters.\n\n Yields:\n AgentResponse: Response for current step, can yield multiple responses.\n - Last response with status FINISHED indicates completion.\n - Status ERROR indicates an error occurred.\n - Status RUNNING indicates continue execution.\n \"\"\"\n raise NotImplementedError(\"Subclasses must implement this method\")\n\n @abstractmethod\n async def _tool_call_handler(\n self,\n messages: list[ChatMessage],\n tool_calls: list[ChatMessage],\n ) -> AsyncGenerator[AgentResponse, None]:\n \"\"\"\n Handle tool calls, subclasses must implement this method.\n\n Args:\n messages: Current message history.\n tool_calls: List of tool call messages.\n\n Yields:\n AgentResponse: Response for current step, can yield multiple responses.\n - Last response with status FINISHED indicates completion.\n - Status ERROR indicates an error occurred.\n - Status RUNNING indicates continue execution.\n \"\"\"\n raise NotImplementedError(\"Subclasses must implement this method\")\n" }, { "path": "cortex/agents/checkpoint_agent/checkpointer.py", "content": "\"\"\"Checkpoint storage implementation.\"\"\"\n\nimport json\nimport logging\nimport os\nfrom pathlib import Path\nfrom typing import Any\n\nlogger = logging.getLogger(__name__)\n\nfrom typing import Generic, TypeVar\n\nfrom pydantic import BaseModel\n\nT = TypeVar(\"T\", bound=BaseModel)\n\n\nclass CheckpointStorage:\n def save_state(self, checkpoint_id: str, state: dict[str, Any]):\n \"\"\"Save state.\"\"\"\n raise NotImplementedError(\"Subclass must implement this method\")\n\n def load_state(self, checkpoint_id: str) -> dict[str, Any] | None:\n \"\"\"Load state from storage.\"\"\"\n raise NotImplementedError(\"Subclass must implement this method\")\n\n def delete_state(self, checkpoint_id: str):\n \"\"\"Delete the specified checkpoint.\"\"\"\n raise NotImplementedError(\"Subclass must implement this method\")\n\n async def asave_state(self, checkpoint_id: str, state: dict[str, Any]):\n \"\"\"Save state asynchronously.\"\"\"\n raise NotImplementedError(\"Subclass must implement this method\")\n\n async def aload_state(self, checkpoint_id: str) -> dict[str, Any] | None:\n \"\"\"Load state from storage asynchronously.\"\"\"\n raise NotImplementedError(\"Subclass must implement this method\")\n\n async def adelete_state(self, checkpoint_id: str):\n \"\"\"Delete the specified checkpoint asynchronously.\"\"\"\n raise NotImplementedError(\"Subclass must implement this method\")\n\n\nclass CheckPointer(Generic[T]):\n \"\"\"Checkpoint manager base class.\"\"\"\n\n def __init__(\n self,\n checkpoint_id: str,\n storage: CheckpointStorage,\n init_state: T | None,\n state_type: type[T],\n ):\n \"\"\"\n Initialize CheckPointer.\n\n Args:\n checkpoint_id: Checkpoint ID for identifying and loading specific checkpoints\n storage: Storage backend\n init_state: Initial state\n state_type: State type, must be a class inheriting from BaseModel\n \"\"\"\n self.checkpoint_id = checkpoint_id\n self.storage = storage\n self.init_state = init_state\n self.state_type = state_type\n self._state: T | None = None\n\n def __enter__(self) -> T:\n state_dict = self.storage.load_state(self.checkpoint_id)\n if state_dict is None:\n self._state = self.init_state\n else:\n # Convert dict to pydantic model\n self._state = self.state_type.model_validate(state_dict)\n return self._state\n\n def __exit__(self, exc_type, exc_value, traceback):\n # Auto-save state on exit\n if self._state is not None and self.checkpoint_id is not None:\n try:\n # Use pydantic's model_dump method to convert to dict\n state_dict = self._state.model_dump()\n self.storage.save_state(self.checkpoint_id, state_dict)\n logger.debug(\"Auto-saved state on exit: checkpoint_id=%s\", self.checkpoint_id)\n except Exception as e:\n logger.error(\"Failed to save state on exit: %s\", str(e))\n # todo handle cancel error\n return False\n\n async def __aenter__(self):\n return self.__enter__()\n\n async def __aexit__(self, exc_type, exc_value, traceback):\n return self.__exit__(exc_type, exc_value, traceback)\n\n\nclass MemoryCheckPointer(CheckpointStorage):\n \"\"\"Memory-based Checkpoint storage.\"\"\"\n\n def __init__(self):\n self._storage: dict[str, dict[str, Any]] = {}\n\n def save_state(self, checkpoint_id: str, state: dict[str, Any]):\n \"\"\"Save state to memory.\"\"\"\n self._storage[checkpoint_id] = state\n logger.debug(\"State saved to memory: checkpoint_id=%s\", checkpoint_id)\n\n def load_state(self, checkpoint_id: str) -> dict[str, Any] | None:\n \"\"\"Load state from memory.\"\"\"\n state = self._storage.get(checkpoint_id)\n if state:\n logger.debug(\"State loaded from memory: checkpoint_id=%s\", checkpoint_id)\n return state\n\n def delete_state(self, checkpoint_id: str):\n \"\"\"Delete the specified checkpoint.\"\"\"\n if checkpoint_id in self._storage:\n del self._storage[checkpoint_id]\n logger.debug(\"Deleted checkpoint from memory: checkpoint_id=%s\", checkpoint_id)\n\n async def asave_state(self, checkpoint_id: str, state: dict[str, Any]):\n \"\"\"Save state to memory.\"\"\"\n self.save_state(checkpoint_id, state)\n\n async def aload_state(self, checkpoint_id: str) -> dict[str, Any] | None:\n \"\"\"Load state from memory.\"\"\"\n return self.load_state(checkpoint_id)\n\n async def adelete_state(self, checkpoint_id: str):\n \"\"\"Delete the specified checkpoint.\"\"\"\n self.delete_state(checkpoint_id)\n\n\nclass FileCheckPointer(CheckpointStorage):\n \"\"\"File-based Checkpoint storage.\"\"\"\n\n def __init__(self, checkpoint_dir: str):\n \"\"\"\n Initialize file storage.\n\n Args:\n checkpoint_dir: Checkpoint storage directory\n \"\"\"\n self.checkpoint_dir = Path(checkpoint_dir)\n\n # Ensure directory exists\n self.checkpoint_dir.mkdir(parents=True, exist_ok=True)\n\n def _get_checkpoint_file(self, checkpoint_id: str) -> Path:\n \"\"\"Get checkpoint file path.\"\"\"\n return self.checkpoint_dir / f\"{checkpoint_id}.json\"\n\n def save_state(self, checkpoint_id: str, state: dict[str, Any]):\n \"\"\"Save state to file.\"\"\"\n checkpoint_file = self._get_checkpoint_file(checkpoint_id)\n try:\n # Write to file\n with open(checkpoint_file, \"w\", encoding=\"utf-8\") as f:\n json.dump(state, f, indent=2, ensure_ascii=False)\n\n logger.debug(\"State saved to file: %s\", checkpoint_file)\n\n except Exception as e:\n logger.error(\"Failed to save state: %s\", str(e))\n raise\n\n def load_state(self, checkpoint_id: str) -> dict[str, Any] | None:\n \"\"\"Load state from file.\"\"\"\n checkpoint_file = self._get_checkpoint_file(checkpoint_id)\n if not checkpoint_file.exists():\n logger.debug(\"Checkpoint file not found: %s\", checkpoint_file)\n return None\n\n try:\n with open(checkpoint_file, \"r\", encoding=\"utf-8\") as f:\n state_dict = json.load(f)\n\n logger.debug(\"State loaded from file: %s\", checkpoint_file)\n return state_dict\n\n except (OSError, json.JSONDecodeError, ValueError) as e:\n logger.error(\"Failed to load state: %s\", str(e))\n return None\n\n def delete_state(self, checkpoint_id: str):\n \"\"\"Delete checkpoint file.\"\"\"\n checkpoint_file = self._get_checkpoint_file(checkpoint_id)\n if checkpoint_file.exists():\n os.remove(checkpoint_file)\n logger.debug(\"Deleted checkpoint file: %s\", checkpoint_file)\n\n async def asave_state(self, checkpoint_id: str, state: dict[str, Any]):\n \"\"\"Save state to file.\"\"\"\n self.save_state(checkpoint_id, state)\n\n async def aload_state(self, checkpoint_id: str) -> dict[str, Any] | None:\n \"\"\"Load state from file.\"\"\"\n return self.load_state(checkpoint_id)\n\n async def adelete_state(self, checkpoint_id: str):\n \"\"\"Delete checkpoint file.\"\"\"\n self.delete_state(checkpoint_id)\n\n\nclass SqliteCheckPointer(CheckpointStorage):\n \"\"\"SQLite-based Checkpoint storage.\"\"\"\n\n def __init__(self, db_path: str):\n \"\"\"\n Initialize SQLite storage.\n\n Args:\n db_path: SQLite database file path\n \"\"\"\n import sqlite3\n\n self.db_path = db_path\n self.conn: sqlite3.Connection | None = None\n\n # Ensure database directory exists\n db_dir = Path(db_path).parent\n db_dir.mkdir(parents=True, exist_ok=True)\n\n # Initialize database table\n self._init_db()\n\n def _init_db(self):\n \"\"\"Initialize database table.\"\"\"\n import sqlite3\n\n conn = sqlite3.connect(self.db_path)\n cursor = conn.cursor()\n\n cursor.execute(\n \"\"\"\n CREATE TABLE IF NOT EXISTS checkpoints (\n checkpoint_id TEXT PRIMARY KEY,\n state TEXT NOT NULL,\n updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\n )\n \"\"\"\n )\n\n conn.commit()\n conn.close()\n\n def __enter__(self):\n import sqlite3\n\n self.conn = sqlite3.connect(self.db_path)\n return self\n\n def __exit__(self, exc_type, exc_value, traceback):\n if self.conn:\n self.conn.close()\n self.conn = None\n return False\n\n async def __aenter__(self):\n return self.__enter__()\n\n async def __aexit__(self, exc_type, exc_value, traceback):\n return self.__exit__(exc_type, exc_value, traceback)\n\n def save_state(self, checkpoint_id: str, state: dict[str, Any]):\n \"\"\"Save state to SQLite.\"\"\"\n import sqlite3\n\n conn = sqlite3.connect(self.db_path)\n try:\n # Convert to JSON string\n state_json = json.dumps(state, ensure_ascii=False)\n\n # Insert or update\n cursor = conn.cursor()\n cursor.execute(\n \"\"\"\n INSERT OR REPLACE INTO checkpoints (checkpoint_id, state, updated_at)\n VALUES (?, ?, CURRENT_TIMESTAMP)\n \"\"\",\n (checkpoint_id, state_json),\n )\n\n conn.commit()\n logger.debug(\"State saved to SQLite: checkpoint_id=%s\", checkpoint_id)\n\n except Exception as e:\n logger.error(\"Failed to save state: %s\", str(e))\n raise\n finally:\n conn.close()\n\n def load_state(self, checkpoint_id: str) -> dict[str, Any] | None:\n \"\"\"Load state from SQLite.\"\"\"\n import sqlite3\n\n conn = sqlite3.connect(self.db_path)\n try:\n cursor = conn.cursor()\n cursor.execute(\n \"SELECT state FROM checkpoints WHERE checkpoint_id = ?\",\n (checkpoint_id,),\n )\n\n row = cursor.fetchone()\n if not row:\n logger.debug(\"Checkpoint not found: checkpoint_id=%s\", checkpoint_id)\n return None\n\n state_json = row[0]\n state_dict = json.loads(state_json)\n\n logger.debug(\"State loaded from SQLite: checkpoint_id=%s\", checkpoint_id)\n return state_dict\n\n except (json.JSONDecodeError, ValueError) as e:\n logger.error(\"Failed to load state: %s\", str(e))\n return None\n finally:\n conn.close()\n\n def delete_state(self, checkpoint_id: str):\n \"\"\"Delete checkpoint.\"\"\"\n import sqlite3\n\n conn = sqlite3.connect(self.db_path)\n try:\n cursor = conn.cursor()\n cursor.execute(\n \"DELETE FROM checkpoints WHERE checkpoint_id = ?\", (checkpoint_id,)\n )\n conn.commit()\n logger.debug(\"Deleted checkpoint: checkpoint_id=%s\", checkpoint_id)\n finally:\n conn.close()\n\n async def asave_state(self, checkpoint_id: str, state: dict[str, Any]):\n \"\"\"Save state to SQLite.\"\"\"\n self.save_state(checkpoint_id, state)\n\n async def aload_state(self, checkpoint_id: str) -> dict[str, Any] | None:\n \"\"\"Load state from SQLite.\"\"\"\n return self.load_state(checkpoint_id)\n\n async def adelete_state(self, checkpoint_id: str):\n \"\"\"Delete checkpoint.\"\"\"\n self.delete_state(checkpoint_id)\n" }, { "path": "cortex/agents/checkpoint_agent/react_agent.py", "content": "\"\"\"ReActAgent - An Agent specifically designed to execute tasks in ReAct (Reasoning + Acting) mode, capable of calling tools to complete tasks.\"\"\"\n\nimport logging\nfrom typing import AsyncGenerator\nfrom uuid import uuid4\n\nfrom cortex.model.definition import ChatMessage\n\nfrom cortex.agents.checkpoint_agent.checkpoint_agent import CheckpointAgent\nfrom cortex.agents.checkpoint_agent.checkpointer import CheckpointStorage\nfrom cortex.agents.react_agent import process_messages\nfrom cortex.agents.types import (\n AgentConfig,\n AgentMessageType,\n AgentResponse,\n AgentRunningStatus,\n)\nfrom cortex.model.provider import ModelProvider\nfrom cortex.tools.toolset import ToolSet\n\nlogger = logging.getLogger(__name__)\n\n\nclass CheckpointReActAgent(CheckpointAgent):\n \"\"\"ReActAgent - An Agent specifically designed to execute tasks in ReAct (Reasoning + Acting) mode.\n\n Features:\n - Can call tools to complete tasks\n - Can call tools multiple times to complete complex tasks\n - Can provide detailed task execution process description\n - Can provide detailed task execution results\n \"\"\"\n\n def __init__(\n self,\n storage: CheckpointStorage,\n context_id: str | None = None,\n provider: ModelProvider | None = None,\n config: AgentConfig | None = None,\n toolset: ToolSet | None = None,\n ):\n # If no toolset is provided, create a default math toolset\n if toolset is None:\n # Note: Cannot call async functions directly here, need to initialize externally\n raise ValueError(\n \"ReActAgent requires a toolset, please use init_react_tools() to create one\"\n )\n\n if not context_id:\n context_id = uuid4().hex\n super().__init__(\n storage=storage,\n thread_id=context_id,\n provider=provider,\n config=config,\n toolset=toolset,\n )\n\n async def _step(\n self,\n messages: list[ChatMessage],\n additional_kwargs: dict | None = None,\n ) -> AsyncGenerator[AgentResponse, None]:\n \"\"\"\n Execute a single step, can yield multiple responses.\n\n Args:\n messages: Current message history\n additional_kwargs: Additional parameters\n\n Yields:\n AgentResponse: Response for the current step\n \"\"\"\n try:\n async for response_message in process_messages(\n self.system_prompt,\n messages,\n self.toolset(),\n self.model_api(),\n getattr(self.model, \"infer_kwargs\", {}).get(\"stream\", False),\n trace_messages=list(messages) if messages else [],\n ):\n yield response_message\n except Exception as e:\n err_text = str(e) or repr(e)\n logger.error(\"@%s execution error: %s\", self.name, err_text, exc_info=True)\n error_response = AgentResponse(\n message=None,\n status=AgentRunningStatus.ERROR.value,\n error_msg=err_text,\n message_type=AgentMessageType.FINAL.value,\n )\n yield error_response\n\n async def _tool_call_handler(\n self, messages: list[ChatMessage], tool_results: list[ChatMessage]\n ) -> AsyncGenerator[AgentResponse, None]:\n \"\"\"\n Handle tool calls.\n \"\"\"\n\n for result in tool_results:\n yield AgentResponse(\n message=result,\n status=AgentRunningStatus.RUNNING.value,\n message_type=AgentMessageType.ACCUMULATED.value,\n )\n" }, { "path": "cortex/agents/input/input.py", "content": "import asyncio\nimport logging\nfrom typing import Generic, TypeVar\n\nfrom pydantic import BaseModel\n\nlogger = logging.getLogger(__name__)\n\nT = TypeVar(\"T\", bound=BaseModel)\n\n\nclass InputChannel(Generic[T]):\n queue: asyncio.Queue[T]\n\n def __init__(self, queue: asyncio.Queue[T]) -> None:\n self.queue = queue\n\n async def get(self) -> list[T]:\n \"\"\"Blocks until at least one message is available, then returns all available messages.\"\"\"\n logger.debug(\"InputChannel waiting for first data\")\n first = await self.queue.get()\n data_list: list[T] = [first]\n count = 1\n while True:\n try:\n data = self.queue.get_nowait()\n data_list.append(data)\n count += 1\n except asyncio.QueueEmpty:\n break\n logger.debug(f\"InputChannel returning {count} data(s)\")\n return data_list\n\n async def get_no_wait(self) -> list[T]:\n data_list: list[T] = []\n count = 0\n while True:\n try:\n data = self.queue.get_nowait()\n data_list.append(data)\n count += 1\n except asyncio.QueueEmpty:\n break\n logger.debug(f\"InputChannel returning {count} data(s)\")\n return data_list\n" }, { "path": "cortex/agents/react_agent.py", "content": "\"\"\"ReActAgent - Specialized Agent for executing ReAct pattern, capable of calling tools to complete tasks.\"\"\"\n\nimport logging\nfrom typing import AsyncGenerator\nfrom uuid import uuid4\n\nfrom cortex.agents.base_agent import BaseAgent\nfrom cortex.agents.base_step_agent import BaseStepAgent\nfrom cortex.agents.types import (\n AgentConfig,\n AgentMessageType,\n AgentResponse,\n AgentRunningStatus,\n)\nfrom cortex.context import BaseContext\nfrom cortex.context.simple_context import SimpleContext\nfrom cortex.model import ChatMessage, MessageType, ModelAPI\nfrom cortex.model.provider import ModelProvider\nfrom cortex.tools.toolset import ToolSet\n\nlogger = logging.getLogger(__name__)\n\n\ndef _check_if_finished(response_message: ChatMessage | None) -> bool:\n \"\"\"Check if task is finished based on model output.\"\"\"\n if not response_message:\n return False\n\n # If there are tool calls, continue execution\n if BaseAgent.has_tool_call(response_message):\n return False\n\n return True\n\n\nasync def process_messages(\n system_prompt: str | None,\n messages: list[ChatMessage],\n toolset: ToolSet,\n model_api: ModelAPI,\n use_stream: bool,\n trace_messages: list[ChatMessage] | None = None,\n) -> AsyncGenerator[AgentResponse, None]:\n if not system_prompt:\n system_prompt = \"\"\"You are a professional task execution assistant capable of calling tools to complete tasks.\nYour task is:\n1. Understand the task proposed by the user\n2. Use the provided tools to complete the task\n3. Call tools multiple times to complete complex tasks\n4. Provide detailed explanations of task execution process\n5. Provide detailed task execution results\n\nPlease ensure:\n- For complex tasks, use tools step by step\n- Provide clear explanations of task execution steps\n- Verify the correctness of task execution results\n- When the task is complete, clearly state \"Task completed\"\n\"\"\"\n\n # Prepare message list (only insert system if none present, to avoid duplicate when runner/caller already added one)\n infer_messages = messages.copy()\n if system_prompt and not any(msg.role == \"system\" for msg in infer_messages):\n infer_messages.insert(0, ChatMessage(role=\"system\", content=system_prompt))\n\n # Get all tool schemas\n tool_schemas = toolset.get_all_schemas()\n\n # Convert tool schemas to model-usable format\n tools_for_model = []\n for tool_name, schema in tool_schemas.items():\n tools_for_model.append(\n {\n \"type\": \"function\",\n \"function\": {\n \"name\": tool_name,\n \"description\": schema.description,\n \"parameters\": schema.parameters,\n },\n }\n )\n\n # Call model\n trace_request = None\n if trace_messages is not None:\n trace_infer_messages = trace_messages.copy()\n if system_prompt and not any(msg.role == \"system\" for msg in trace_infer_messages):\n trace_infer_messages.insert(\n 0, ChatMessage(role=\"system\", content=system_prompt)\n )\n trace_request = {\n \"messages\": trace_infer_messages,\n \"sent_messages\": infer_messages,\n \"tools\": tools_for_model if tools_for_model else None,\n }\n\n if use_stream:\n # Streaming output mode\n delta_count = 0\n response_message = None\n\n async for model_msg in model_api.chat_completion_stream(\n messages=infer_messages,\n tools=tools_for_model if tools_for_model else None,\n trace_request=trace_request,\n ):\n delta_count += 1\n event = model_msg.message\n response_message = (\n event # Save last message (model yields accumulated_message at end)\n )\n\n if model_msg.message_type == MessageType.DELTA:\n # This is a delta event, yield directly\n delta_response = AgentResponse(\n message=event,\n status=AgentRunningStatus.RUNNING.value,\n message_type=AgentMessageType.STREAM.value,\n )\n yield delta_response\n\n else:\n # Non-streaming output mode\n model_msg = await model_api.chat_completion(\n messages=infer_messages,\n tools=tools_for_model if tools_for_model else None,\n trace_request=trace_request,\n )\n response_message = model_msg.message\n\n is_finished = _check_if_finished(response_message)\n response_status = (\n AgentRunningStatus.FINISHED.value\n if is_finished\n else AgentRunningStatus.RUNNING.value\n )\n message_type = (\n AgentMessageType.FINAL.value\n if is_finished\n else AgentMessageType.ACCUMULATED.value\n )\n\n model_response = AgentResponse(\n message=response_message,\n status=response_status,\n message_type=message_type,\n )\n yield model_response\n\n\nclass ReActAgent(BaseStepAgent):\n \"\"\"ReActAgent - Specialized Agent for executing ReAct pattern, capable of calling tools to complete tasks.\n\n Features:\n - Can call tools to complete tasks\n - Can call tools multiple times to complete complex tasks\n - Can provide detailed explanations of task execution process\n - Can provide detailed task execution results\n \"\"\"\n\n def __init__(\n self,\n context: BaseContext | None = None,\n provider: ModelProvider | None = None,\n config: AgentConfig | None = None,\n toolset: ToolSet | None = None,\n ):\n if context is None:\n context = SimpleContext(uuid4().hex)\n # If toolset is not provided, create default math toolset\n if toolset is None:\n # Note: Cannot call async functions directly here, need to initialize externally\n raise ValueError(\n \"ReActAgent requires a toolset, please use init_react_tools() to create one\"\n )\n\n super().__init__(\n context=context, provider=provider, config=config, toolset=toolset\n )\n\n async def _step(\n self,\n messages: list[ChatMessage],\n additional_kwargs: dict | None = None,\n ) -> AsyncGenerator[AgentResponse, None]:\n \"\"\"\n Execute single step operation, can yield multiple responses.\n\n Args:\n messages: Current message history\n additional_kwargs: Additional parameters\n\n Yields:\n AgentResponse: Response for current step\n \"\"\"\n\n trace_messages: list[ChatMessage] | None = None\n try:\n trace_messages = list(self.context.get_all())\n except Exception: # noqa: BLE001\n trace_messages = None\n\n if trace_messages is not None and self._force_final_answer_enabled and self._force_prompt_inserted:\n prompt_message = self._make_force_final_answer_message()\n if not trace_messages or not (\n getattr(trace_messages[-1], \"role\", None) == prompt_message.role\n and getattr(trace_messages[-1], \"content\", None) == prompt_message.content\n ):\n trace_messages.append(prompt_message)\n\n async for response_message in process_messages(\n self.system_prompt,\n messages,\n self.toolset(),\n self.model_api(),\n getattr(self.model, \"infer_kwargs\", {}).get(\"stream\", False),\n trace_messages=trace_messages,\n ):\n try:\n yield response_message\n\n # Only check and execute tool calls for non-STREAM types (i.e., ACCUMULATED or FINAL)\n # In streaming output, STREAM type delta messages don't contain complete tool_call\n if response_message.message_type == AgentMessageType.STREAM.value:\n continue\n\n # Check for tool calls and execute\n tool_result_messages = []\n message = response_message.message\n if self.has_tool_call(message):\n tool_result_messages = await self.run_tool_call(message)\n if tool_result_messages:\n logger.info(\n \"@%s Detected %s tool call results\",\n self.name,\n len(tool_result_messages),\n )\n\n # Yield tool result responses\n for tool_result_msg in tool_result_messages:\n tool_response = AgentResponse(\n message=tool_result_msg,\n status=AgentRunningStatus.RUNNING.value,\n message_type=AgentMessageType.FINAL.value,\n )\n yield tool_response\n\n except Exception as e:\n err_text = str(e) or repr(e)\n logger.error(\"@%s Execution error: %s\", self.name, err_text, exc_info=True)\n error_response = AgentResponse(\n message=None,\n status=AgentRunningStatus.ERROR.value,\n error_msg=err_text,\n message_type=AgentMessageType.FINAL.value,\n )\n yield error_response\n" }, { "path": "cortex/agents/types.py", "content": "from enum import Enum\n\nfrom cortex.model.definition import ChatMessage, ModelParams\nfrom cortex.model.utils import merge_delta_message\nfrom pydantic import BaseModel, Field\n\nfrom cortex.tools.base import ToolSchema\n\n\nclass RunnerType(str, Enum):\n \"\"\"Agent runner type enum.\"\"\"\n\n LOCAL = \"local\"\n REMOTE = \"remote\"\n\n\nclass AgentRunningStatus(str, Enum):\n \"\"\"Agent running status enum.\"\"\"\n\n FINISHED = \"finished\"\n STOPPED = \"stopped\"\n ERROR = \"error\"\n RUNNING = \"running\"\n\n\nclass AgentMessageType(str, Enum):\n \"\"\"Agent message type enum.\"\"\"\n\n STREAM = \"stream\" # Streaming output\n ACCUMULATED = \"accumulated\" # Accumulated output\n FINAL = \"final\" # Final output\n\n\nclass AgentResponseType(str, Enum):\n \"\"\"Agent response type enum.\"\"\"\n\n RESPONSE = \"response\"\n TOOL_CALL = \"tool_call\"\n TOOL_RESULT = \"tool_result\"\n\n\nclass AgentConfig(BaseModel):\n \"\"\"Declarative configuration for Agent.\"\"\"\n\n model: ModelParams\n name: str = Field(default=\"\")\n agent_type: str | None = None\n system_prompt: str | None = None\n description: str | None = None\n tools: list[ToolSchema | str] = Field(default_factory=list)\n max_steps: int = 10\n extra_config: dict | None = None\n runner_type: RunnerType = RunnerType.LOCAL\n endpoint: str | None = None\n unfinished_mode: bool = False\n use_share_context: bool = False\n\n\nclass AgentResponse(BaseModel):\n \"\"\"Agent response model.\"\"\"\n\n agent_name: str | None = None\n message: ChatMessage | None = None\n message_type: AgentMessageType = (\n AgentMessageType.FINAL\n ) # delta: streaming output, accumulated: accumulated output, final: final output\n status: AgentRunningStatus = AgentRunningStatus.RUNNING\n error_msg: str | None = None\n metadata: dict[str, object] | None = None\n\n def get_type(self) -> AgentResponseType:\n \"\"\"Get response type.\"\"\"\n if self.message is None:\n return AgentResponseType.RESPONSE\n if self.message.tool_call_id is not None:\n return AgentResponseType.TOOL_CALL\n if self.message.tool_calls is not None:\n return AgentResponseType.TOOL_CALL\n return AgentResponseType.RESPONSE\n\n def __add__(self, other: \"AgentResponse\") -> \"AgentResponse\":\n if not isinstance(other, AgentResponse):\n return NotImplemented\n\n # Merge delta_message dictionaries\n merged_delta_dict = merge_delta_message(\n self.message.model_dump() if self.message else None,\n other.message.model_dump() if other.message else None,\n )\n merged_delta = ChatMessage(**merged_delta_dict)\n\n # Create new field dictionary, default to self's fields\n new_fields = self.model_dump()\n new_fields[\"message\"] = merged_delta\n\n # Iterate other fields, override with other's value if present\n for field, _ in self.model_dump().items():\n if field == \"message\":\n continue # Already processed\n other_value = getattr(other, field)\n if other_value:\n new_fields[field] = other_value\n\n return AgentResponse(**new_fields)\n" }, { "path": "cortex/context/__init__.py", "content": "\"\"\"Context module for managing conversation and session context.\"\"\"\n\nfrom cortex.context.base_context import BaseContext\nfrom cortex.context.file_context import FileContext\nfrom cortex.context.simple_context import SimpleContext\n\n__all__ = [\"BaseContext\", \"make_simple_context\", \"make_file_context\"]\n\n\ndef make_simple_context(session_id: str) -> BaseContext:\n return SimpleContext(session_id)\n\n\ndef make_file_context(path: str, session_id: str) -> BaseContext:\n return FileContext(session_id, path)\n" }, { "path": "cortex/context/base_context.py", "content": "\"\"\"Base context management class.\"\"\"\n\nfrom abc import ABC, abstractmethod\nfrom typing import List\n\nfrom cortex.model.definition import ChatMessage\n\n\nclass BaseContext(ABC):\n \"\"\"Base context management class, providing basic interface for session message management.\"\"\"\n\n def __init__(self, session_id: str):\n \"\"\"\n Initialize base context.\n\n Args:\n session_id: Session ID\n \"\"\"\n self.session_id = session_id\n\n @abstractmethod\n def add(self, messages: list[ChatMessage]) -> None:\n \"\"\"\n Add chat messages to context.\n\n Args:\n messages: List of chat messages to add\n \"\"\"\n ...\n\n @abstractmethod\n def get_all(self) -> List[ChatMessage]:\n \"\"\"\n Get all chat messages.\n\n Returns:\n List[ChatMessage]: List of all chat messages\n \"\"\"\n ...\n" }, { "path": "cortex/context/file_context.py", "content": "\"\"\"File-based context management class.\"\"\"\n\nimport asyncio\nimport json\nimport os\nfrom pathlib import Path\nfrom typing import List\n\nfrom cortex.model.definition import ChatMessage\n\nfrom cortex.context.base_context import BaseContext\n\n\nclass FileContext(BaseContext):\n \"\"\"File-based context management class, each session_id corresponds to a file.\"\"\"\n\n def __init__(\n self,\n session_id: str,\n storage_dir: str = \"contexts\",\n batch_size: int = 5,\n delay_seconds: float = 2.0,\n ):\n \"\"\"\n Initialize file context.\n\n Args:\n session_id: Session ID\n storage_dir: Storage directory, defaults to \"contexts\"\n batch_size: Batch write size, write immediately when this count is reached\n delay_seconds: Delay write time (seconds)\n \"\"\"\n super().__init__(session_id)\n self.storage_dir = Path(storage_dir)\n self.storage_dir.mkdir(exist_ok=True)\n self.file_path = self.storage_dir / f\"{session_id}.jsonl\"\n self._messages: List[ChatMessage] = []\n self._pending_messages: List[ChatMessage] = []\n self._batch_size = batch_size\n self._delay_seconds = delay_seconds\n self._write_task = None\n self._load_messages()\n\n def __del__(self):\n \"\"\"Automatically call flush when object is garbage collected\"\"\"\n try:\n self.flush()\n except Exception:\n # Ignore exceptions in destructor to avoid affecting garbage collection\n pass\n\n def _load_messages(self) -> None:\n \"\"\"Load messages from file\"\"\"\n if self.file_path.exists():\n try:\n with open(self.file_path, \"r\", encoding=\"utf-8\") as f:\n self._messages = []\n for line in f:\n line = line.strip()\n if line:\n msg_data = json.loads(line)\n self._messages.append(ChatMessage(**msg_data))\n except (json.JSONDecodeError, KeyError, TypeError):\n # If file is corrupted or format is incorrect, start fresh\n self._messages = []\n\n def _save_messages(self) -> None:\n \"\"\"Save messages to file\"\"\"\n # Rewrite entire file\n all_messages = self._messages + self._pending_messages\n with open(self.file_path, \"w\", encoding=\"utf-8\") as f:\n for msg in all_messages:\n json_line = json.dumps(msg.model_dump(), ensure_ascii=False)\n f.write(json_line + \"\\n\")\n\n # Move pending messages to main message list\n self._messages.extend(self._pending_messages)\n self._pending_messages = []\n\n async def _delayed_write(self) -> None:\n \"\"\"Delayed write task\"\"\"\n await asyncio.sleep(self._delay_seconds)\n if self._pending_messages:\n self._save_messages()\n self._write_task = None\n\n def _schedule_write(self) -> None:\n \"\"\"Schedule write task\"\"\"\n # If batch size is reached, write immediately\n if len(self._pending_messages) >= self._batch_size:\n if self._write_task:\n self._write_task.cancel()\n self._write_task = None\n self._save_messages()\n else:\n # If no write task in progress, create a delayed write task\n if not self._write_task:\n self._write_task = asyncio.create_task(self._delayed_write())\n\n def add(self, messages: list[ChatMessage]) -> None:\n \"\"\"Add chat messages to context\n\n Args:\n messages: List of chat messages to add\n \"\"\"\n self._pending_messages.extend(messages)\n self._schedule_write()\n\n def get_all(self) -> List[ChatMessage]:\n \"\"\"Get all chat messages\n\n Returns:\n List[ChatMessage]: List of all chat messages\n \"\"\"\n return (self._messages + self._pending_messages).copy()\n\n def clear(self) -> None:\n \"\"\"Clear context messages\"\"\"\n if self._write_task:\n self._write_task.cancel()\n self._write_task = None\n self._messages = []\n self._pending_messages = []\n if self.file_path.exists():\n os.remove(self.file_path)\n\n def flush(self) -> None:\n \"\"\"Force write all pending messages\"\"\"\n if self._write_task:\n self._write_task.cancel()\n self._write_task = None\n if self._pending_messages:\n self._save_messages()\n" }, { "path": "cortex/context/simple_context.py", "content": "from typing import List\n\nfrom cortex.model.definition import ChatMessage\n\nfrom cortex.context import BaseContext\n\nsimple_contexts: dict[str, list[ChatMessage]] = {}\n\n\nclass SimpleContext(BaseContext):\n \"\"\"Simple context management class for managing session messages.\"\"\"\n\n def __init__(self, session_id: str):\n super().__init__(session_id)\n\n def add(self, msg: list[ChatMessage]) -> None:\n \"\"\"Add chat message list to context.\"\"\"\n if self.session_id in simple_contexts:\n simple_contexts[self.session_id].extend(msg)\n else:\n simple_contexts[self.session_id] = msg\n\n def get_all(self) -> List[ChatMessage]:\n \"\"\"Get all chat messages.\"\"\"\n return simple_contexts.get(self.session_id, [])\n" }, { "path": "cortex/env.py", "content": "from __future__ import annotations\n\nimport os\nfrom pathlib import Path\n\n\ndef _repo_root() -> Path:\n return Path(__file__).resolve().parent.parent\n\n\ndef _parse_env_line(line: str) -> tuple[str, str] | None:\n stripped = line.strip()\n if not stripped or stripped.startswith(\"#\"):\n return None\n if stripped.startswith(\"export \"):\n stripped = stripped[7:].lstrip()\n if \"=\" not in stripped:\n return None\n key, value = stripped.split(\"=\", 1)\n key = key.strip()\n if not key:\n return None\n value = value.strip()\n if len(value) >= 2 and value[0] == value[-1] and value[0] in {\"'\", '\"'}:\n value = value[1:-1]\n return key, value\n\n\ndef _load_env_fallback(path: Path) -> None:\n try:\n content = path.read_text(encoding=\"utf-8\")\n except Exception:\n return\n for line in content.splitlines():\n parsed = _parse_env_line(line)\n if not parsed:\n continue\n key, value = parsed\n os.environ.setdefault(key, value)\n\n\ndef load_env() -> None:\n env_path = _repo_root() / \".env\"\n if not env_path.exists():\n return\n try:\n from dotenv import load_dotenv # type: ignore\n except Exception:\n _load_env_fallback(env_path)\n return\n load_dotenv(env_path, override=False)\n" }, { "path": "cortex/examples/agents/ask_input_agent.py", "content": "\"\"\"AskInputAgent - Fixed-flow Agent for asking user input and repeating what the user says\"\"\"\n\nimport json\nimport logging\nimport uuid\nfrom typing import AsyncGenerator\nfrom uuid import uuid4\n\nfrom cortex.model.definition import ChatMessage, ChatToolCall, ContentBlockType, Function\n\nfrom cortex.agents.base_agent import BaseAgent\nfrom cortex.agents.input.input import InputChannel\nfrom cortex.agents.types import (\n AgentConfig,\n AgentMessageType,\n AgentResponse,\n AgentRunningStatus,\n)\nfrom cortex.context import BaseContext, make_simple_context\nfrom cortex.model import ModelParams\nfrom cortex.tools.client_tool import ClientTool\nfrom cortex.tools.toolset import ToolSet\nfrom cortex.tools.types import ToolType\n\nlogger = logging.getLogger(__name__)\n\n\nasync def init_ask_input_tools() -> ToolSet:\n \"\"\"Initialize ask_input toolset\"\"\"\n toolset = ToolSet()\n\n # Register ask_input tool (ClientTool)\n ask_input_tool = ClientTool(\n name=\"ask_input\",\n description=\"Ask user for input. Used for scenarios requiring user interaction such as obtaining user feedback, confirmation, modification suggestions, etc. Parameters: prompt (required) - prompt message to display to the user; context (optional) - context information to help users understand the current situation.\",\n tool_type=ToolType.ASK_INPUT,\n channel=toolset.channel,\n timeout=300.0, # User input may take a long time\n client_params={\n \"properties\": {\n \"prompt\": {\n \"type\": \"string\",\n \"description\": \"Prompt message to display to the user, explaining what the user needs to do (confirm, modify, provide information, etc.)\",\n },\n \"context\": {\n \"type\": \"string\",\n \"description\": \"Context information to help users understand the current situation, such as current plan content, items that need confirmation, etc.\",\n },\n },\n \"required\": [\"prompt\"],\n },\n )\n toolset.register(ask_input_tool)\n logger.info(\"Registered ask_input tool\")\n\n return toolset\n\n\nclass AskInputAgent(BaseAgent):\n \"\"\"AskInputAgent - Fixed-flow Agent for asking user input and repeating what the user says\n\n Fixed flow:\n 1. Send a function call with tool_name as ask_input\n 2. After receiving ask_input result, send a message repeating what the user said\n \"\"\"\n\n def __init__(\n self, context: BaseContext, config: AgentConfig, toolset: ToolSet | None = None\n ):\n super().__init__(config=config, toolset=toolset)\n self.context = context\n\n async def _run(\n self,\n messages: list[ChatMessage] | InputChannel[ChatMessage],\n additional_kwargs: dict | None = None,\n ) -> AsyncGenerator[AgentResponse, None]:\n \"\"\"\n Run agent, execute fixed flow\n\n Args:\n messages: Input message list or input channel\n additional_kwargs: Additional parameters\n\n Yields:\n AgentResponse: Agent response object\n \"\"\"\n try:\n # Step 1: Create a ChatMessage containing tool_calls, call ask_input tool\n tool_call_id = f\"call_{uuid.uuid4().hex[:8]}\"\n\n # Create tool call arguments\n tool_args = json.dumps({\"prompt\": \"Please enter some content\"})\n\n # Create ChatToolCall\n tool_call = ChatToolCall(\n id=tool_call_id,\n function=Function(\n name=\"ask_input\",\n arguments=tool_args,\n ),\n )\n\n # Create ChatMessage containing tool_calls\n tool_call_message = ChatMessage(\n role=\"assistant\",\n tool_calls=[tool_call],\n )\n\n # Yield tool call response\n tool_call_response = AgentResponse(\n message=tool_call_message,\n status=AgentRunningStatus.RUNNING.value,\n message_type=AgentMessageType.FINAL.value,\n )\n yield tool_call_response\n\n # Execute tool call\n tool_result_messages = await self.run_tool_call(tool_call_message)\n\n if not tool_result_messages:\n error_response = AgentResponse(\n message=None,\n status=AgentRunningStatus.ERROR.value,\n error_msg=\"ask_input tool call failed, no result returned\",\n message_type=AgentMessageType.FINAL.value,\n )\n yield error_response\n return\n\n # Yield tool result response\n tool_result_message = tool_result_messages[0]\n tool_result_response = AgentResponse(\n message=tool_result_message,\n status=AgentRunningStatus.RUNNING.value,\n message_type=AgentMessageType.FINAL.value,\n )\n yield tool_result_response\n\n # Step 2: Get user input from tool result and repeat what the user said\n user_input = None\n if tool_result_message.content:\n # Extract user input\n if isinstance(tool_result_message.content, list):\n for block in tool_result_message.content:\n if isinstance(block, dict):\n text_content = block.get(\n ContentBlockType.TEXT.value\n ) or block.get(\"text\")\n if text_content:\n user_input = str(text_content)\n break\n elif isinstance(block, str):\n user_input = block\n break\n elif isinstance(tool_result_message.content, str):\n user_input = tool_result_message.content\n\n if user_input is None:\n user_input = \"Failed to get user input\"\n\n # Create message repeating what the user said\n repeat_message = ChatMessage(\n role=\"assistant\",\n content=[\n {\n \"type\": ContentBlockType.TEXT.value,\n ContentBlockType.TEXT.value: f\"You said: {user_input}\",\n }\n ],\n )\n\n # Yield final response\n final_response = AgentResponse(\n message=repeat_message,\n status=AgentRunningStatus.FINISHED.value,\n message_type=AgentMessageType.FINAL.value,\n )\n yield final_response\n\n except Exception as e:\n logger.error(f\"@{self.name} execution error: {e}\", exc_info=True)\n error_response = AgentResponse(\n message=None,\n status=AgentRunningStatus.ERROR.value,\n error_msg=str(e),\n message_type=AgentMessageType.FINAL.value,\n )\n yield error_response\n\n\nasync def make_ask_input_agent(\n config: AgentConfig, context_id: str | None = None\n) -> BaseAgent:\n \"\"\"Create AskInputAgent\"\"\"\n toolset = await init_ask_input_tools()\n context = make_simple_context(context_id)\n if context_id is None:\n context_id = uuid4().hex\n return AskInputAgent(context=context, config=config, toolset=toolset)\n\n\ndef get_ask_input_agent_config() -> AgentConfig:\n \"\"\"Get AskInputAgent configuration\"\"\"\n return AgentConfig(\n name=\"AskInputAgent\",\n description=\"Fixed-flow Agent for asking user input and repeating what the user says. Does not require model calls.\",\n system_prompt=None, # No system prompt needed\n model=ModelParams(\n name=\"gpt-4o-mini\", # Although model calls are not needed, configuration requires it\n infer_kwargs={\"max_tokens\": 100, \"temperature\": 0.7, \"stream\": False},\n ),\n )\n" }, { "path": "cortex/examples/agents/deep_reasearch_agent.py", "content": "\"\"\"DeepResearchAgent - Deep research Agent\"\"\"\n\nfrom cortex.agents.base_agent import BaseAgent\nfrom cortex.agents.react_agent import ReActAgent\nfrom cortex.agents.types import AgentConfig\nfrom cortex.model import ModelParams\nfrom cortex.tools.agent_tool import AgentTool\nfrom cortex.tools.toolset import ToolSet\n\n\nasync def init_deep_research_tools() -> ToolSet:\n \"\"\"Initialize deep research tools\"\"\"\n # Register web_search tool\n toolset = ToolSet()\n await toolset.register_from_mcp_server(\n mcp_server=\"http://xxx/mcp\",\n tool_names=[\"web_search\"],\n )\n\n # plan agent tool\n plan_agent_tool = AgentTool(\n name=\"PlanAgent\",\n description=\"PlanAgent - Planning Agent responsible for creating plans\",\n channel=toolset.channel,\n timeout=3000.0,\n )\n toolset.register(plan_agent_tool)\n\n return toolset\n\n\ndef get_deep_research_agent_config() -> AgentConfig:\n \"\"\"Get deep research Agent configuration\"\"\"\n return AgentConfig(\n name=\"DeepResearchAgent\",\n description=\"DeepResearchAgent - Deep research Agent responsible for deep research tasks\",\n system_prompt=\"\"\"You are a powerful \"Deep Search Agent\", an intelligent system with capabilities for initial thinking, reflection, and adaptation. When facing a complex problem, you don't simply execute searches, but rather think deeply, plan, and flexibly utilize tools (such as web search, webpage access, code execution) to find answers like a smart assistant.\nThe key lies in \"depth\" and \"intelligence\", which means the Agent needs to have the following characteristics:\n1. **Tool Coordination**: Able to carefully analyze problems, reference plans already made in history, and strategically use tools to collect, process, and present information.\n2. **Data Extraction and Visualization**: Able to extract data relevant to the problem from massive amounts of information, and use the visualization (visualize_data) tool to transform data into intuitive charts.\n3. **Reflection and Adaptation**: This is the most important capability! During the search process, if problems are encountered (such as unsatisfactory search results, insufficient information, or uncertainty), the Agent won't give up easily, but will proactively reflect and adjust strategies. For example:\n - Change search keywords.\n - View more search results.\n - Determine whether current information is sufficient to answer the question; if not, continue searching for missing information.\n - Evaluate the reliability of information sources.\n - Use different tools or information sources for cross-validation.\n - Determine whether there is critical data that needs to be displayed through charts; if so, use visualization tools to present it.\n - Check whether the collected information fully meets all requirements of the original question.\n\n\nYou must always follow these rules to complete tasks:\n1. Always provide tool calls, otherwise it will fail\n2. Always use correct tool parameters. Don't use variable names in action parameters, use specific values instead\n3. Only call tools when needed: if information is not needed, don't call search tools, try to solve the problem yourself\n4. Never repeat a tool call with exactly the same parameters that has already been used\n5. Only use the visualize_data tool for visualization, while the execute_python_code tool can only be used for complex data calculations or file processing\n6. If information involves key numerical values, data presentation, data comparison, process chains, multiple stages, entity relationships, timelines, etc., you must use the visualize_data tool to generate charts\n7. For data that has already been visualized, don't call the visualization tool repeatedly\n8. Do not use the execute_python_code tool to output large amounts of text, do not use the execute_python_code tool to output reports\n9. Never express gratitude for any tool call results (such as search results)\n10. Multi-language support: You support responding in Chinese, English, Japanese, Korean, Traditional Chinese, Spanish, and Portuguese, automatically identifying the \"user's input language\" and matching the output.\n11. Search additional requirements: When the problem you need to search is related to travel/public opinion, you need to generate queries in the language corresponding to the travel/public opinion location for searching, and also generate an identical query in the user's language for searching.\n\n\nStart now! If you complete the task correctly, you will receive a $1,000,000 reward.\n \"\"\",\n model=ModelParams(\n name=\"gpt-5.1\",\n infer_kwargs={\"max_tokens\": 2000, \"temperature\": 0.7, \"stream\": False},\n ),\n max_steps=10,\n )\n\n\nasync def make_deep_research_agent(session_id: str, config: AgentConfig) -> BaseAgent:\n \"\"\"Create deep research Agent\"\"\"\n toolset = await init_deep_research_tools()\n return ReActAgent(config=config, toolset=toolset)\n" }, { "path": "cortex/examples/agents/main_agent.py", "content": "\"\"\"MainAgent - Main coordination Agent.\"\"\"\n\nimport logging\nfrom uuid import uuid4\n\nfrom cortex.agents.base_agent import BaseAgent\nfrom cortex.agents.react_agent import ReActAgent\nfrom cortex.agents.types import AgentConfig\nfrom cortex.context import make_simple_context\nfrom cortex.examples.agents.math_agent import get_math_agent_config, make_math_agent\nfrom cortex.examples.agents.search_agent import (\n get_search_agent_config,\n make_search_agent,\n)\nfrom cortex.model import ModelParams\nfrom cortex.tools.agent_tool import AgentTool\nfrom cortex.tools.toolset import ToolSet\n\nlogger = logging.getLogger(__name__)\n\n\nasync def init_main_tools() -> ToolSet:\n \"\"\"Initialize main coordination tools.\"\"\"\n toolset = ToolSet()\n search_agent = await make_search_agent(config=get_search_agent_config())\n math_agent = await make_math_agent(config=get_math_agent_config())\n\n search_tool_params = search_agent.as_tool()\n math_tool_params = math_agent.as_tool()\n toolset.register(\n AgentTool(\n name=search_tool_params[\"name\"],\n description=search_tool_params[\"description\"],\n timeout=search_tool_params[\"timeout\"]\n if \"timeout\" in search_tool_params\n else 300,\n channel=toolset.channel,\n )\n )\n toolset.register(\n AgentTool(\n name=math_tool_params[\"name\"],\n description=math_tool_params[\"description\"],\n timeout=math_tool_params[\"timeout\"]\n if \"timeout\" in math_tool_params\n else 300,\n channel=toolset.channel,\n )\n )\n return toolset\n\n\nasync def make_main_agent(\n config: AgentConfig, context_id: str | None = None\n) -> BaseAgent:\n \"\"\"Create MainAgent.\"\"\"\n toolset = await init_main_tools()\n if context_id is None:\n context_id = uuid4().hex\n context = make_simple_context(context_id)\n return ReActAgent(context=context, config=config, toolset=toolset)\n\n\ndef get_main_agent_config() -> AgentConfig:\n \"\"\"Get MainAgent configuration.\"\"\"\n return AgentConfig(\n name=\"MainAgent\",\n description=\"Main coordination Agent responsible for coordinating and calling other specialized Agents to complete tasks. Can select appropriate Agents based on task requirements (e.g., MathAgent for mathematical calculations, SearchAgent for information search) and coordinate multiple Agents to complete complex tasks.\",\n system_prompt=\"You are a main coordination Agent responsible for coordinating and calling other specialized Agents to complete tasks.\",\n model=ModelParams(\n name=\"gpt-4o-mini\",\n infer_kwargs={\"max_tokens\": 2000, \"temperature\": 0.7, \"stream\": False},\n ),\n )\n" }, { "path": "cortex/examples/agents/math_agent.py", "content": "\"\"\"MathAgent - An Agent that can solve mathematical problems.\"\"\"\n\nimport logging\nimport math\nfrom uuid import uuid4\n\nfrom cortex.agents.base_agent import BaseAgent\nfrom cortex.agents.react_agent import ReActAgent\nfrom cortex.agents.types import AgentConfig\nfrom cortex.context import make_simple_context\nfrom cortex.model import ModelParams\nfrom cortex.tools.function_tool import FunctionTool\nfrom cortex.tools.toolset import ToolSet\n\nlogger = logging.getLogger(__name__)\n\n\n# Basic math operation tool functions\ndef add(a: float, b: float) -> float:\n \"\"\"Add two numbers.\n\n Args:\n a: First number\n b: Second number\n\n Returns:\n Sum of the two numbers\n \"\"\"\n return a + b\n\n\ndef subtract(a: float, b: float) -> float:\n \"\"\"Subtract two numbers.\n\n Args:\n a: Minuend\n b: Subtrahend\n\n Returns:\n Difference of the two numbers\n \"\"\"\n return a - b\n\n\ndef multiply(a: float, b: float) -> float:\n \"\"\"Multiply two numbers.\n\n Args:\n a: First number\n b: Second number\n\n Returns:\n Product of the two numbers\n \"\"\"\n return a * b\n\n\ndef divide(a: float, b: float) -> float:\n \"\"\"Divide two numbers.\n\n Args:\n a: Dividend\n b: Divisor (cannot be 0)\n\n Returns:\n Quotient of the two numbers\n\n Raises:\n ValueError: When divisor is 0\n \"\"\"\n if b == 0:\n raise ValueError(\"Divisor cannot be 0\")\n return a / b\n\n\ndef power(base: float, exponent: float) -> float:\n \"\"\"Calculate power operation.\n\n Args:\n base: Base number\n exponent: Exponent\n\n Returns:\n Base raised to the power of exponent\n \"\"\"\n return base**exponent\n\n\ndef sqrt(number: float) -> float:\n \"\"\"Calculate square root.\n\n Args:\n number: Number to calculate square root (must be >= 0)\n\n Returns:\n Square root of the number\n\n Raises:\n ValueError: When number is less than 0\n \"\"\"\n if number < 0:\n raise ValueError(\"Cannot calculate square root of negative number\")\n return math.sqrt(number)\n\n\ndef calculate_expression(expression: str) -> float:\n \"\"\"Calculate a mathematical expression (using eval, for simple expressions only).\n\n Args:\n expression: Mathematical expression string, e.g., \"2 + 3 * 4\"\n\n Returns:\n Calculation result\n\n Warning:\n This function uses eval, only for simple mathematical expressions, do not use for untrusted input\n \"\"\"\n # Only allow numbers, operators, and parentheses\n allowed_chars = set(\"0123456789+-*/.() \")\n if not all(c in allowed_chars for c in expression):\n raise ValueError(\"Expression contains disallowed characters\")\n try:\n return float(eval(expression))\n except Exception as e:\n raise ValueError(f\"Expression calculation error: {str(e)}\")\n\n\nasync def init_math_tools() -> ToolSet:\n \"\"\"Initialize math toolset.\"\"\"\n toolset = ToolSet()\n\n # Register basic math operation tools\n math_functions = [\n (\"add\", add, \"Addition: add two numbers\"),\n (\"subtract\", subtract, \"Subtraction: subtract two numbers\"),\n (\"multiply\", multiply, \"Multiplication: multiply two numbers\"),\n (\"divide\", divide, \"Division: divide two numbers\"),\n (\"power\", power, \"Power: calculate a number raised to a power\"),\n (\"sqrt\", sqrt, \"Square root: calculate the square root of a number\"),\n (\n \"calculate_expression\",\n calculate_expression,\n \"Calculate expression: evaluate a simple mathematical expression string\",\n ),\n ]\n\n for tool_name, tool_func, tool_desc in math_functions:\n function_tool = FunctionTool(\n name=tool_name,\n func=tool_func,\n description=tool_desc,\n )\n toolset.register(function_tool)\n logger.info(f\"Registered math tool: {tool_name}\")\n\n return toolset\n\n\nasync def make_math_agent(\n config: AgentConfig, context_id: str | None = None\n) -> BaseAgent:\n \"\"\"Create MathAgent.\"\"\"\n toolset = await init_math_tools()\n if context_id is None:\n context_id = uuid4().hex\n context = make_simple_context(context_id)\n return ReActAgent(context=context, config=config, toolset=toolset)\n\n\ndef get_math_agent_config() -> AgentConfig:\n \"\"\"Get MathAgent configuration.\"\"\"\n return AgentConfig(\n name=\"MathAgent\",\n description=\"An Agent specialized for mathematical calculations. Supports basic math operations (addition, subtraction, multiplication, division, power, square root, etc.), can handle complex mathematical expressions, supports multi-step calculations, and provides detailed calculation process explanations. Suitable for arithmetic, algebra, geometry calculations, and mathematical expression solving.\",\n system_prompt=\"You are a professional mathematical calculation assistant.\",\n model=ModelParams(\n name=\"gpt-4o-mini\",\n infer_kwargs={\"max_tokens\": 2000, \"temperature\": 0.7, \"stream\": False},\n ),\n )\n" }, { "path": "cortex/examples/agents/plan_agent.py", "content": "\"\"\"PlanAgent - Creates plans based on task objectives, and finalizes plans after user confirmation or modification\"\"\"\n\nimport logging\n\nfrom cortex.agents.base_agent import BaseAgent\nfrom cortex.agents.react_agent import ReActAgent\nfrom cortex.agents.types import AgentConfig\nfrom cortex.context import make_simple_context\nfrom cortex.model import ModelParams\nfrom cortex.tools.client_tool import ClientTool\nfrom cortex.tools.toolset import ToolSet\nfrom cortex.tools.types import ToolType\n\nlogger = logging.getLogger(__name__)\n\n\nasync def init_plan_tools() -> ToolSet:\n \"\"\"Initialize plan toolset\"\"\"\n toolset = ToolSet()\n\n # Register web_search tool\n await toolset.register_from_mcp_server(\n mcp_server=\"http://xxx/mcp\",\n tool_names=[\"web_search\"],\n )\n\n # Register ask_input tool (ClientTool)\n ask_input_tool = ClientTool(\n name=\"ask_input\",\n description=\"Ask users for input, confirmation, or modification suggestions. Used for scenarios requiring user interaction such as obtaining user feedback, confirming plans, modifying suggestions, etc. Parameters: prompt (required) - prompt message to display to the user; context (optional) - context information to help users understand the current situation.\",\n tool_type=ToolType.ASK_INPUT,\n channel=toolset.channel,\n timeout=300.0, # User input may take a long time\n client_params={\n \"properties\": {\n \"prompt\": {\n \"type\": \"string\",\n \"description\": \"Prompt message to display to the user, explaining what the user needs to do (confirm, modify, provide information, etc.)\",\n },\n \"context\": {\n \"type\": \"string\",\n \"description\": \"Context information to help users understand the current situation, such as current plan content, items that need confirmation, etc.\",\n },\n },\n \"required\": [\"prompt\"],\n },\n )\n toolset.register(ask_input_tool)\n logger.info(\"Registered ask_input tool\")\n\n return toolset\n\n\nasync def make_plan_agent(context_id: str, config: AgentConfig) -> BaseAgent:\n \"\"\"Create PlanAgent\"\"\"\n toolset = await init_plan_tools()\n context = make_simple_context(context_id)\n return ReActAgent(context=context, config=config, toolset=toolset)\n\n\ndef get_plan_agent_config() -> AgentConfig:\n \"\"\"Get PlanAgent configuration\"\"\"\n return AgentConfig(\n name=\"PlanAgent\",\n description=\"An Agent specialized in creating plans based on task objectives. Capable of analyzing task requirements, searching for relevant information, creating detailed plans, and finalizing plans after user confirmation or modification. Suitable for project planning, task decomposition, action plan creation, and similar scenarios.\",\n system_prompt=\"\"\"You are a professional planning assistant. Your responsibility is to create detailed, feasible plans based on the task objectives provided by the user.\n\nWorkflow:\n1. **Understand Task Objectives**: Carefully analyze the task objectives provided by the user, understand the core requirements, expected results, and constraints of the task.\n\n2. **Information Collection** (if needed):\n - If the task involves the need for latest information or professional knowledge, use the web_search tool to search for relevant information\n - Collect background knowledge, best practices, case references, etc. related to the task\n\n3. **User Confirmation and Modification**:\n - Break down the task into clear steps\n - Use the ask_input tool to show the preliminary plan to the user\n - Clearly explain the plan content in the prompt and ask the user if modifications are needed\n - Provide detailed plan content in context for the user to review\n - Adjust the plan based on user feedback (confirmation, modification suggestions, etc.)\n\n4. **Create Final Plan**:\n - Create the final plan based on user confirmation or modification feedback\n - Ensure the plan is complete, clear, and executable\n - Provide a summary of the plan and execution suggestions\n\nImportant Principles:\n- Be sure to use the ask_input tool to show the preliminary plan to the user and adjust the plan based on user feedback\n- Be sure to output the final plan; the final plan does not need modification, just output it directly\n- Plans should be specific and executable, avoiding being too abstract\n- Consider practical feasibility and resource constraints\"\"\",\n model=ModelParams(\n name=\"gpt-5.1\",\n infer_kwargs={\"max_tokens\": 2000, \"temperature\": 0.7, \"stream\": False},\n ),\n )\n" }, { "path": "cortex/examples/agents/search_agent.py", "content": "\"\"\"SearchAgent - An Agent that can integrate search information\"\"\"\n\nimport logging\nfrom uuid import uuid4\n\nfrom cortex.agents.base_agent import BaseAgent\nfrom cortex.agents.react_agent import ReActAgent\nfrom cortex.agents.types import AgentConfig\nfrom cortex.context import make_simple_context\nfrom cortex.model import ModelParams\nfrom cortex.tools.toolset import ToolSet\n\nlogger = logging.getLogger(__name__)\n\n\nasync def init_search_tools() -> ToolSet:\n \"\"\"Initialize search tools\"\"\"\n toolset = ToolSet()\n await toolset.register_from_mcp_server(\n mcp_server=\"http://xxx/mcp\",\n tool_names=[\"web_search\"],\n )\n return toolset\n\n\nasync def make_search_agent(\n config: AgentConfig, context_id: str | None = None\n) -> BaseAgent:\n \"\"\"Create SearchAgent\"\"\"\n toolset = await init_search_tools()\n if context_id is None:\n context_id = uuid4().hex\n context = make_simple_context(context_id)\n return ReActAgent(context=context, config=config, toolset=toolset)\n\n\ndef get_search_agent_config() -> AgentConfig:\n \"\"\"Get SearchAgent configuration\"\"\"\n return AgentConfig(\n name=\"SearchAgent\",\n description=\"An Agent specialized in searching and integrating information. Uses search tools to find latest information, integrates multiple search results, provides comprehensive and accurate answers with citations. Suitable for finding latest information, integrating information from multiple sources, answering questions requiring real-time data, and information retrieval and organization.\",\n system_prompt=\"You are a professional information search and integration assistant.\",\n model=ModelParams(\n name=\"gpt-4o-mini\",\n infer_kwargs={\"max_tokens\": 2000, \"temperature\": 0.7, \"stream\": False},\n ),\n )\n" }, { "path": "cortex/examples/demo_agent_cli.py", "content": "\"\"\"Demo Agent CLI - TUI interface built using TUI module to display AgentEvent.\"\"\"\n\nimport asyncio\nimport logging\nfrom pathlib import Path\n\nfrom agentkit.trace import LocalStorageTracer\n\nfrom cortex.agents.agent_factory import AgentFactory\nfrom cortex.examples.agents.ask_input_agent import (\n get_ask_input_agent_config,\n make_ask_input_agent,\n)\nfrom cortex.examples.agents.deep_reasearch_agent import (\n get_deep_research_agent_config,\n make_deep_research_agent,\n)\nfrom cortex.examples.agents.main_agent import get_main_agent_config, make_main_agent\nfrom cortex.examples.agents.math_agent import get_math_agent_config, make_math_agent\nfrom cortex.examples.agents.plan_agent import get_plan_agent_config, make_plan_agent\nfrom cortex.examples.agents.search_agent import (\n get_search_agent_config,\n make_search_agent,\n)\nfrom cortex.orchestrator.orchestrator import Orchestrator\nfrom cortex.tui import AgentTUIApp\n\nlogger = logging.getLogger(__name__)\n\n\nasync def main():\n \"\"\"\n Main function.\n\n Example:\n python -m cortex.examples.demo_agent_cli\n \"\"\"\n # Create logs directory\n logs_dir = Path(\"./logs\")\n logs_dir.mkdir(exist_ok=True)\n\n # Configure logger to output to file\n log_file = logs_dir / \"demo_agent_cli.log\"\n file_handler = logging.FileHandler(log_file, encoding=\"utf-8\")\n file_handler.setLevel(logging.DEBUG)\n file_handler.setFormatter(\n logging.Formatter(\n \"%(asctime)s %(levelname)s %(filename)s:%(lineno)d %(message)s\"\n )\n )\n\n # Configure root logger to output only to file, not console\n root_logger = logging.getLogger()\n root_logger.setLevel(logging.DEBUG)\n root_logger.handlers.clear() # Clear default console handler\n root_logger.addHandler(file_handler)\n\n # Initialize Orchestrator\n agent_factory = AgentFactory()\n agent_factory.register_agent(\n name=\"DeepResearchAgent\",\n make_agent_func=make_deep_research_agent,\n default_config=get_deep_research_agent_config(),\n )\n agent_factory.register_agent(\n name=\"PlanAgent\",\n make_agent_func=make_plan_agent,\n default_config=get_plan_agent_config(),\n )\n agent_factory.register_agent(\n name=\"AskInputAgent\",\n make_agent_func=make_ask_input_agent,\n default_config=get_ask_input_agent_config(),\n )\n\n agent_factory.register_agent(\n name=\"MainAgent\",\n make_agent_func=make_main_agent,\n default_config=get_main_agent_config(),\n )\n agent_factory.register_agent(\n name=\"SearchAgent\",\n make_agent_func=make_search_agent,\n default_config=get_search_agent_config(),\n )\n agent_factory.register_agent(\n name=\"MathAgent\",\n make_agent_func=make_math_agent,\n default_config=get_math_agent_config(),\n )\n orchestrator = Orchestrator(agent_factory)\n\n tracer = LocalStorageTracer(storage_dir=\"./traces\")\n app = AgentTUIApp(orchestrator=orchestrator, workdir=\"./logs\", tracer=tracer)\n await app.run_async()\n\n\nif __name__ == \"__main__\":\n asyncio.run(main())\n" }, { "path": "cortex/examples/demo_agent_with_orchestrator.py", "content": "\"\"\"Demo Agent with Orchestrator - Demonstrates how to use the Orchestrator pattern to coordinate multiple Agents\"\"\"\n\nimport argparse\nimport asyncio\nimport logging\n\nfrom agentkit.trace import LocalStorageTracer, SpanContext\nfrom cortex.model.definition import ChatMessage\n\nfrom cortex.agents.agent_factory import AgentFactory\nfrom cortex.examples.agents.main_agent import get_main_agent_config, make_main_agent\nfrom cortex.examples.agents.math_agent import get_math_agent_config, make_math_agent\nfrom cortex.examples.agents.search_agent import (\n get_search_agent_config,\n make_search_agent,\n)\nfrom cortex.orchestrator import AgentEvent\nfrom cortex.orchestrator.orchestrator import Orchestrator\nfrom cortex.orchestrator.types import AgentEventType, AgentRequest\n\nlogger = logging.getLogger(__name__)\n\n\nasync def main(\n request: str | None = None,\n output_file: str | None = None,\n):\n \"\"\"\n Main function demonstrating how to use Orchestrator\n\n Args:\n request: User request content. If None, uses default test message\n output_file: Output file path. If None, outputs to stdout\n \"\"\"\n\n agent_factory = AgentFactory()\n agent_factory.register_agent(\n name=\"MainAgent\",\n make_agent_func=make_main_agent,\n default_config=get_main_agent_config(),\n )\n agent_factory.register_agent(\n name=\"SearchAgent\",\n make_agent_func=make_search_agent,\n default_config=get_search_agent_config(),\n )\n agent_factory.register_agent(\n name=\"MathAgent\",\n make_agent_func=make_math_agent,\n default_config=get_math_agent_config(),\n )\n orchestrator = Orchestrator(agent_factory)\n\n messages = [ChatMessage(role=\"user\", content=request)]\n async for event in orchestrator.run(\n agent_name=\"MainAgent\",\n event=AgentEvent(\n type=AgentEventType.REQUEST,\n request=AgentRequest(\n agent_name=\"MainAgent\",\n messages=messages,\n ),\n ),\n agent_config=None,\n ):\n logger.info(\"------ event: %s\", event.model_dump_json())\n if output_file:\n with open(output_file, \"a\", encoding=\"utf-8\") as f:\n f.write(event.model_dump_json(ensure_ascii=False) + \"\\n\")\n else:\n print(event.model_dump_json(ensure_ascii=False))\n\n\nif __name__ == \"__main__\":\n # Parse command line arguments\n parser = argparse.ArgumentParser(\n description=\"Demo Agent with Orchestrator - Demonstrates how to use the Orchestrator pattern to coordinate multiple Agents\",\n formatter_class=argparse.RawDescriptionHelpFormatter,\n epilog=\"\"\"\nExample usage:\n # Use default request, output to stdout\n python cortex/examples/demo_agent_with_orchestrator.py\n\n # Specify request content\n python cortex/examples/demo_agent_with_orchestrator.py --request \"Please calculate the result of 123 + 456 + 789\"\n\n # Output to file\n python cortex/examples/demo_agent_with_orchestrator.py --request \"Please calculate the result of 123 + 456 + 789\" --output result.txt\n\n # Use grep to filter responses (stdout with prefix)\n python cortex/examples/demo_agent_with_orchestrator.py | grep \"\\\\[RESPONSE\\\\]\"\n\nDemo examples (can be uncommented in code or specified via --request parameter):\n 1. Math calculation - Simple addition: \"Please calculate the result of 123 + 456 + 789\"\n 2. Math calculation - Complex operation: \"Please calculate the result of (123 + 456) * 789 / 100\"\n 3. Math calculation - Multi-step: \"Please calculate 15 + 27 + 39 + 41 step by step, first calculate the sum of the first two numbers, then add the remaining numbers\"\n 4. Information search: \"Please search for the latest information about China's central bank gold reserves in 2025\"\n 5. Mixed task: \"Please first search for China's GDP growth rate in 2024, then calculate what the GDP would be if it grows by 5% in 2025\"\n \"\"\",\n )\n parser.add_argument(\n \"--request\",\n \"-r\",\n type=str,\n default=None,\n help=\"User request content (if not specified, uses default test message)\",\n )\n parser.add_argument(\n \"--output\",\n \"-o\",\n type=str,\n default=None,\n help=\"Output file path (if not specified, outputs to stdout with prefix for grep)\",\n )\n parser.add_argument(\n \"--log-level\",\n type=str,\n default=\"DEBUG\",\n choices=[\"DEBUG\", \"INFO\", \"WARNING\", \"ERROR\"],\n help=\"Log level (default: INFO)\",\n )\n\n args = parser.parse_args()\n\n # Set log level and corresponding source file and line number\n log_level = getattr(logging, args.log_level.upper())\n logging.basicConfig(\n level=log_level,\n format=\"%(asctime)s %(levelname)s %(filename)s:%(lineno)d %(message)s\",\n encoding=\"utf-8\", # Add UTF-8 encoding support for proper character display\n )\n\n print(\"-----------------------------------------------------\")\n tracer = LocalStorageTracer(storage_dir=\"./traces\")\n ctx = SpanContext(tracer=tracer, app_name=\"demo_agent_with_orchestrator\")\n with ctx.span(name=\"demo_agent_with_orchestrator\") as span:\n trace_id = ctx.get_current_trace_id()\n logger.info(\"agent_cortex trace_id %s\", trace_id)\n asyncio.run(main(request=args.request, output_file=args.output))\n" }, { "path": "cortex/examples/demo_agent_with_tool.py", "content": "\"\"\"Demo Agent with Tool implementation.\"\"\"\n\nimport argparse\nimport asyncio\nimport logging\nimport math\nimport random\nimport string\nimport uuid\n\nfrom agentkit.trace import LocalStorageTracer, SpanContext\nfrom cortex.model.definition import ChatMessage\n\nfrom cortex.agents.react_agent import ReActAgent\nfrom cortex.agents.types import AgentConfig\nfrom cortex.context.simple_context import SimpleContext\nfrom cortex.model import ModelParams\nfrom cortex.tools.function_tool import FunctionTool\nfrom cortex.tools.toolset import ToolSet\n\nlogger = logging.getLogger(__name__)\n\n\ndef add_numbers(a: int, b: int) -> int:\n \"\"\"Add two numbers.\n\n Args:\n a: First number\n b: Second number\n\n Returns:\n Sum of the two numbers\n \"\"\"\n return a + b\n\n\nasync def multiply_numbers(a: int, b: int) -> int:\n \"\"\"Multiply two numbers.\n\n Args:\n a: First number\n b: Second number\n\n Returns:\n Product of the two numbers\n \"\"\"\n await asyncio.sleep(5)\n return a * b\n\n\ndef get_random_string(length: int = 10) -> str:\n \"\"\"Generate a random string.\n\n Args:\n length: String length, defaults to 10\n\n Returns:\n Randomly generated string\n \"\"\"\n return \"\".join(random.choices(string.ascii_letters + string.digits, k=length))\n\n\ndef calculate_area(radius: float) -> float:\n \"\"\"Calculate the area of a circle.\n\n Args:\n radius: Circle radius\n\n Returns:\n Area of the circle\n \"\"\"\n return math.pi * radius * radius\n\n\nasync def init_tools():\n \"\"\"Initialize toolset.\"\"\"\n toolset = ToolSet()\n available_functions = [\n (\"add_numbers\", add_numbers, \"Addition tool\"),\n (\"multiply_numbers\", multiply_numbers, \"Multiplication tool\"),\n (\"get_random_string\", get_random_string, \"Random string generator tool\"),\n (\"calculate_area\", calculate_area, \"Circle area calculator tool\"),\n ]\n\n # Register function tools to ToolSet\n for tool_name, tool_func, tool_desc in available_functions:\n function_tool = FunctionTool(\n name=tool_name,\n func=tool_func,\n description=tool_desc,\n )\n toolset.register(function_tool)\n logger.info(\"Registered tool: %s\", tool_name)\n\n await toolset.register_from_mcp_server(\n mcp_server=\"http://xxx/mcp\",\n tool_names=[\"web_search\"],\n )\n return toolset\n\n\nasync def main(user_input: str):\n \"\"\"Main function demonstrating how to use DemoAgentWithTool.\"\"\"\n # agent config\n agent_config = AgentConfig(\n name=\"demo_agent:ReActAgent\",\n description=\"Agent demonstrating how to use ToolSet and anymodel.\",\n system_prompt=\"You are a helpful assistant that can use the provided tools to help users.\",\n model=ModelParams(\n name=\"gpt-4o-mini\",\n infer_kwargs={\"max_tokens\": 2000, \"temperature\": 0.7, \"stream\": False},\n ),\n max_steps=10,\n )\n\n # toolset\n toolset = await init_tools()\n\n # agent instance\n context = SimpleContext(session_id=str(uuid.uuid4()))\n agent = ReActAgent(context=context, config=agent_config, toolset=toolset)\n\n logger.info(\"=== %s ===\\n\", agent.name)\n logger.info(\"Registered tools: %s\\n\", agent.toolset().list_tools())\n logger.info(\"Tool schema: %s\\n\", agent.toolset().get_all_schemas())\n\n # messages\n messages = [\n ChatMessage(role=\"user\", content=user_input),\n ]\n\n # run agent\n async for response in agent.run(messages):\n logger.info(\"------ response: %s\", response.model_dump_json())\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser(description=\"Demo Agent with Tool\")\n parser.add_argument(\n \"--input\",\n \"-i\",\n type=str,\n default=\"China central bank gold reserves 2025\",\n help=\"User input query (optional, defaults to 'China central bank gold reserves 2025')\",\n )\n args = parser.parse_args()\n\n logging.basicConfig(\n level=logging.DEBUG,\n format=\"%(asctime)s %(levelname)s %(filename)s:%(lineno)d %(message)s\",\n encoding=\"utf-8\", # Add UTF-8 encoding support\n )\n tracer = LocalStorageTracer(storage_dir=\"./traces\")\n ctx = SpanContext(tracer=tracer, app_name=\"demo_agent_with_tool\")\n with ctx.span(name=\"demo_agent_with_tool\") as span:\n asyncio.run(main(args.input))\n" }, { "path": "cortex/examples/demo_checkpoint.py", "content": "import asyncio\nimport logging\nfrom uuid import uuid4\n\nfrom cortex.agents.agent_factory import AgentFactory\nfrom cortex.agents.base_agent import BaseAgent\nfrom cortex.agents.checkpoint_agent.checkpointer import (\n CheckpointStorage,\n SqliteCheckPointer,\n)\nfrom cortex.agents.checkpoint_agent.react_agent import CheckpointReActAgent\nfrom cortex.agents.types import (\n AgentConfig,\n)\nfrom cortex.examples.agents.math_agent import init_math_tools\nfrom cortex.examples.agents.search_agent import init_search_tools\nfrom cortex.model import ModelParams\nfrom cortex.orchestrator.orchestrator import Orchestrator\nfrom cortex.server.http_server import HttpServer\nfrom cortex.tools.toolset import ToolSet\nfrom cortex.tools.ublock_agent_tool import UnblockAgentTool\n\nlogger = logging.getLogger(__name__)\n\n\ndef make_storage() -> CheckpointStorage:\n return SqliteCheckPointer(db_path=\"cp.db\")\n\n\ndef get_search_agent_config() -> AgentConfig:\n return AgentConfig(\n name=\"SearchAgent\",\n description=\"Agent specialized in searching and aggregating information. Uses search tools to find the latest information, aggregates multiple search results, provides comprehensive and accurate answers with source citations. Suitable for finding latest information, aggregating information from multiple sources, answering questions requiring real-time data, and information retrieval organization.\",\n system_prompt=\"You are a professional information search and aggregation assistant.\",\n model=ModelParams(\n name=\"gpt-4o-mini\",\n infer_kwargs={\"max_tokens\": 10000, \"temperature\": 0.7, \"stream\": False},\n ),\n )\n\n\ndef get_math_agent_config() -> AgentConfig:\n return AgentConfig(\n name=\"MathAgent\",\n description=\"Agent specialized in mathematical calculations. Supports basic math operations (addition, subtraction, multiplication, division, exponentiation, square root, etc.), can handle complex mathematical expressions, supports multi-step calculations, and provides detailed calculation process explanations. Suitable for arithmetic operations, algebraic calculations, geometric calculations, and mathematical expression solving.\",\n system_prompt=\"You are a professional mathematical calculation assistant.\",\n model=ModelParams(\n name=\"gpt-4o-mini\",\n infer_kwargs={\"max_tokens\": 10000, \"temperature\": 0.7, \"stream\": False},\n ),\n )\n\n\nasync def make_search_agent(\n config: AgentConfig, context_id: str | None = None\n) -> BaseAgent:\n \"\"\"Create SearchAgent\"\"\"\n toolset = await init_search_tools()\n if context_id is None:\n context_id = uuid4().hex\n\n return CheckpointReActAgent(\n storage=make_storage(),\n context_id=context_id,\n config=config,\n toolset=toolset,\n )\n\n\nasync def make_math_agent(\n config: AgentConfig, context_id: str | None = None\n) -> BaseAgent:\n toolset = await init_math_tools()\n if context_id is None:\n context_id = uuid4().hex\n\n return CheckpointReActAgent(\n storage=make_storage(),\n context_id=context_id,\n config=config,\n toolset=toolset,\n )\n\n\nasync def make_main_agent(\n config: AgentConfig, context_id: str | None = None\n) -> BaseAgent:\n \"\"\"Create MainAgent\"\"\"\n toolset = ToolSet()\n\n # toolset.register(\n # UnblockClientTool(\n # name=\"ask_for_user\",\n # description=\"Ask user for input. Used to get user feedback, confirmation, modification suggestions, etc. Parameters: prompt (required) - message shown to user; context (optional) - context information to help user understand the situation.\",\n # channel=toolset.channel,\n # tool_type=ToolType.ASK_INPUT,\n # timeout=300.0,\n # client_params={\n # \"properties\": {\n # \"prompt\": {\n # \"type\": \"string\",\n # \"description\": \"Message shown to user, explaining what user needs to do (confirm, modify, provide information, etc.)\",\n # },\n # \"context\": {\n # \"type\": \"string\",\n # \"description\": \"Context information to help user understand the current situation, e.g., current plan content, items to confirm, etc.\",\n # },\n # },\n # \"required\": [\"prompt\"],\n # },\n # )\n # )\n\n search_agent = await make_search_agent(config=get_search_agent_config())\n math_agent = await make_math_agent(config=get_math_agent_config())\n\n search_tool_params = search_agent.as_tool()\n math_tool_params = math_agent.as_tool()\n toolset.register(\n UnblockAgentTool(\n name=search_tool_params[\"name\"],\n description=search_tool_params[\"description\"],\n timeout=search_tool_params[\"timeout\"]\n if \"timeout\" in search_tool_params\n else 300,\n channel=toolset.channel,\n )\n )\n toolset.register(\n UnblockAgentTool(\n name=math_tool_params[\"name\"],\n description=math_tool_params[\"description\"],\n timeout=math_tool_params[\"timeout\"]\n if \"timeout\" in math_tool_params\n else 300,\n channel=toolset.channel,\n )\n )\n\n if context_id is None:\n context_id = uuid4().hex\n return CheckpointReActAgent(\n storage=make_storage(),\n context_id=context_id,\n config=config,\n toolset=toolset,\n )\n\n\nasync def main():\n logging.basicConfig(\n level=logging.INFO,\n format=\"%(asctime)s %(levelname)s %(filename)s:%(lineno)d %(message)s\",\n encoding=\"utf-8\", # Add UTF-8 encoding support\n )\n\n agent_factory = AgentFactory()\n agent_factory.register_agent(\n name=\"MainAgent\",\n make_agent_func=make_main_agent,\n default_config=AgentConfig(\n name=\"MainAgent\",\n description=\"Main coordination Agent responsible for coordinating and calling other specialized Agents to complete tasks. Can select appropriate Agents based on task requirements (e.g., MathAgent for mathematical calculations, SearchAgent for information search) and coordinate multiple Agents to complete complex tasks.\",\n system_prompt=\"You are a main coordination Agent responsible for coordinating and calling other specialized Agents to complete tasks. You can select appropriate Agents based on task requirements (e.g., MathAgent for mathematical calculations, SearchAgent for information search) and coordinate multiple Agents to complete complex tasks.\",\n model=ModelParams(\n name=\"gpt-4o-mini\",\n infer_kwargs={\"max_tokens\": 10000, \"temperature\": 0.7, \"stream\": False},\n ),\n unfinished_mode=False,\n ),\n )\n agent_factory.register_agent(\n name=\"SearchAgent\",\n make_agent_func=make_search_agent,\n default_config=get_search_agent_config(),\n )\n agent_factory.register_agent(\n name=\"MathAgent\",\n make_agent_func=make_math_agent,\n default_config=get_math_agent_config(),\n )\n\n # with open(\"graph.mermaid\", \"w\", encoding=\"utf-8\") as f:\n # g = GraphReActAgent(\n # make_checkpointer=lambda: AsyncSqliteSaver.from_conn_string(\n # \"checkpoints.db\"\n # ),\n # config=get_search_agent_config(),\n # toolset=ToolSet(),\n # )\n # f.write(g.draw())\n\n orch = Orchestrator(agent_factory)\n http_server = HttpServer(orch)\n await http_server.start()\n\n\nif __name__ == \"__main__\":\n asyncio.run(main())\n" }, { "path": "cortex/examples/demo_toolset_channel.py", "content": "\"\"\"Channel usage example: Demonstrates complete async communication flow.\"\"\"\n\nimport asyncio\nimport logging\nimport time\nfrom typing import Any, Dict\n\nfrom cortex.tools.base import ToolSchema\nfrom cortex.tools.channel import Channel\nfrom cortex.tools.types import ToolParameters\n\nlogger = logging.getLogger(__name__)\n\n\n# Mock external processing system (server, message queue, etc.)\nclass MockExternalServer:\n \"\"\"Mock external server for processing requests.\"\"\"\n\n def __init__(self):\n \"\"\"Initialize mock server.\"\"\"\n self.received_requests: Dict[str, Dict[str, Any]] = {}\n\n async def process_request(self, tool_name: str, request_id: str, data: Any) -> Any:\n \"\"\"\n Process request (mock async processing).\n\n Args:\n tool_name: Tool name\n request_id: Request ID\n data: Request data\n\n Returns:\n Any: Processing result\n \"\"\"\n # Simulate network delay\n await asyncio.sleep(0.5)\n\n # Store request\n self.received_requests[request_id] = {\n \"tool_name\": tool_name,\n \"data\": data,\n \"timestamp\": time.time(),\n }\n\n # Mock processing logic\n if tool_name == \"calculator\":\n if \"operation\" in data and \"operands\" in data:\n operation = data[\"operation\"]\n operands = data[\"operands\"]\n if operation == \"add\":\n result = sum(operands)\n elif operation == \"multiply\":\n result = 1\n for x in operands:\n result *= x\n else:\n result = f\"Unknown operation: {operation}\"\n return {\"result\": result, \"operation\": operation}\n\n elif tool_name == \"data_processor\":\n if \"data\" in data:\n processed = f\"Processed: {data['data']}\"\n return {\"output\": processed, \"length\": len(data[\"data\"])}\n\n # Default return\n return {\"status\": \"success\", \"data\": data}\n\n\nasync def demo_basic_usage():\n \"\"\"Demo 1: Basic usage flow.\"\"\"\n logger.info(\"=\" * 60)\n logger.info(\"Demo 1: Basic usage flow - Send and wait for response\")\n logger.info(\"=\" * 60)\n\n # Create mock server\n server = MockExternalServer()\n\n # Define send callback function\n async def send_to_server(\n tool_name: str, tool_schema: ToolSchema, tool_parameters: ToolParameters\n ):\n \"\"\"Callback function to send data to server.\"\"\"\n # Extract request_id from tool_parameters.kwargs\n request_id = tool_parameters.kwargs.pop(\"_request_id\", None)\n if request_id is None:\n request_id = f\"req_{hash(tool_name)}_{hash(tool_parameters.parameters)}\"\n\n # Build request data\n request_data = {\n \"parameters\": tool_parameters.parameters,\n **tool_parameters.kwargs,\n }\n\n logger.info(\n f\" [Send] tool={tool_name}, request_id={request_id}, data={request_data}\"\n )\n result = await server.process_request(tool_name, request_id, request_data)\n\n # Simulate setting response via Channel\n # In real scenarios, this might be done asynchronously via network, message queue, etc.\n channel.set_response(request_id, result)\n\n # Create Channel and register send callback\n channel = Channel(on_send=send_to_server)\n\n # Send request 1: Calculator tool\n logger.info(\"\\nSend request 1: Calculator - Addition\")\n request_id1, response1 = await channel.send_request(\n tool_name=\"calculator\",\n data=ToolParameters(\n parameters=\"\", kwargs={\"operation\": \"add\", \"operands\": [10, 20, 30]}\n ),\n tool_schema=ToolSchema(name=\"calculator\", description=\"Calculator tool\"),\n timeout=5.0,\n )\n logger.info(f\" Request ID: {request_id1}\")\n logger.info(f\" Response: {response1}\")\n\n # Send request 2: Data processor\n logger.info(\"\\nSend request 2: Data processor\")\n request_id2, response2 = await channel.send_request(\n tool_name=\"data_processor\",\n data=ToolParameters(parameters=\"\", kwargs={\"data\": \"Hello, World!\"}),\n tool_schema=ToolSchema(name=\"data_processor\", description=\"Data processor tool\"),\n timeout=5.0,\n )\n logger.info(f\" Request ID: {request_id2}\")\n logger.info(f\" Response: {response2}\")\n\n logger.info(f\"\\nNumber of requests received by server: {len(server.received_requests)}\")\n\n\nasync def demo_custom_request_id():\n \"\"\"Demo 2: Using custom request_id.\"\"\"\n logger.info(\"\\n\" + \"=\" * 60)\n logger.info(\"Demo 2: Using custom request_id\")\n logger.info(\"=\" * 60)\n\n server = MockExternalServer()\n channel = Channel()\n\n # Define send callback\n async def send_handler(\n tool_name: str, tool_schema: ToolSchema, tool_parameters: ToolParameters\n ):\n \"\"\"Send handler.\"\"\"\n # Extract request_id from tool_parameters.kwargs\n request_id = tool_parameters.kwargs.pop(\"_request_id\", None)\n if request_id is None:\n request_id = f\"req_{hash(tool_name)}_{hash(tool_parameters.parameters)}\"\n\n # Build request data\n request_data = {\n \"parameters\": tool_parameters.parameters,\n **tool_parameters.kwargs,\n }\n\n logger.info(f\" [Send] request_id={request_id}, data={request_data}\")\n result = await server.process_request(tool_name, request_id, request_data)\n channel.set_response(request_id, result)\n\n # Send request with custom request_id\n custom_id = \"req_custom_001\"\n logger.info(f\"\\nUsing custom request_id: {custom_id}\")\n request_id, response = await channel.send_request(\n tool_name=\"calculator\",\n data=ToolParameters(\n parameters=\"\", kwargs={\"operation\": \"multiply\", \"operands\": [2, 3, 4]}\n ),\n tool_schema=ToolSchema(name=\"calculator\", description=\"Calculator tool\"),\n request_id=custom_id,\n on_send=send_handler,\n timeout=5.0,\n )\n logger.info(f\" Returned request_id: {request_id}\")\n logger.info(f\" Response: {response}\")\n\n\nasync def demo_manual_response():\n \"\"\"Demo 3: Manually set response.\"\"\"\n logger.info(\"\\n\" + \"=\" * 60)\n logger.info(\"Demo 3: Manually set response (separate send and response)\")\n logger.info(\"=\" * 60)\n\n channel = Channel()\n saved_request_id = {\"id\": None} # Used to save request_id\n\n # Define send callback (only send, don't set response immediately)\n async def send_only(\n tool_name: str, tool_schema: ToolSchema, tool_parameters: ToolParameters\n ):\n \"\"\"Only send data, don't set response.\"\"\"\n # Extract request_id from tool_parameters.kwargs\n request_id = tool_parameters.kwargs.pop(\"_request_id\", None)\n if request_id is None:\n request_id = f\"req_{hash(tool_name)}_{hash(tool_parameters.parameters)}\"\n\n saved_request_id[\"id\"] = request_id # Save request_id for later use\n\n # Build request data\n request_data = {\n \"parameters\": tool_parameters.parameters,\n **tool_parameters.kwargs,\n }\n\n logger.info(f\" [Send only] request_id={request_id}, data={request_data}\")\n # In real scenarios, this might be sent to message queue, WebSocket, etc.\n # Response is set asynchronously by another thread/process/service\n\n # Start a background task to simulate async response\n async def delayed_response(delay: float = 1.0):\n \"\"\"Delayed response setting.\"\"\"\n await asyncio.sleep(delay)\n req_id = saved_request_id[\"id\"]\n if req_id:\n result = {\n \"status\": \"completed\",\n \"request_id\": req_id,\n \"data\": \"processed_data\",\n }\n logger.info(f\" [Set response] request_id={req_id}, response={result}\")\n channel.set_response(req_id, result)\n\n # Create send request task\n send_task = asyncio.create_task(\n channel.send_request(\n tool_name=\"async_processor\",\n data=ToolParameters(parameters=\"\", kwargs={\"task\": \"process_data\"}),\n tool_schema=ToolSchema(\n name=\"async_processor\", description=\"Async processor tool\"\n ),\n on_send=send_only,\n timeout=5.0,\n )\n )\n\n # Start delayed response task (auto-processed after request is sent)\n response_task = asyncio.create_task(delayed_response(delay=0.8))\n\n # Wait for request to complete\n request_id, response = await send_task\n logger.info(f\" Request ID: {request_id}\")\n logger.info(f\" Response: {response}\")\n\n await response_task\n\n\nasync def demo_error_handling():\n \"\"\"Demo 4: Error handling.\"\"\"\n logger.info(\"\\n\" + \"=\" * 60)\n logger.info(\"Demo 4: Error handling\")\n logger.info(\"=\" * 60)\n\n channel = Channel()\n\n async def send_with_error(\n tool_name: str, tool_schema: ToolSchema, tool_parameters: ToolParameters\n ):\n \"\"\"Send and set error response.\"\"\"\n # Extract request_id from tool_parameters.kwargs\n request_id = tool_parameters.kwargs.pop(\"_request_id\", None)\n if request_id is None:\n request_id = f\"req_{hash(tool_name)}_{hash(tool_parameters.parameters)}\"\n\n logger.info(f\" [Send] request_id={request_id}\")\n await asyncio.sleep(0.3)\n # Simulate processing failure, set error response\n channel.set_response(request_id, None, error=\"Processing failed: Invalid input\")\n\n logger.info(\"\\nSending request (will return error)...\")\n try:\n request_id, response = await channel.send_request(\n tool_name=\"failing_tool\",\n data=ToolParameters(parameters=\"\", kwargs={\"invalid\": \"data\"}),\n tool_schema=ToolSchema(name=\"failing_tool\", description=\"Failing tool\"),\n on_send=send_with_error,\n timeout=5.0,\n )\n logger.info(f\" Response: {response}\")\n except Exception as e:\n logger.info(f\" Caught error: {type(e).__name__}: {e}\")\n\n # Demo timeout\n logger.info(\"\\nSending request (will timeout)...\")\n\n async def slow_sender(\n tool_name: str, tool_schema: ToolSchema, tool_parameters: ToolParameters\n ):\n \"\"\"Slow sender, won't set response.\"\"\"\n # Extract request_id from tool_parameters.kwargs\n request_id = tool_parameters.kwargs.pop(\"_request_id\", None)\n if request_id is None:\n request_id = f\"req_{hash(tool_name)}_{hash(tool_parameters.parameters)}\"\n\n logger.info(f\" [Send] request_id={request_id} (but won't set response)\")\n await asyncio.sleep(2.0) # Exceeds timeout\n\n try:\n request_id, response = await channel.send_request(\n tool_name=\"slow_tool\",\n data=ToolParameters(parameters=\"\", kwargs={\"slow\": \"data\"}),\n tool_schema=ToolSchema(name=\"slow_tool\", description=\"Slow tool\"),\n on_send=slow_sender,\n timeout=1.0, # 1 second timeout\n )\n except TimeoutError as e:\n logger.info(f\" Caught timeout error: {e}\")\n\n\nasync def demo_concurrent_requests():\n \"\"\"Demo 5: Concurrent requests.\"\"\"\n logger.info(\"\\n\" + \"=\" * 60)\n logger.info(\"Demo 5: Concurrent request handling\")\n logger.info(\"=\" * 60)\n\n server = MockExternalServer()\n channel = Channel()\n\n async def concurrent_sender(\n tool_name: str, tool_schema: ToolSchema, tool_parameters: ToolParameters\n ):\n \"\"\"Concurrent send handler.\"\"\"\n # Extract request_id from tool_parameters.kwargs\n request_id = tool_parameters.kwargs.pop(\"_request_id\", None)\n if request_id is None:\n request_id = f\"req_{hash(tool_name)}_{hash(tool_parameters.parameters)}\"\n\n # Build request data\n request_data = {\n \"parameters\": tool_parameters.parameters,\n **tool_parameters.kwargs,\n }\n\n result = await server.process_request(tool_name, request_id, request_data)\n channel.set_response(request_id, result)\n\n # Create multiple concurrent requests\n logger.info(\"\\nSending 5 concurrent requests...\")\n tasks = []\n for i in range(5):\n task = channel.send_request(\n tool_name=\"calculator\",\n data=ToolParameters(\n parameters=\"\",\n kwargs={\"operation\": \"add\", \"operands\": [i, i + 1, i + 2]},\n ),\n tool_schema=ToolSchema(name=\"calculator\", description=\"Calculator tool\"),\n on_send=concurrent_sender,\n timeout=5.0,\n )\n tasks.append(task)\n\n # Wait for all requests to complete\n results = await asyncio.gather(*tasks)\n logger.info(f\"\\nCompleted {len(results)} requests:\")\n for i, (req_id, response) in enumerate(results, 1):\n logger.info(\n f\" Request {i}: request_id={req_id}, result={response.get('result', 'N/A')}\"\n )\n\n\nasync def main():\n \"\"\"Main function: Run all demos.\"\"\"\n logger.info(\"\\n\" + \"=\" * 60)\n logger.info(\"Channel Complete Usage Flow Demo\")\n logger.info(\"=\" * 60)\n\n await demo_basic_usage()\n await demo_custom_request_id()\n await demo_manual_response()\n await demo_error_handling()\n await demo_concurrent_requests()\n\n logger.info(\"\\n\" + \"=\" * 60)\n logger.info(\"All demos completed!\")\n logger.info(\"=\" * 60)\n\n\nif __name__ == \"__main__\":\n asyncio.run(main())\n" }, { "path": "cortex/examples/server.py", "content": "import asyncio\nimport logging\n\nfrom agentkit.trace import LocalStorageTracer\n\nfrom cortex.agents.agent_factory import AgentFactory\nfrom cortex.agents.types import AgentConfig\nfrom cortex.examples.agents.main_agent import make_main_agent\nfrom cortex.examples.agents.math_agent import get_math_agent_config, make_math_agent\nfrom cortex.examples.agents.search_agent import (\n get_search_agent_config,\n make_search_agent,\n)\nfrom cortex.model import ModelParams\nfrom cortex.orchestrator.orchestrator import Orchestrator\nfrom cortex.server.http_server import HttpServer\n\nlogger = logging.getLogger(__name__)\n\n\nasync def main():\n agent_factory = AgentFactory()\n agent_factory.register_agent(\n name=\"MainAgent\",\n make_agent_func=make_main_agent,\n default_config=AgentConfig(\n name=\"MainAgent\",\n description=\"Main coordination Agent responsible for coordinating and calling other specialized Agents to complete tasks. Can select appropriate Agents based on task requirements (e.g., MathAgent for mathematical calculations, SearchAgent for information search) and coordinate multiple Agents to complete complex tasks.\",\n system_prompt=\"You are a main coordination Agent responsible for coordinating and calling other specialized Agents to complete tasks.\",\n model=ModelParams(\n name=\"gpt-4o-mini\",\n infer_kwargs={\"max_tokens\": 2000, \"temperature\": 0.7, \"stream\": False},\n ),\n ),\n )\n agent_factory.register_agent(\n name=\"SearchAgent\",\n make_agent_func=make_search_agent,\n default_config=get_search_agent_config(),\n )\n agent_factory.register_agent(\n name=\"MathAgent\",\n make_agent_func=make_math_agent,\n default_config=get_math_agent_config(),\n )\n orch = Orchestrator(agent_factory)\n tracer = LocalStorageTracer(storage_dir=\"./traces\")\n http_server = HttpServer(orch, tracer=tracer)\n await http_server.start()\n\n\nif __name__ == \"__main__\":\n asyncio.run(main())\n" }, { "path": "cortex/model/__init__.py", "content": "\"\"\"Model system for Agent components.\"\"\"\n\nfrom enum import Enum\nfrom typing import AsyncGenerator\n\nfrom agentkit.trace import get_current_context\nfrom .definition import ChatMessage, ModelParams\nfrom cortex.model.utils import merge_delta_message\nfrom pydantic import BaseModel\n\nfrom cortex.model.provider import ModelProvider\n\n__all__ = [\n \"ModelAPI\",\n \"ModelMessage\",\n \"MessageType\",\n \"ModelParams\",\n]\n\n\nclass MessageType(str, Enum):\n \"\"\"Message type enumeration\"\"\"\n\n DELTA = \"delta\" # Streaming output\n ACCUMULATED = \"accumulated\" # Accumulated output\n\n\nclass ModelMessage(BaseModel):\n \"\"\"Model message\"\"\"\n\n message: ChatMessage\n message_type: MessageType\n\n\nclass ModelAPI:\n \"\"\"Model API\"\"\"\n\n def __init__(self, provider: ModelProvider):\n self.provider = provider\n\n async def chat_completion(\n self,\n messages: list[ChatMessage],\n tools: list | None = None,\n log_file: str | None = None,\n trace_request: dict | None = None,\n ) -> ModelMessage:\n \"\"\"Call model API, return accumulated message\"\"\"\n ctx = get_current_context()\n with ctx.llm_span(name=\"ModelAPI.chat_completion\") as span:\n span.update_payload_data(\n request=trace_request\n or {\n # \"model_params\": self.params(),\n \"messages\": messages,\n \"tools\": tools,\n },\n tools=tools,\n )\n response = await self.provider.chat_completion(\n messages=messages, tools=tools, log_file=log_file\n )\n span.update_payload_data(\n response=ModelMessage(\n message=response, message_type=MessageType.ACCUMULATED\n ),\n )\n return ModelMessage(message=response, message_type=MessageType.ACCUMULATED)\n\n async def chat_completion_stream(\n self,\n messages: list[ChatMessage],\n tools: list | None = None,\n log_file: str | None = None,\n trace_request: dict | None = None,\n ) -> AsyncGenerator[ModelMessage, None]:\n \"\"\"Call model API, return streaming messages\"\"\"\n accumulated_message = None\n async for event in self.provider.chat_completion_stream(\n messages=messages, tools=tools, log_file=log_file\n ):\n merged_dict = merge_delta_message(\n accumulated_message.model_dump() if accumulated_message else None,\n event.model_dump() if event else None,\n )\n accumulated_message = ChatMessage(**merged_dict)\n yield ModelMessage(message=event, message_type=MessageType.DELTA)\n\n # Yield the accumulated complete message\n if accumulated_message:\n yield ModelMessage(\n message=accumulated_message, message_type=MessageType.ACCUMULATED\n )\n\n ctx = get_current_context()\n with ctx.llm_span(name=\"ModelAPI.chat_completion_stream\") as span:\n span.update_payload_data(\n request=trace_request\n or {\n # \"model_params\": self.params(),\n \"messages\": messages,\n \"tools\": tools,\n },\n )\n span.update_payload_data(\n response=ModelMessage(\n message=accumulated_message, message_type=MessageType.ACCUMULATED\n ),\n )\n" }, { "path": "cortex/model/definition.py", "content": "from enum import Enum\nfrom typing import Any, Dict, List, Optional\nfrom typing_extensions import TypedDict\n\nfrom pydantic import BaseModel, Field\n\n\nclass MessageRole(Enum):\n USER = \"user\"\n HUMAN = \"human\"\n ASSISTANT = \"assistant\"\n SYSTEM = \"system\"\n TOOL_DESC = \"tool-description\"\n TOOL_RESPONSE = \"tool-response\"\n TOOL = \"tool\"\n\n\nclass ContentBlockType(Enum):\n TEXT = \"text\"\n THINK = \"thinking\"\n REDACTED_THINK = \"redacted_thinking\"\n TOOLCALL = \"toolcall\"\n TOOLUSE = \"tool_use\"\n TOOLRESULT = \"tool_result\"\n IMAGE = \"image\"\n IMAGE_URL = \"image_url\"\n VIDEO = \"video\"\n VIDEO_URL = \"video_url\"\n AUDIO = \"input_audio\"\n AUDIO_URL = \"audio_url\"\n DOCURL = \"doc_url\"\n\n\nclass ModelParams(BaseModel):\n name: str\n response_format: BaseModel | None = None\n toolcall_parser_version: str | None = None\n parallel_tool_calls: bool = True\n infer_kwargs: dict = {}\n user_role: str = MessageRole.USER.value\n tool_role: str = MessageRole.TOOL.value\n explicit_model_family_value: str | None = None\n explicit_api_base: str | None = None\n explicit_api_key: str | None = None\n\n\nclass Function(BaseModel):\n arguments: str | None = None\n \"\"\"\n The arguments to call the function with, as generated by the model in JSON\n format. Note that the model does not always generate valid JSON, and may\n hallucinate parameters not defined by your function schema. Validate the\n arguments in your code before calling your function.\n \"\"\"\n\n name: str | None = None\n \"\"\"The name of the function to call.\"\"\"\n\n\nclass ChatToolCall(BaseModel):\n index: int | None = None\n \"\"\"The index of the tool call.\"\"\"\n\n id: str | None = None\n \"\"\"The ID of the tool call.\"\"\"\n\n function: Function\n \"\"\"The function that the model called.\"\"\"\n\n type: str | None = \"function\"\n \"\"\"The type of the tool. Currently, only `function` is supported.\"\"\"\n\n\nclass ExtraInfo(TypedDict, total=False):\n \"\"\"Extra information for ChatMessage.\"\"\"\n\n cache_msg_id: str\n \"\"\"The cache message ID.\n\n Try best to leverage the cache functionality of the backend model service.\n But it's not guaranteed to be successful.\n Currently only Anthropic model of models-proxy backend supports this functionality.\n \"\"\"\n\n usage: Dict[str, Any]\n \"\"\"The usage information of the model. \"\"\"\n\n finish_reason: str\n \"\"\"The finish reason of the model.\n\n For OpenAI model, it's called \"finish_reason\".\n For Anthropic model, it's called \"stop_reason\".\n Here we use the unified name \"finish_reason\" for all cases.\n \"\"\"\n\n\nclass ChatMessage(BaseModel):\n id: Optional[str] = None\n role: str | None = None\n content: str | List[Dict] | None = None\n tool_call_id: Optional[str] = None\n tool_calls: Optional[List[ChatToolCall]] = Field(default_factory=list)\n\n # for train use\n extra_info: Optional[ExtraInfo] = None\n\n @classmethod\n def from_dict(cls, data: dict) -> \"ChatMessage\":\n return cls(**data)\n\n def to_dict(self) -> dict:\n return self.model_dump()\n" }, { "path": "cortex/model/provider.py", "content": "from abc import ABC, abstractmethod\nfrom typing import AsyncGenerator\n\nfrom .definition import ChatMessage\n\n\nclass ModelProvider(ABC):\n @abstractmethod\n async def chat_completion_stream(\n self,\n messages: list[ChatMessage],\n tools: list | None = None,\n log_file: str | None = None,\n ) -> AsyncGenerator[ChatMessage, None]:\n pass\n\n @abstractmethod\n async def chat_completion(\n self,\n messages: list[ChatMessage],\n tools: list | None = None,\n log_file: str | None = None,\n ) -> ChatMessage:\n pass\n" }, { "path": "cortex/model/stepfun_chat.py", "content": "\"\"\"StepFun Chat API Client with reasoning support.\"\"\"\n\nimport json\nfrom typing import Any, AsyncGenerator, Iterable\n\nimport httpx\nfrom pydantic import BaseModel\n\n\n# StepFun API supported infer_kwargs parameters\nSUPPORTED_INFER_KWARGS = frozenset([\n \"temperature\", # 0.0-2.0, default 0.5\n \"top_p\", # default 0.9\n \"max_tokens\",\n \"n\", # default 1\n \"stop\",\n \"frequency_penalty\", # 0.0-1.0, default 0\n \"response_format\",\n \"reasoning_format\", # StepFun specific: \"general\" or \"deepseek-style\"\n])\n\n\n# ============ Structured Response Types ============\n\nclass Function(BaseModel):\n \"\"\"Function call in tool_calls.\"\"\"\n name: str | None = None\n arguments: str | None = None\n\n\nclass ToolCall(BaseModel):\n \"\"\"Tool call object.\"\"\"\n id: str | None = None\n index: int | None = None\n type: str = \"function\"\n function: Function | None = None\n\n\nclass Message(BaseModel):\n \"\"\"Chat completion message with reasoning support.\"\"\"\n role: str | None = None\n content: str | None = None\n reasoning: str | None = None # StepFun specific\n tool_calls: list[ToolCall] | None = None\n\n\nclass Delta(BaseModel):\n \"\"\"Streaming delta with reasoning support.\"\"\"\n role: str | None = None\n content: str | None = None\n reasoning: str | None = None # StepFun specific\n tool_calls: list[ToolCall] | None = None\n\n\nclass Choice(BaseModel):\n \"\"\"Chat completion choice.\"\"\"\n index: int = 0\n message: Message | None = None\n finish_reason: str | None = None\n\n\nclass StreamChoice(BaseModel):\n \"\"\"Streaming choice with delta.\"\"\"\n index: int = 0\n delta: Delta | None = None\n finish_reason: str | None = None\n\n\nclass Usage(BaseModel):\n \"\"\"Token usage information.\"\"\"\n prompt_tokens: int = 0\n completion_tokens: int = 0\n total_tokens: int = 0\n\n\nclass ChatCompletion(BaseModel):\n \"\"\"Chat completion response.\"\"\"\n id: str | None = None\n object: str = \"chat.completion\"\n created: int | None = None\n model: str | None = None\n choices: list[Choice] = []\n usage: Usage | None = None\n\n\nclass ChatCompletionChunk(BaseModel):\n \"\"\"Streaming chat completion chunk.\"\"\"\n id: str | None = None\n object: str = \"chat.completion.chunk\"\n created: int | None = None\n model: str | None = None\n choices: list[StreamChoice] = []\n usage: Usage | None = None\n\n\n# ============ Client ============\n\nclass StepFunClient:\n \"\"\"StepFun API client using httpx.\n \n Returns raw OpenAI format data, with an additional reasoning field in message.\n \"\"\"\n\n DEFAULT_API_BASE = \"https://api.stepfun.com\"\n CHAT_COMPLETIONS_ENDPOINT = \"/v1/chat/completions\"\n\n def __init__(\n self,\n api_key: str,\n api_base: str | None = None,\n timeout: float = 120.0,\n ):\n self.api_key = api_key\n self.api_base = (api_base or self.DEFAULT_API_BASE).rstrip(\"/\")\n self.timeout = timeout\n\n def _build_headers(self) -> dict[str, str]:\n return {\n \"Authorization\": f\"Bearer {self.api_key}\",\n \"Content-Type\": \"application/json\",\n }\n\n def _build_request_body(\n self,\n model: str,\n messages: Iterable[dict[str, Any]],\n tools: Iterable[dict[str, Any]] | None = None,\n infer_kwargs: dict[str, Any] | None = None,\n stream: bool = False,\n ) -> dict[str, Any]:\n \"\"\"Build request body for chat completion.\"\"\"\n body: dict[str, Any] = {\n \"model\": model,\n \"messages\": list(messages),\n \"stream\": stream,\n }\n\n if tools:\n body[\"tools\"] = list(tools)\n\n # Add supported parameters from infer_kwargs\n if infer_kwargs:\n for key in SUPPORTED_INFER_KWARGS:\n if key in infer_kwargs:\n body[key] = infer_kwargs[key]\n\n return body\n\n async def chat_completion(\n self,\n model: str,\n messages: Iterable[dict[str, Any]],\n tools: Iterable[dict[str, Any]] | None = None,\n infer_kwargs: dict[str, Any] | None = None,\n ) -> ChatCompletion:\n \"\"\"Make a non-streaming chat completion request.\n \n Returns:\n ChatCompletion structured object, choices[0].message contains reasoning field\n \"\"\"\n url = f\"{self.api_base}{self.CHAT_COMPLETIONS_ENDPOINT}\"\n body = self._build_request_body(\n model=model,\n messages=messages,\n tools=tools,\n infer_kwargs=infer_kwargs,\n stream=False,\n )\n\n async with httpx.AsyncClient(timeout=self.timeout) as client:\n response = await client.post(\n url,\n headers=self._build_headers(),\n json=body,\n )\n response.raise_for_status()\n data = response.json()\n\n return ChatCompletion.model_validate(data)\n\n async def chat_completion_stream(\n self,\n model: str,\n messages: Iterable[dict[str, Any]],\n tools: Iterable[dict[str, Any]] | None = None,\n infer_kwargs: dict[str, Any] | None = None,\n ) -> AsyncGenerator[ChatCompletionChunk, None]:\n \"\"\"Make a streaming chat completion request.\n \n Yields:\n ChatCompletionChunk structured object, delta contains reasoning field\n \"\"\"\n url = f\"{self.api_base}{self.CHAT_COMPLETIONS_ENDPOINT}\"\n body = self._build_request_body(\n model=model,\n messages=messages,\n tools=tools,\n infer_kwargs=infer_kwargs,\n stream=True,\n )\n\n async with httpx.AsyncClient(timeout=self.timeout) as client:\n async with client.stream(\n \"POST\",\n url,\n headers=self._build_headers(),\n json=body,\n ) as response:\n response.raise_for_status()\n\n async for line in response.aiter_lines():\n if not line or not line.startswith(\"data: \"):\n continue\n\n data = line[6:] # Remove \"data: \" prefix\n if data == \"[DONE]\":\n break\n\n try:\n chunk = json.loads(data)\n yield ChatCompletionChunk.model_validate(chunk)\n except json.JSONDecodeError:\n continue\n" }, { "path": "cortex/model/stepfun_provider.py", "content": "\"\"\"StepFun Model Provider implementation.\"\"\"\n\nimport asyncio\nimport logging\nimport re\nfrom typing import Any, AsyncGenerator, Callable, Awaitable\n\nimport httpx\n\nfrom .definition import ChatMessage\nfrom cortex.model.definition import ChatToolCall, ContentBlockType, ExtraInfo, MessageRole, ModelParams\n\nfrom cortex.model.provider import ModelProvider\n\nfrom .stepfun_chat import (\n ChatCompletion,\n ChatCompletionChunk,\n Message,\n Delta,\n StepFunClient,\n)\n\n# Regex pattern for matching ... tags\nTHINK_TAG_PATTERN = re.compile(r\"(.*?)\", re.DOTALL)\n\n\nclass StepFunModelProvider(ModelProvider):\n \"\"\"Model provider for StepFun API with reasoning support.\"\"\"\n\n # Constants for matching tags\n _THINK_OPEN_TAG = \"\"\n _THINK_CLOSE_TAG = \"\"\n\n def __init__(self, model_params: ModelParams):\n self.model_params = model_params\n # Used to track tag state in streaming scenarios\n self._stream_in_think_tag = False\n # Used to buffer potentially incomplete tags\n self._stream_tag_buffer = \"\"\n self._logger = logging.getLogger(__name__)\n\n async def _call_with_retry(\n self,\n description: str,\n func: Callable[[], Awaitable[Any]],\n max_attempts: int = 5,\n ) -> Any:\n \"\"\"Retry wrapper with linear backoff (2s,4s,6s,8s).\"\"\"\n last_exc: Exception | None = None\n for attempt in range(1, max_attempts + 1):\n try:\n return await func()\n except Exception as exc: # noqa: BLE001\n last_exc = exc\n detail = \"\"\n if isinstance(exc, httpx.HTTPStatusError):\n try:\n detail = f\" body={exc.response.text}\"\n except Exception:\n detail = \"\"\n if attempt >= max_attempts:\n break\n delay = 2 * attempt\n self._logger.warning(\n \"StepFun %s failed (attempt %s/%s): %s%s; retrying in %ss\",\n description,\n attempt,\n max_attempts,\n exc,\n detail,\n delay,\n )\n await asyncio.sleep(delay)\n assert last_exc is not None\n raise last_exc\n\n async def _stream_with_retry(\n self,\n description: str,\n stream_factory: Callable[[], AsyncGenerator[Any, None]],\n max_attempts: int = 5,\n ) -> AsyncGenerator[Any, None]:\n \"\"\"Retry wrapper for streaming calls.\"\"\"\n last_exc: Exception | None = None\n for attempt in range(1, max_attempts + 1):\n try:\n async for chunk in stream_factory():\n yield chunk\n return\n except Exception as exc: # noqa: BLE001\n last_exc = exc\n if attempt >= max_attempts:\n break\n delay = 2 * attempt\n self._logger.warning(\n \"StepFun %s stream failed (attempt %s/%s): %s; retrying in %ss\",\n description,\n attempt,\n max_attempts,\n exc,\n delay,\n )\n await asyncio.sleep(delay)\n assert last_exc is not None\n raise last_exc\n\n def _create_client(self) -> StepFunClient:\n \"\"\"Create a StepFunClient from model params.\"\"\"\n api_key = self.model_params.explicit_api_key\n if not api_key:\n raise ValueError(\"StepFun API key is required\")\n\n transport_timeout = None\n if isinstance(self.model_params.infer_kwargs, dict):\n transport_timeout = (\n self.model_params.infer_kwargs.get(\"request_timeout\")\n or self.model_params.infer_kwargs.get(\"timeout\")\n )\n\n return StepFunClient(\n api_key=api_key,\n api_base=self.model_params.explicit_api_base,\n timeout=float(transport_timeout) if transport_timeout else 120.0,\n )\n\n def _chat_messages_to_openai(\n self,\n messages: list[ChatMessage] | list[dict],\n ) -> list[dict[str, Any]]:\n \"\"\"Convert ChatMessage list to OpenAI format messages.\n \n References openai_model.py's chat_message_to_openai_messages implementation,\n handles complex ChatMessage formats (including content blocks, tool results, etc.).\n \"\"\"\n openai_messages = []\n for message in messages:\n if isinstance(message, dict):\n message = ChatMessage(**message)\n\n if isinstance(message.content, str) or message.content is None:\n new_message = message.model_dump()\n elif isinstance(message.content, list):\n new_content = []\n for block in message.content:\n if block[\"type\"] == ContentBlockType.TEXT.value:\n # Merge consecutive text blocks\n if (\n new_content\n and new_content[-1][\"type\"] == ContentBlockType.TEXT.value\n ):\n new_content[-1][\"text\"] += block.get(\"text\", block.get(block[\"type\"], \"\"))\n else:\n new_content.append(block)\n\n elif block[\"type\"] == ContentBlockType.THINK.value:\n # Convert thinking content to text wrapped in tags\n think_content = block.get(block[\"type\"], \"\")\n if (\n new_content\n and new_content[-1][\"type\"] == ContentBlockType.TEXT.value\n ):\n new_content[-1][\"text\"] += f\"{think_content}\"\n else:\n new_content.append({\n \"type\": ContentBlockType.TEXT.value,\n \"text\": f\"{think_content}\",\n })\n\n elif block[\"type\"] == ContentBlockType.REDACTED_THINK.value:\n redacted_content = block.get(\"data\", \"\")\n if (\n new_content\n and new_content[-1][\"type\"] == ContentBlockType.TEXT.value\n ):\n new_content[-1][\"text\"] += f\"{redacted_content}\"\n else:\n new_content.append({\n \"type\": ContentBlockType.TEXT.value,\n \"text\": f\"{redacted_content}\",\n })\n\n elif block[\"type\"] == ContentBlockType.TOOLRESULT.value:\n # Tool result requires special handling\n tool_block_content = block.get(\"content\", [])\n if isinstance(tool_block_content, str):\n tool_block_content = [{\n \"type\": ContentBlockType.TEXT.value,\n \"text\": tool_block_content,\n }]\n \n # First text as tool role message\n openai_messages.append({\n \"role\": message.role,\n \"content\": tool_block_content[0].get(\"text\", \"\") if tool_block_content else \"\",\n \"tool_call_id\": block.get(\"tool_use_id\", message.tool_call_id),\n })\n \n # Remaining content as user message\n if len(tool_block_content) > 1:\n extra_message = {\n \"role\": MessageRole.USER.value,\n \"content\": tool_block_content[1:],\n }\n openai_messages.append(extra_message)\n else:\n # Other block types (e.g., image_url, etc.) keep as is\n new_content.append(block)\n\n new_message = message.model_dump()\n new_message[\"content\"] = new_content\n\n # Handle tool_calls\n if message.tool_calls:\n normalized_tool_calls = []\n for tc in message.tool_calls:\n tc_dict = tc.model_dump()\n # Ensure function and arguments are present for API validation\n func = tc_dict.get(\"function\") or {}\n if func.get(\"arguments\") in (None, \"\"):\n func[\"arguments\"] = \"{}\"\n if not func.get(\"name\"):\n func[\"name\"] = tc_dict.get(\"id\") or \"unknown\"\n tc_dict[\"function\"] = func\n normalized_tool_calls.append(tc_dict)\n new_message[\"tool_calls\"] = normalized_tool_calls\n openai_messages.append(new_message)\n elif new_message.get(\"content\"):\n openai_messages.append(new_message)\n\n return openai_messages\n\n def _extract_think_from_content(self, content: str | None) -> tuple[str | None, str | None]:\n \"\"\"Extract content from ... tags in content.\n \n Args:\n content: Original content string\n \n Returns:\n tuple[reasoning, remaining_content]: \n - reasoning: Content inside tags, or None if not present\n - remaining_content: Content after removing tags\n \"\"\"\n if not content or not isinstance(content, str):\n return None, content\n \n # Find all ... tags\n matches = THINK_TAG_PATTERN.findall(content)\n if not matches:\n return None, content\n \n # Merge all think content\n reasoning = \"\".join(matches)\n \n # Remove all ... tags\n remaining_content = THINK_TAG_PATTERN.sub(\"\", content).strip()\n \n return reasoning if reasoning else None, remaining_content if remaining_content else None\n\n def _message_to_chat_message(self, message: Message) -> ChatMessage:\n \"\"\"Convert StepFun Message (with reasoning) to ChatMessage.\n \n StepFun message format:\n - role: str\n - content: str | None (may contain ... tags)\n - reasoning: str | None (StepFun specific, may be empty)\n - tool_calls: list[ToolCall] | None\n \n reasoning may be in the standalone reasoning field or in tags in content.\n \"\"\"\n reasoning = message.reasoning\n content = message.content\n\n # If reasoning field is empty, try to extract tag content from content\n if not reasoning and isinstance(content, str):\n extracted_reasoning, remaining_content = self._extract_think_from_content(content)\n if extracted_reasoning:\n reasoning = extracted_reasoning\n content = remaining_content\n\n # Process reasoning and content separately, build content blocks\n new_content: list[dict] = []\n \n # Handle reasoning field, convert to thinking content block\n if reasoning:\n new_content.append({\n \"type\": ContentBlockType.THINK.value,\n ContentBlockType.THINK.value: reasoning,\n })\n \n # Handle content field\n if isinstance(content, str) and content:\n new_content.append({\n \"type\": ContentBlockType.TEXT.value,\n \"text\": content,\n })\n elif isinstance(content, list):\n new_content.extend(content)\n\n # Determine final content: use list if there are blocks, otherwise keep original value\n final_content: str | list | None\n if new_content:\n final_content = new_content\n else:\n final_content = content # Could be None or empty string\n\n # Handle tool_calls\n tool_calls = None\n if message.tool_calls:\n tool_calls = [\n ChatToolCall(**tc.model_dump(exclude_none=True))\n for tc in message.tool_calls\n ]\n\n return ChatMessage(\n role=message.role,\n content=final_content,\n tool_calls=tool_calls,\n )\n\n def _process_stream_content_for_think(self, content: str | None) -> tuple[str | None, str | None]:\n \"\"\"Process content in streaming scenarios to identify tags.\n \n Uses character-level state machine to handle tags spanning multiple chunks:\n - Tags may be split across chunks (e.g., '')\n - Uses _stream_tag_buffer to cache potentially incomplete tag fragments\n - Uses _stream_in_think_tag to track whether inside a tag\n \n Args:\n content: Current chunk's content\n \n Returns:\n tuple[reasoning, text_content]:\n - reasoning: If inside tag, returns current reasoning content\n - text_content: If not inside tag, returns plain text content\n \"\"\"\n if not content:\n return None, None\n \n # Merge cached content with new content\n buffer = self._stream_tag_buffer + content\n self._stream_tag_buffer = \"\"\n \n reasoning_parts = []\n text_parts = []\n \n i = 0\n while i < len(buffer):\n if not self._stream_in_think_tag:\n # Not inside tag, look for opening tag\n # Check if this could be the start of a tag\n if buffer[i] == '<':\n # Check if remaining content is sufficient to determine\n remaining = buffer[i:]\n if remaining.startswith(self._THINK_OPEN_TAG):\n # Complete tag\n self._stream_in_think_tag = True\n i += len(self._THINK_OPEN_TAG)\n elif self._THINK_OPEN_TAG.startswith(remaining):\n # Possibly incomplete tag (e.g., ' tag, treat as plain text\n text_parts.append(buffer[i])\n i += 1\n else:\n text_parts.append(buffer[i])\n i += 1\n else:\n # Inside tag, look for closing tag\n if buffer[i] == '<':\n # Check if this could be a tag\n remaining = buffer[i:]\n if remaining.startswith(self._THINK_CLOSE_TAG):\n # Complete tag\n self._stream_in_think_tag = False\n i += len(self._THINK_CLOSE_TAG)\n elif self._THINK_CLOSE_TAG.startswith(remaining):\n # Possibly incomplete tag\n self._stream_tag_buffer = remaining\n break\n else:\n # Not a tag, treat as reasoning content\n reasoning_parts.append(buffer[i])\n i += 1\n else:\n reasoning_parts.append(buffer[i])\n i += 1\n \n reasoning = \"\".join(reasoning_parts) if reasoning_parts else None\n text_content = \"\".join(text_parts) if text_parts else None\n \n return reasoning, text_content\n\n def _delta_to_chat_message(\n self,\n delta: Delta,\n chunk_id: str | None = None,\n ) -> ChatMessage:\n \"\"\"Convert StepFun Delta (with reasoning) to ChatMessage.\n \n StepFun stream delta format:\n - role: str | None\n - content: str | None (may contain ... tag fragments)\n - reasoning: str | None (StepFun specific, may be empty)\n - tool_calls: list[ToolCall] | None\n \n reasoning may be in the standalone reasoning field or in tags in content.\n In streaming scenarios, tags may span multiple chunks.\n \"\"\"\n reasoning = delta.reasoning\n content = delta.content\n\n # If reasoning field is empty, try to extract tag content from content\n if not reasoning and isinstance(content, str):\n extracted_reasoning, remaining_content = self._process_stream_content_for_think(content)\n reasoning = extracted_reasoning\n content = remaining_content\n\n # Process reasoning and content separately\n new_content: list[dict] = []\n \n if reasoning:\n new_content.append({\n \"type\": ContentBlockType.THINK.value,\n ContentBlockType.THINK.value: reasoning,\n })\n \n if isinstance(content, str) and content:\n new_content.append({\n \"type\": ContentBlockType.TEXT.value,\n \"text\": content,\n })\n\n # Determine final content\n final_content: str | list | None\n if new_content:\n final_content = new_content\n else:\n final_content = None # Return None when streaming has no content\n\n # Handle tool_calls\n tool_calls = None\n if delta.tool_calls:\n tool_calls = [\n ChatToolCall(**tc.model_dump(exclude_none=True))\n for tc in delta.tool_calls\n ]\n\n return ChatMessage(\n id=chunk_id,\n role=delta.role,\n content=final_content,\n tool_calls=tool_calls,\n )\n\n async def chat_completion(\n self,\n messages: list[ChatMessage],\n tools: list | None = None,\n log_file: str | None = None,\n ) -> ChatMessage:\n \"\"\"Make a non-streaming chat completion request.\"\"\"\n client = self._create_client()\n openai_messages = self._chat_messages_to_openai(messages)\n infer_kwargs = dict(self.model_params.infer_kwargs or {})\n infer_kwargs.pop(\"request_timeout\", None)\n infer_kwargs.pop(\"timeout\", None)\n\n response: ChatCompletion = await self._call_with_retry(\n description=\"chat_completion\",\n func=lambda: client.chat_completion(\n model=self.model_params.name,\n messages=openai_messages,\n tools=tools,\n infer_kwargs=infer_kwargs,\n ),\n )\n\n # Extract and convert message\n if not response.choices:\n return ChatMessage()\n\n choice = response.choices[0]\n if not choice.message:\n return ChatMessage()\n\n result = self._message_to_chat_message(choice.message)\n\n # Add usage and finish_reason\n result.extra_info = ExtraInfo()\n if response.usage:\n result.extra_info[\"usage\"] = response.usage.model_dump()\n if choice.finish_reason:\n result.extra_info[\"finish_reason\"] = choice.finish_reason\n\n return result\n\n async def chat_completion_stream(\n self,\n messages: list[ChatMessage],\n tools: list | None = None,\n log_file: str | None = None,\n ) -> AsyncGenerator[ChatMessage, None]:\n \"\"\"Make a streaming chat completion request.\"\"\"\n client = self._create_client()\n openai_messages = self._chat_messages_to_openai(messages)\n infer_kwargs = dict(self.model_params.infer_kwargs or {})\n infer_kwargs.pop(\"request_timeout\", None)\n infer_kwargs.pop(\"timeout\", None)\n\n async def stream_factory() -> AsyncGenerator[ChatCompletionChunk, None]:\n # Reset streaming state each attempt\n self._stream_in_think_tag = False\n self._stream_tag_buffer = \"\"\n async for chunk in client.chat_completion_stream(\n model=self.model_params.name,\n messages=openai_messages,\n tools=tools,\n infer_kwargs=infer_kwargs,\n ):\n yield chunk\n\n async for chunk in self._stream_with_retry(\n description=\"chat_completion_stream\",\n stream_factory=stream_factory,\n ):\n chunk: ChatCompletionChunk\n if not chunk.choices:\n continue\n\n choice = chunk.choices[0]\n if not choice.delta:\n continue\n\n result = self._delta_to_chat_message(choice.delta, chunk.id)\n\n # Skip this chunk if no valid content\n if result.content is None and result.role is None and result.tool_calls is None:\n continue\n\n # Add usage and finish_reason (if present)\n if chunk.usage or choice.finish_reason:\n result.extra_info = ExtraInfo()\n if chunk.usage:\n result.extra_info[\"usage\"] = chunk.usage.model_dump()\n if choice.finish_reason:\n result.extra_info[\"finish_reason\"] = choice.finish_reason\n\n yield result\n" }, { "path": "cortex/model/utils.py", "content": "import copy\n\n\ndef merge_delta_message(d1: dict | None, d2: dict | None) -> dict:\n \"\"\"\n Merge two message dictionaries.\n\n Args:\n d1: Accumulated message dictionary, can be dict or None\n d2: Delta message dictionary, can be dict or None\n\n Returns:\n Merged dictionary\n\n Note:\n role and id fields will retain d1's values (if present), so during streaming\n the role from the first chunk can be correctly preserved.\n \"\"\"\n if d1 is None:\n return d2\n if d2 is None:\n return d1\n\n result = copy.deepcopy(d1) # Create a copy of d1\n for key, value in d2.items():\n if key in result:\n if key == \"index\":\n pass\n elif key == \"role\" or key == \"id\":\n # Retain d1's value for role and id (if present), otherwise use d2's value\n if not result[key] and value:\n result[key] = value\n elif key == \"tool_calls\":\n if value and not isinstance(value, list):\n value = [value]\n toolcall = {}\n if result[key]:\n for v in result[key]:\n toolcall[v[\"index\"]] = v\n if value:\n for v in value:\n toolcall_id = v[\"index\"]\n if toolcall_id in toolcall:\n toolcall[toolcall_id] = merge_delta_message(\n toolcall[toolcall_id], v\n )\n else:\n toolcall[toolcall_id] = v\n result[key] = []\n for k, v in toolcall.items():\n result[key].append(v)\n elif isinstance(result[key], dict) and isinstance(value, dict):\n # Recursively merge nested dictionaries\n result[key] = merge_delta_message(result[key], value)\n elif isinstance(result[key], str) and isinstance(value, str) and value:\n # type/id are fixed enum-like values, don't concatenate (avoids \"functionfunction...\")\n if key in (\"type\", \"id\"):\n result[key] = value\n else:\n result[key] += value\n elif isinstance(result[key], list) and isinstance(value, list):\n result[key] += value\n elif isinstance(result[key], (int, float)) and isinstance(\n value, (int, float)\n ):\n result[key] += value\n elif value:\n # Type mismatch, override with d2's value\n result[key] = value\n else:\n result[key] = value\n # Merge data like [{'type': 'text', 'text': '...'}], combining consecutive items with the same type into one\n for key, value in result.items():\n if isinstance(value, list):\n new_value = []\n last_type = None\n last_index = None\n for i in range(len(value)):\n if \"type\" in value[i]:\n current_type = value[i][\"type\"]\n current_index = value[i].get(\"index\", None)\n if current_type == last_type and current_index == last_index:\n # Merge all non-type fields\n for k, v in value[i].items():\n if k != \"type\":\n if k in new_value[-1]:\n new_value[-1][k] += v\n else:\n new_value[-1][k] = v\n else:\n new_value.append(value[i])\n last_type = current_type\n last_index = current_index\n result[key] = new_value\n\n return result\n" }, { "path": "cortex/orchestrator/__init__.py", "content": "\"\"\"Orchestrator module - Provides generator merging and agent coordination functionality.\"\"\"\n\nfrom .local_runner import LocalRunner\nfrom .orchestrator import Orchestrator\nfrom .remote_runner import RemoteRunner\nfrom .runner import Runner\nfrom .types import AgentEvent, AgentRequest, ClientToolCall, ClientToolCallType\n\n__all__ = [\n \"Orchestrator\",\n \"AgentEvent\",\n \"AgentRequest\",\n \"ClientToolCall\",\n \"ClientToolCallType\",\n \"Runner\",\n \"LocalRunner\",\n \"RemoteRunner\",\n]\n" }, { "path": "cortex/orchestrator/local_runner.py", "content": "\"\"\"Local Runner implementation\"\"\"\n\nimport asyncio\nfrom typing import AsyncGenerator, Optional\nfrom uuid import uuid4\n\nfrom cortex.model.definition import ChatMessage, Function\nfrom cortex.agents.agent_factory import AgentFactory\nfrom cortex.agents.base_agent import BaseAgent\nfrom cortex.agents.input.input import InputChannel\nfrom cortex.agents.types import (\n AgentConfig,\n AgentMessageType,\n AgentResponse,\n AgentRunningStatus,\n)\nfrom cortex.orchestrator.runner import Runner\nfrom cortex.orchestrator.types import (\n AgentEvent,\n AgentEventType,\n ClientToolCall,\n ClientToolCallType,\n)\nfrom cortex.tools.base import ToolSchema\nfrom cortex.tools.types import ToolParameters, ToolType\nfrom cortex.utils.generator_merger import GeneratorMerger\n\n\nclass LocalRunner(Runner):\n \"\"\"Local Agent runner.\"\"\"\n\n _continue_input: bool = False\n\n def __init__(\n self,\n agent_factory: AgentFactory,\n task_id: str,\n parent_task_id: str | None = None,\n root_task_id: str | None = None,\n tool_call_id: str | None = None,\n ):\n \"\"\"Initialize LocalRunner.\n\n Args:\n task_id: Task ID\n parent_task_id: Parent task ID\n root_task_id: Root task ID\n continue_input: Whether to continue input\n tool_call_id: Tool call ID\n \"\"\"\n super().__init__(task_id, parent_task_id, root_task_id)\n self._agent_factory = agent_factory\n self._tool_call_id = tool_call_id\n self._task_id = task_id\n self._merger = GeneratorMerger(\n on_generator_complete=self._on_generator_complete\n )\n self._agent_factory = agent_factory\n self._agent: BaseAgent | None = None\n self._agent_name: str | None = None\n\n self._result: AgentEvent | None = None\n\n async def init(\n self,\n agent_name: str,\n context_id: str | None = None,\n config: AgentConfig | None = None,\n ) -> None:\n \"\"\"Initialize Runner.\n\n Args:\n agent_name: Agent name\n context_id: Context ID\n config: AgentConfig\n \"\"\"\n if config is not None and not config.use_share_context:\n context_id = uuid4().hex\n\n self._agent_name = agent_name\n self._agent = await self._agent_factory.make_agent(\n agent_name, context_id, config\n )\n if self._agent.config.unfinished_mode:\n self._continue_input = True\n self._message_queue: asyncio.Queue[ChatMessage] = asyncio.Queue()\n self._messages: InputChannel[ChatMessage] = InputChannel(\n self._message_queue\n )\n else:\n self._messages: list[ChatMessage] = []\n\n self._agent.toolset().set_on_send(self._on_client_tool_send)\n\n async def send(self, event: AgentEvent) -> None:\n \"\"\"Send AgentEvent.\n\n Args:\n event: AgentEvent to send\n \"\"\"\n if event.type == AgentEventType.REQUEST:\n # Add messages from request to queue\n if event.request and event.request.messages:\n for message in event.request.messages:\n if self._continue_input:\n await self._message_queue.put(message)\n else:\n self._messages.append(message)\n if event.type == AgentEventType.CLIENT_TOOL_RESULT:\n result = event.client_tool_result\n if result is not None:\n self._agent.toolset().set_response(\n result.message.tool_call_id,\n result.message.content,\n result.error_msg,\n )\n\n async def run(self) -> AsyncGenerator[AgentEvent, None]:\n \"\"\"Run and return AgentEvent generator.\n\n Yields:\n AgentEvent: Agent event\n \"\"\"\n if self._task_id is None:\n raise ValueError(\"task_id not set, please call send() first to send REQUEST event\")\n\n async def agent_generator() -> AsyncGenerator[AgentEvent, None]:\n if self._agent is None:\n raise ValueError(\"Agent not initialized, please call init() first\")\n\n async for response in self._agent.run(self._messages):\n self._on_agent_finished(response)\n yield AgentEvent(\n agent_name=self._agent_name,\n task_id=self._task_id,\n parent_task_id=self._parent_task_id,\n root_task_id=self._root_task_id,\n type=AgentEventType.RESPONSE,\n response=response,\n )\n\n self._messages = []\n\n self._merger.add_async_generator(\n agent_generator, generator_id=f\"local_runner_{self._task_id}\"\n )\n\n async for event in self._merger.merge():\n yield event\n\n def _on_agent_finished(self, response: AgentResponse) -> None:\n \"\"\"Handle Agent finished event.\n\n Args:\n response: AgentResponse\n \"\"\"\n if response.status != AgentRunningStatus.FINISHED:\n return\n\n if response.message_type == AgentMessageType.STREAM:\n return\n\n if self._tool_call_id is None:\n self._result = AgentEvent(\n agent_name=self._agent_name,\n task_id=self._task_id,\n parent_task_id=self._parent_task_id,\n root_task_id=self._root_task_id,\n type=AgentEventType.RESPONSE,\n response=response,\n )\n return\n\n response.message.tool_call_id = self._tool_call_id\n self._result = AgentEvent(\n agent_name=self._agent_name,\n task_id=self._task_id,\n parent_task_id=self._parent_task_id,\n root_task_id=self._root_task_id,\n type=AgentEventType.CLIENT_TOOL_RESULT,\n client_tool_result=response,\n )\n\n async def _on_generator_complete(\n self, generator_id: str, generator_type: str, error: Optional[Exception]\n ) -> None:\n \"\"\"Handle generator complete event.\n\n Args:\n generator_id: Generator ID\n generator_type: Generator Type\n error: Error\n \"\"\"\n # Generator complete event handling (no special handling needed currently)\n\n async def _on_client_tool_send(\n self, tool_name: str, tool_schema: ToolSchema, tool_parameters: ToolParameters\n ) -> None:\n \"\"\"Handle client tool send event.\n\n Args:\n tool_name: Tool name\n tool_schema: Tool Schema\n tool_parameters: Tool Parameters (contains tool_call_id in kwargs)\n \"\"\"\n # Extract tool_call_id from tool_parameters.kwargs\n tool_call_id = tool_parameters.kwargs.get(\"tool_call_id\")\n if tool_call_id is None:\n tool_call_id = (\n f\"tool_call_{hash(tool_name)}_{hash(tool_parameters.parameters)}\"\n )\n\n async def tool_generator() -> AsyncGenerator[AgentResponse, None]:\n tool_type = ClientToolCallType.TOOL\n if tool_schema.tool_type == ToolType.AGENT:\n tool_type = ClientToolCallType.AGENT\n elif tool_schema.tool_type == ToolType.ASK_INPUT:\n tool_type = ClientToolCallType.ASK_INPUT\n\n # tool params event\n yield AgentEvent(\n task_id=self._task_id,\n parent_task_id=self._parent_task_id,\n root_task_id=self._root_task_id,\n agent_name=self._agent_name,\n type=AgentEventType.CLIENT_TOOL_CALL,\n client_tool_call=ClientToolCall(\n tool_call_id=tool_call_id,\n function=Function(\n arguments=tool_parameters.parameters,\n name=tool_schema.name,\n ),\n type=tool_type,\n extra=tool_parameters.kwargs,\n ),\n )\n\n self._merger.add_async_generator(tool_generator, generator_id=tool_call_id)\n\n def get_result(self) -> AgentEvent:\n \"\"\"Get AgentEvent.\n\n Returns:\n AgentEvent: AgentEvent\n \"\"\"\n return self._result\n" }, { "path": "cortex/orchestrator/orchestrator.py", "content": "\"\"\"Orchestrator for coordinating execution of multiple Agents.\"\"\"\n\nimport logging\nfrom enum import Enum\nfrom typing import AsyncGenerator, Optional\nfrom uuid import uuid4\n\nfrom cortex.agents.agent_factory import AgentFactory\nfrom cortex.agents.types import AgentConfig, AgentResponse, AgentRunningStatus\nfrom cortex.orchestrator.local_runner import LocalRunner\nfrom cortex.orchestrator.runner import Runner\nfrom cortex.orchestrator.types import (\n AgentEvent,\n AgentEventType,\n AgentRequest,\n ClientToolCallType,\n)\nfrom cortex.tools.agent_tool import AgentTool\nfrom cortex.utils.generator_merger import GeneratorMerger\n\nlogger = logging.getLogger(__name__)\n\n\nclass OrchMode(str, Enum):\n \"\"\"Orchestrator mode.\"\"\"\n\n MULTI = \"multi\"\n SINGLE = \"single\"\n\n\nclass Orchestrator:\n \"\"\"Orchestrator for coordinating execution of multiple Agents.\"\"\"\n\n def __init__(self, agent_factory: AgentFactory):\n \"\"\"Initialize Orchestrator.\n\n Args:\n agent_factory: AgentFactory instance for creating Agents\n \"\"\"\n self._agent_factory = agent_factory\n # Manage runner by task_id\n self._runners: dict[str, Runner] = {}\n # Manage generator_merger by root_task_id\n self._mergers: dict[str, GeneratorMerger] = {}\n # Record task_id to root_task_id mapping\n self._task_to_root: dict[str, str] = {}\n # Record task_id to parent_task_id mapping\n self._task_to_parent: dict[str, str] = {}\n # Record runner's task_id (for cleanup)\n self._runner_task_ids: dict[str, str] = {} # runner_id -> task_id\n\n def list_agents(self) -> list[AgentConfig]:\n \"\"\"List all Agent configurations.\"\"\"\n return self._agent_factory.list_agents()\n\n async def run(\n self,\n agent_name: str,\n event: AgentEvent,\n agent_config: AgentConfig | None = None,\n mode: OrchMode = OrchMode.MULTI,\n context_id: str | None = None,\n ) -> AsyncGenerator[AgentEvent, None]:\n \"\"\"Run Agent and return response stream.\n\n Args:\n agent_name: Agent name\n messages: Message list\n agent_config: Agent configuration\n\n Yields:\n AgentEvent: Agent event\n \"\"\"\n task_id = event.task_id or f\"root_{uuid4().hex}\"\n root_task_id = event.root_task_id or task_id\n parent_task_id = event.parent_task_id or None\n event.task_id = task_id\n event.root_task_id = root_task_id\n event.parent_task_id = parent_task_id\n if context_id is None:\n context_id = task_id\n\n # Create root runner\n runner = await self._create_runner(\n task_id=task_id,\n parent_task_id=parent_task_id,\n root_task_id=root_task_id,\n )\n\n # Initialize runner\n await runner.init(agent_name, context_id, agent_config)\n\n async def on_generator_complete_with_mode(\n generator_id: str, _generator_type: str, error: Optional[Exception]\n ) -> None:\n await self._on_generator_complete(\n generator_id, _generator_type, error, mode\n )\n\n # Create generator_merger to merge events\n merger = GeneratorMerger(on_generator_complete=on_generator_complete_with_mode)\n self._mergers[root_task_id] = merger\n\n need_run_root_runner = True\n while need_run_root_runner:\n need_run_root_runner = False\n await self.run_root_runner(runner, merger, root_task_id, event)\n event = None\n\n # Get events from merger and yield AgentResponse\n async for event in merger.merge():\n if not isinstance(event, AgentEvent):\n continue\n\n # Handle client tool call\n if (\n mode == OrchMode.MULTI\n and event.type == AgentEventType.CLIENT_TOOL_CALL\n ):\n # Check if it's AGENT type, only AGENT type needs special handling\n client_tool_call = event.client_tool_call\n if (\n client_tool_call\n and client_tool_call.type == ClientToolCallType.AGENT\n ):\n await self._handle_client_tool_call(\n context_id, event, root_task_id\n )\n need_run_root_runner = True\n else:\n # Non-AGENT type CLIENT_TOOL_CALL (e.g., ASK_INPUT, TOOL) yield directly\n yield event\n else:\n yield event\n\n self._cleanup_runner(task_id)\n\n async def run_root_runner(\n self,\n runner: Runner,\n merger: GeneratorMerger,\n root_task_id: str,\n event: AgentEvent,\n ):\n # Add runner's run() to merger\n async def runner_generator() -> AsyncGenerator[AgentEvent, None]:\n async for event in runner.run():\n yield event\n\n merger.add_async_generator(runner_generator, generator_id=root_task_id)\n if event is None:\n return\n if event.type == AgentEventType.REQUEST and len(event.request.messages) > 0:\n # Send REQUEST event\n await runner.send(event)\n elif event.type == AgentEventType.CLIENT_TOOL_RESULT:\n await runner.send(event)\n\n async def send_event(self, event: AgentEvent) -> None:\n \"\"\"Receive external input AgentEvent, find runner by task_id and send.\n\n Args:\n event: AgentEvent\n \"\"\"\n task_id = event.task_id\n runner = self._runners.get(task_id)\n if runner is None:\n logger.error(\"Cannot find runner for task_id=%s\", task_id)\n raise ValueError(\"Cannot find runner for task_id=%s\", task_id)\n await runner.send(event)\n\n async def _create_runner(\n self,\n task_id: str,\n parent_task_id: str | None = None,\n root_task_id: str | None = None,\n tool_call_id: str | None = None,\n ) -> Runner:\n \"\"\"Create Runner instance.\n\n Args:\n task_id: Task ID\n parent_task_id: Parent task ID\n root_task_id: Root task ID\n tool_call_id: Tool call ID\n Returns:\n Runner instance\n \"\"\"\n runner = LocalRunner(\n agent_factory=self._agent_factory,\n task_id=task_id,\n parent_task_id=parent_task_id,\n root_task_id=root_task_id,\n tool_call_id=tool_call_id,\n )\n self._runners[task_id] = runner\n if root_task_id:\n self._task_to_root[task_id] = root_task_id\n if parent_task_id:\n self._task_to_parent[task_id] = parent_task_id\n return runner\n\n async def _handle_client_tool_call(\n self, context_id: str, event: AgentEvent, root_task_id: str\n ):\n \"\"\"Handle client tool call event.\n\n Args:\n event: AgentEvent, type is CLIENT_TOOL_CALL\n context_id: Context ID\n root_task_id: Root task ID\n\n Returns:\n None\n \"\"\"\n client_tool_call = event.client_tool_call\n if client_tool_call is None:\n return\n\n # Determine if it's an agent based on client_tool_call.type\n is_agent = client_tool_call.type == ClientToolCallType.AGENT\n if not is_agent:\n # Not an agent, return directly\n return\n\n agent_name = client_tool_call.function.name\n # Get agent config\n try:\n agent_config = self._agent_factory.get_default_agent_config(agent_name)\n except ValueError:\n agent_config = None\n\n if agent_config is None:\n return\n\n child_task_id = f\"child_{uuid4().hex}\"\n # Create child runner\n child_runner = await self._create_runner(\n task_id=child_task_id,\n parent_task_id=event.task_id,\n root_task_id=root_task_id,\n tool_call_id=client_tool_call.tool_call_id,\n )\n\n # Initialize child runner\n await child_runner.init(agent_name, context_id, agent_config)\n\n # Get merger for root_task_id\n merger = self._mergers.get(root_task_id)\n if merger is None:\n raise ValueError(f\"Cannot find merger for root_task_id={root_task_id}\")\n\n # Prepare messages (convert tool call to user message)\n tool_call_messages = AgentTool.parse_messages(\n client_tool_call.function.arguments\n )\n if tool_call_messages is None:\n tool_call_messages = []\n\n # Add child runner's run() to merger\n async def child_runner_generator() -> AsyncGenerator[AgentEvent, None]:\n # Send REQUEST event to child runner\n await child_runner.send(\n AgentEvent(\n task_id=child_task_id,\n parent_task_id=event.task_id,\n root_task_id=root_task_id,\n type=AgentEventType.REQUEST,\n request=AgentRequest(\n agent_name=agent_name,\n config=agent_config,\n messages=tool_call_messages,\n ),\n )\n )\n\n # Run child runner and yield events\n async for child_event in child_runner.run():\n yield child_event\n\n merger.add_async_generator(child_runner_generator, generator_id=child_task_id)\n self._runner_task_ids[child_task_id] = child_task_id\n\n async def _on_generator_complete(\n self,\n generator_id: str,\n _generator_type: str,\n error: Optional[Exception],\n mode: OrchMode,\n ) -> None:\n \"\"\"Handle generator complete event (callback).\n\n Args:\n generator_id: Generator ID (usually task_id)\n _generator_type: Generator Type (unused, kept for interface compatibility)\n error: Error message (if any)\n \"\"\"\n # Find corresponding runner\n task_id = generator_id\n runner = self._runners.get(task_id)\n if runner is None:\n return\n\n if mode == OrchMode.SINGLE:\n self._cleanup_runner(task_id)\n return\n\n # Get parent_task_id\n parent_task_id = self._task_to_parent.get(task_id)\n if parent_task_id is None:\n # No parent, means this is root runner, no need to send result\n return\n\n # Has parent, need to send result to parent runner\n parent_runner = self._runners.get(parent_task_id)\n if parent_runner is None:\n # Parent runner doesn't exist, cleanup current runner\n self._cleanup_runner(task_id)\n return\n\n # Create result response\n if error:\n result_response = AgentResponse(\n status=AgentRunningStatus.ERROR,\n error_msg=str(error),\n )\n else:\n # Completed normally\n runner_result = runner.get_result()\n if runner_result is None:\n result_response = AgentResponse(\n status=AgentRunningStatus.ERROR,\n error_msg=f\"runner {task_id} did not return result\",\n )\n else:\n result_response = runner_result.client_tool_result\n\n # Send CLIENT_TOOL_RESULT event to parent runner\n await parent_runner.send(\n AgentEvent(\n task_id=parent_task_id,\n parent_task_id=parent_runner.get_parent_task_id(),\n root_task_id=parent_runner.get_root_task_id(),\n type=AgentEventType.CLIENT_TOOL_RESULT,\n client_tool_result=result_response,\n )\n )\n\n # Cleanup completed runner\n self._cleanup_runner(task_id)\n\n def _cleanup_runner(self, task_id: str) -> None:\n \"\"\"Cleanup completed runner.\n\n Args:\n task_id: Task ID\n \"\"\"\n # Remove from runners\n if task_id in self._runners:\n del self._runners[task_id]\n\n # Remove from mappings\n if task_id in self._task_to_root:\n del self._task_to_root[task_id]\n\n if task_id in self._task_to_parent:\n del self._task_to_parent[task_id]\n\n if task_id in self._runner_task_ids:\n del self._runner_task_ids[task_id]\n\n # If it's root_task_id, cleanup merger\n if task_id in self._mergers:\n del self._mergers[task_id]\n" }, { "path": "cortex/orchestrator/remote_runner.py", "content": "from typing import AsyncGenerator\n\nfrom cortex.agents.types import AgentConfig, AgentRunningStatus\nfrom cortex.orchestrator.runner import Runner\nfrom cortex.orchestrator.types import AgentEvent, AgentEventType\nfrom cortex.server.channel.channel import Channel\n\n\nclass RemoteRunner(Runner):\n \"\"\"Remote Agent runner that communicates with remote services through Channel.\"\"\"\n\n def __init__(self, channel: Channel):\n \"\"\"\n Initialize remote runner.\n\n Args:\n channel: Channel instance for communicating with remote services\n \"\"\"\n self.channel = channel\n self._task_id: str | None = None\n self._running: bool = False\n\n async def init(\n self,\n agent_name: str,\n context_id: str | None = None,\n config: AgentConfig | None = None,\n ) -> None:\n \"\"\"\n Initialize Runner.\n\n Args:\n agent_name: Agent name\n context_id: Context ID\n config: AgentConfig\n \"\"\"\n # Remote runner doesn't need local configuration, configuration is on the remote server side\n # If needed, you can save config or perform other initialization here\n\n async def send(self, event: AgentEvent) -> None:\n \"\"\"\n Send AgentEvent to remote service.\n\n Args:\n event: The AgentEvent to send\n \"\"\"\n if event.type == AgentEventType.REQUEST:\n # Save task_id\n if self._task_id is None:\n self._task_id = event.task_id\n\n # Send event to remote service\n await self.channel.send(event.model_dump())\n\n async def run(self) -> AsyncGenerator[AgentEvent, None]:\n \"\"\"\n Run and return an AgentEvent generator, receiving events from remote service.\n\n Yields:\n AgentEvent: Agent events\n \"\"\"\n if self._task_id is None:\n raise ValueError(\"task_id is not set, please call send() to send REQUEST event first\")\n\n self._running = True\n\n try:\n while self._running:\n # Receive events from remote service\n data = await self.channel.receive()\n event = AgentEvent.model_validate(data)\n\n # Only yield events related to current task_id\n if event.task_id == self._task_id:\n yield event\n\n # Stop receiving if error or completion signal is received\n if event.type == AgentEventType.ERROR:\n break\n elif event.type == AgentEventType.RESPONSE and event.response:\n if event.response.status in (\n AgentRunningStatus.FINISHED,\n AgentRunningStatus.STOPPED,\n AgentRunningStatus.ERROR,\n ):\n break\n finally:\n self._running = False\n" }, { "path": "cortex/orchestrator/runner.py", "content": "\"\"\"Runner interface definition.\"\"\"\n\nfrom abc import ABC, abstractmethod\nfrom typing import AsyncGenerator\n\nfrom cortex.agents.types import AgentConfig\nfrom cortex.orchestrator.types import AgentEvent\n\n\nclass Runner(ABC):\n \"\"\"Runner interface for handling AgentEvent sending and running.\"\"\"\n\n def __init__(\n self,\n task_id: str,\n parent_task_id: str | None = None,\n root_task_id: str | None = None,\n ):\n \"\"\"\n Initialize Runner.\n\n Args:\n task_id: Task ID\n parent_task_id: Parent task ID\n root_task_id: Root task ID\n \"\"\"\n self._task_id: str = task_id\n self._parent_task_id: str | None = parent_task_id\n self._root_task_id: str | None = root_task_id\n\n def get_parent_task_id(self) -> str | None:\n \"\"\"\n Get parent task ID.\n\n Returns:\n str | None: Parent task ID\n \"\"\"\n return self._parent_task_id\n\n def get_root_task_id(self) -> str | None:\n \"\"\"\n Get root task ID.\n\n Returns:\n str | None: Root task ID\n \"\"\"\n return self._root_task_id\n\n @abstractmethod\n async def init(\n self,\n agent_name: str,\n context_id: str | None = None,\n config: AgentConfig | None = None,\n ) -> None:\n \"\"\"\n Initialize Runner.\n\n Args:\n agent_name: Agent name\n context_id: Context ID\n config: AgentConfig\n \"\"\"\n raise NotImplementedError\n\n @abstractmethod\n async def send(self, event: AgentEvent) -> None:\n \"\"\"\n Send AgentEvent.\n\n Args:\n event: The AgentEvent to send\n \"\"\"\n raise NotImplementedError\n\n @abstractmethod\n async def run(self) -> AsyncGenerator[AgentEvent, None]:\n \"\"\"\n Run and return an AgentEvent generator.\n\n Yields:\n AgentEvent: Agent events\n \"\"\"\n raise NotImplementedError\n\n @abstractmethod\n def get_result(self) -> AgentEvent:\n \"\"\"\n Get AgentEvent.\n\n Returns:\n AgentEvent: AgentEvent\n \"\"\"\n raise NotImplementedError\n" }, { "path": "cortex/orchestrator/types.py", "content": "from enum import Enum\nfrom uuid import uuid4\n\nfrom cortex.model.definition import ChatMessage, Function\nfrom pydantic import BaseModel, Field\n\nfrom cortex.agents.types import AgentConfig, AgentResponse\n\n\nclass AgentEventType(str, Enum):\n \"\"\"Agent event type enum.\"\"\"\n\n REQUEST = \"request\"\n RESPONSE = \"response\"\n ERROR = \"error\"\n SIGNAL = \"signal\"\n CLIENT_TOOL_CALL = \"client_tool_call\"\n CLIENT_TOOL_RESULT = \"client_tool_result\"\n\n\nclass AgentRequest(BaseModel):\n \"\"\"Agent request model.\"\"\"\n\n agent_name: str\n config: AgentConfig | None = None\n messages: list[ChatMessage] | None = None\n\n\nclass ClientToolCallType(str, Enum):\n \"\"\"Client tool call type enum.\"\"\"\n\n AGENT = \"agent\"\n TOOL = \"tool\"\n ASK_INPUT = \"ask_input\"\n\n\nclass ClientToolCall(BaseModel):\n \"\"\"Client tool call model.\"\"\"\n\n tool_call_id: str\n function: Function\n type: ClientToolCallType\n extra: dict | None = None\n\n\nclass AgentEvent(BaseModel):\n \"\"\"Agent event model.\"\"\"\n\n event_id: str = Field(default_factory=lambda: f\"{uuid4().hex}\")\n task_id: str | None = None\n parent_task_id: str | None = None\n root_task_id: str | None = None\n type: AgentEventType\n metadata: dict | None = None\n agent_name: str | None = None\n\n # Input\n request: AgentRequest | None = None\n\n # Output\n response: AgentResponse | None = None\n error: str | None = None\n\n # client tool call\n client_tool_call: ClientToolCall | None = None\n client_tool_result: AgentResponse | None = None\n" }, { "path": "cortex/runtime_config.py", "content": "from __future__ import annotations\n\nimport os\nfrom pathlib import Path\nfrom typing import Any\n\ntry:\n import yaml\nexcept Exception: # noqa: BLE001\n yaml = None\n\n_RUNTIME_CONFIG_CACHE: dict[str, Any] | None = None\n\n\ndef _repo_root() -> Path:\n # `cortex/` lives under the project root in this repo.\n return Path(__file__).resolve().parent.parent\n\n\ndef _pick_first_present(cfg: dict[str, Any], keys: tuple[str, ...]) -> Any:\n for key in keys:\n if key in cfg:\n return cfg.get(key)\n return None\n\n\ndef _as_int(value: Any) -> int | None:\n if value is None:\n return None\n if isinstance(value, bool):\n return None\n if isinstance(value, (int, float)):\n return int(value)\n if isinstance(value, str):\n text = value.strip()\n if not text:\n return None\n try:\n return int(text)\n except Exception: # noqa: BLE001\n return None\n return None\n\n\ndef load_runtime_config() -> dict[str, Any]:\n \"\"\"Load runtime config from YAML, cached for process lifetime.\n\n Resolution order:\n 1) `$STEP_DEEPRESEARCH_CONFIG` (if set)\n 2) `/config.yaml` (if exists)\n \"\"\"\n global _RUNTIME_CONFIG_CACHE # noqa: PLW0603\n if _RUNTIME_CONFIG_CACHE is not None:\n return _RUNTIME_CONFIG_CACHE\n\n if yaml is None:\n _RUNTIME_CONFIG_CACHE = {}\n return _RUNTIME_CONFIG_CACHE\n\n env_path = (os.getenv(\"STEP_DEEPRESEARCH_CONFIG\") or \"\").strip()\n candidates: list[Path] = []\n if env_path:\n candidates.append(Path(env_path))\n candidates.append(_repo_root() / \"config.yaml\")\n\n for path in candidates:\n try:\n if not path.exists() or not path.is_file():\n continue\n with path.open(\"r\", encoding=\"utf-8\") as f:\n loaded = yaml.safe_load(f) or {}\n if isinstance(loaded, dict):\n _RUNTIME_CONFIG_CACHE = loaded\n return _RUNTIME_CONFIG_CACHE\n except Exception: # noqa: BLE001\n continue\n\n _RUNTIME_CONFIG_CACHE = {}\n return _RUNTIME_CONFIG_CACHE\n\n\ndef get_context_limit_overrides() -> tuple[int | None, int | None]:\n \"\"\"Return (upper, lower) context limit overrides from runtime config.\"\"\"\n cfg = load_runtime_config()\n upper = _pick_first_present(\n cfg, (\"context_upper_limit\", \"final_answer_context_upper_limit\")\n )\n lower = _pick_first_present(\n cfg, (\"context_lower_limit\", \"final_answer_context_lower_limit\")\n )\n return _as_int(upper), _as_int(lower)\n\n" }, { "path": "cortex/server/channel/channel.py", "content": "from abc import ABC, abstractmethod\n\n\nclass Channel(ABC):\n def __init__(self, context_id: str) -> None:\n self.context_id: str = context_id\n\n @abstractmethod\n async def send(self, event: dict[str, object]):\n pass\n\n @abstractmethod\n async def receive(self) -> dict[str, object]:\n pass\n\n @abstractmethod\n async def heartbeat(self):\n pass\n\n @abstractmethod\n async def close(self):\n pass\n" }, { "path": "cortex/server/channel/error.py", "content": "class ChannelClosedError(Exception):\n pass\n" }, { "path": "cortex/server/channel/memory_channel.py", "content": "import asyncio\nfrom typing import Dict\n\nfrom loguru import logger\n\nfrom cortex.server.channel.channel import Channel\nfrom cortex.server.channel.error import ChannelClosedError\n\n\nclass MemoryChannel(Channel):\n def __init__(self, context_id: str) -> None:\n super().__init__(context_id)\n self.send_queue = asyncio.Queue()\n self.receive_queue = asyncio.Queue()\n self.is_closed = False\n logger.debug(f\"MemoryChannel created with context_id: {context_id}\")\n\n async def send(self, event: Dict[str, object]):\n if self.is_closed:\n raise ChannelClosedError(\"Channel is closed\")\n await self.send_queue.put(event)\n logger.debug(f\"MemoryChannel {self.context_id}: sent event\")\n\n async def receive(self) -> Dict[str, object]:\n if self.is_closed:\n raise ChannelClosedError(\"Channel is closed\")\n\n data = await self.receive_queue.get()\n logger.debug(f\"MemoryChannel {self.context_id}: received event\")\n return data\n\n async def heartbeat(self):\n pass\n\n async def close(self):\n logger.debug(f\"MemoryChannel {self.context_id}: closing\")\n self.is_closed = True\n" }, { "path": "cortex/server/channel/ws_channel.py", "content": "\"\"\"WebSocket channel implementation for agent communication.\"\"\"\n\nimport asyncio\nimport logging\nimport time\nfrom typing import cast\n\nfrom fastapi import WebSocket, WebSocketDisconnect\n\nfrom cortex.server.channel.channel import Channel\nfrom cortex.server.channel.error import ChannelClosedError\n\nlogger = logging.getLogger(__name__)\n\n\nclass WebSocketChannel(Channel):\n _closed: bool = False\n last_heartbeat_time: float = time.time()\n\n def __init__(self, ws: WebSocket) -> None:\n self.ws: WebSocket = ws\n\n async def send(self, event: dict[str, object]) -> None:\n if self._closed:\n raise ChannelClosedError\n await self.ws.send_json(event)\n\n async def receive(self) -> dict[str, object]:\n if self._closed:\n raise ChannelClosedError\n while True:\n try:\n data = cast(dict[str, object], await self.ws.receive_json())\n self.last_heartbeat_time = time.time()\n if data.get(\"type\") == \"ping\":\n await self.ws.send_json({\"type\": \"pong\"})\n continue\n\n if data.get(\"type\") == \"pong\":\n continue\n return data\n except WebSocketDisconnect:\n self._closed = True\n raise ChannelClosedError\n\n async def heartbeat(self) -> None:\n await asyncio.sleep(10)\n while not self._closed:\n await self.ws.send_json({\"type\": \"ping\"})\n await asyncio.sleep(10)\n # TODO Check if last_heartbeat_time has timed out\n\n async def close(self) -> None:\n logger.debug(\"WebSocketChannel closing\")\n if not self._closed:\n self._closed = True\n await self.ws.close()\n" }, { "path": "cortex/server/http_server.py", "content": "import asyncio\nimport logging\nimport uuid\n\nimport uvicorn\nfrom agentkit.trace import SpanContext, Tracer\nfrom fastapi import FastAPI, Request\nfrom fastapi.middleware.cors import CORSMiddleware\nfrom fastapi.responses import StreamingResponse\nfrom fastapi.websockets import WebSocket\nfrom pydantic import ValidationError\nfrom starlette.middleware.base import BaseHTTPMiddleware\n\nfrom cortex.agents.types import AgentConfig\nfrom cortex.orchestrator.orchestrator import Orchestrator, OrchMode\nfrom cortex.orchestrator.types import AgentEvent, AgentEventType\nfrom cortex.server.channel.ws_channel import WebSocketChannel\nfrom cortex.server.log.log import setup_logging\nfrom cortex.server.log.trace import set_trace_id\n\nsetup_logging()\nlogger = logging.getLogger(__name__)\n\n\ndef extract_and_set_trace_id(headers: dict) -> str:\n \"\"\"\n Extract x-step-trace from headers, auto-generate UUID if not present,\n and set it to ContextVar\n\n Args:\n headers: HTTP headers dict\n\n Returns:\n trace_id string\n \"\"\"\n trace_id = headers.get(\"Step-Trace-ID\")\n if not trace_id:\n trace_id = str(uuid.uuid4())\n\n set_trace_id(trace_id)\n return trace_id\n\n\nclass TraceMiddleware(BaseHTTPMiddleware):\n \"\"\"Extract x-step-trace from HTTP header and inject into log context, auto-generate if not present\"\"\"\n\n async def dispatch(self, request: Request, call_next):\n # Use unified function to extract and set trace_id\n extract_and_set_trace_id(request.headers)\n response = await call_next(request)\n return response\n\n\nclass HttpServer:\n orch: Orchestrator\n tracer: Tracer\n\n def __init__(self, orch: Orchestrator, tracer: Tracer):\n self.orch = orch\n self.tracer = tracer\n\n async def start(self, host: str = \"0.0.0.0\", port: int = 8001) -> None:\n app = self._build_app()\n logger.info(\"Starting HTTP server, listening on 0.0.0.0:8001\")\n # Configure uvicorn to use Python logging\n server_config = uvicorn.Config(\n app,\n host=host,\n port=port,\n log_config=None, # Disable uvicorn default log config\n access_log=True, # Enable access log\n )\n await uvicorn.Server(server_config).serve()\n\n def _build_app(self) -> FastAPI:\n app = FastAPI(title=\"Agent Server\")\n\n # Add CORS middleware, support cross-origin requests from any domain\n app.add_middleware(\n CORSMiddleware,\n allow_origins=[\"*\"], # Allow all domains\n allow_credentials=True,\n allow_methods=[\"*\"], # Allow all HTTP methods\n allow_headers=[\"*\"], # Allow all request headers\n )\n\n # Add Trace middleware to extract trace_id from x-step-trace header and inject into log\n app.add_middleware(TraceMiddleware)\n\n @app.get(\"/agents\")\n async def list_agents() -> list[AgentConfig]:\n return self.orch.list_agents()\n\n @app.websocket(\"/multi/ws/{agent_name}/{context_id}\")\n async def multi_call(websocket: WebSocket, agent_name: str, context_id: str):\n await websocket_handler(websocket, agent_name, context_id, OrchMode.MULTI)\n\n @app.websocket(\"/single/ws/{agent_name}/{context_id}\")\n async def single_call(websocket: WebSocket, agent_name: str, context_id: str):\n await websocket_handler(websocket, agent_name, context_id, OrchMode.SINGLE)\n\n @app.post(\"/multi/sse/{agent_name}/{context_id}\")\n async def multi_call_sse(agent_name: str, context_id: str, request: AgentEvent):\n return StreamingResponse(\n sse_handler(\n agent_name, context_id, OrchMode.MULTI, request\n ), # Async generator (core)\n media_type=\"text/event-stream\", # SSE response type\n headers={\n \"Cache-Control\": \"no-cache\",\n \"Connection\": \"keep-alive\",\n # Cross-origin support (needed if frontend and backend have different domains)\n \"Access-Control-Allow-Origin\": \"*\",\n \"Access-Control-Allow-Credentials\": \"true\",\n },\n )\n\n @app.post(\"/single/sse/{agent_name}/{context_id}\")\n async def single_call_sse(\n agent_name: str, context_id: str, request: AgentEvent\n ):\n return StreamingResponse(\n sse_handler(\n agent_name, context_id, OrchMode.SINGLE, request\n ), # Async generator (core)\n media_type=\"text/event-stream\", # SSE response type\n headers={\n \"Cache-Control\": \"no-cache\",\n \"Connection\": \"keep-alive\",\n # Cross-origin support (needed if frontend and backend have different domains)\n \"Access-Control-Allow-Origin\": \"*\",\n \"Access-Control-Allow-Credentials\": \"true\",\n },\n )\n\n async def websocket_handler(\n websocket: WebSocket, agent_name: str, context_id: str, mode: OrchMode\n ):\n # Use unified function to extract and set trace_id from headers\n extract_and_set_trace_id(websocket.headers)\n\n await websocket.accept()\n\n channel = WebSocketChannel(websocket)\n task_id = str(uuid.uuid4())\n request = AgentEvent.model_validate(await channel.receive())\n if request.task_id is None:\n request.task_id = task_id\n if request.type != AgentEventType.REQUEST:\n raise ValueError(\"first event type must be request\")\n\n async def send_to_agent_loop():\n while True:\n try:\n data = await channel.receive()\n event = AgentEvent.model_validate(data)\n await self.orch.send_event(event)\n except ValidationError as e:\n logger.error(f\"Pydantic validation error: {e}\")\n continue\n\n async def send_to_client_loop():\n ctx = SpanContext(tracer=self.tracer, app_name=agent_name)\n with ctx.span(name=f\"orchestrator_{agent_name}_{context_id}\"):\n async for event in self.orch.run(\n agent_name, request, request.request.config, mode, context_id\n ):\n logger.warning(f'event: {event.model_dump()}')\n await channel.send(event.model_dump())\n\n agent_task = asyncio.create_task(\n send_to_client_loop(), name=\"send_to_client_loop\"\n )\n recv_task = asyncio.create_task(\n send_to_agent_loop(), name=\"send_to_agent_loop\"\n )\n\n done, pending = await asyncio.wait(\n [recv_task, agent_task], return_when=asyncio.FIRST_COMPLETED\n )\n for task in pending:\n logger.info(f\"Cancelling pending task: {task.get_name()}\")\n task.cancel()\n\n for task in done:\n try:\n task.result()\n logger.info(f\"Completed task: {task.get_name()}\")\n except Exception as e:\n logger.error(f\"Task error in {task.get_name()}: {e}\", exc_info=True)\n\n await channel.close()\n logger.info(\"WebSocket connection closed\")\n\n async def sse_handler(\n agent_name: str, context_id: str, mode: OrchMode, request: AgentEvent\n ):\n ctx = SpanContext(tracer=self.tracer, app_name=agent_name)\n with ctx.span(name=f\"http_sse_handler_{agent_name}_{context_id}\"):\n config = None\n if request.request is not None:\n config = request.request.config\n async for event in self.orch.run(\n agent_name, request, config, mode, context_id\n ):\n yield f\"data: {event.model_dump_json()}\\n\\n\"\n\n @app.get(\"/health\")\n async def health_check() -> dict[str, str]:\n return {\"status\": \"ok\"}\n\n return app\n" }, { "path": "cortex/server/log/log.py", "content": "# context.py\n\nimport logging\nimport os\nfrom datetime import datetime\nfrom pathlib import Path\n\nfrom pythonjsonlogger import jsonlogger\n\nfrom cortex.server.log.trace import TraceIdFilter\n\n\ndef setup_logging(log_dir: str = \"./logs\", log_level: int = logging.WARNING):\n \"\"\"\n Configure logging system to output to both console and file.\n\n Args:\n log_dir: Log file directory, defaults to ./logs\n log_level: Log level, defaults to INFO\n \"\"\"\n # Use rename_fields parameter to alias field names\n formatter = jsonlogger.JsonFormatter(\n \"%(asctime)s %(levelname)s %(trace_id)s %(name)s %(message)s\",\n rename_fields={\n \"asctime\": \"time\", # Timestamp\n \"levelname\": \"level\", # Log level\n \"trace_id\": \"traceid\", # Trace ID (camelCase)\n \"name\": \"name\", # Logger name\n \"message\": \"msg\", # Log message\n },\n )\n\n # Set StreamHandler encoding to UTF-8 for proper Unicode handling\n import sys\n\n stream_handler = logging.StreamHandler(sys.stdout)\n stream_handler.setFormatter(formatter)\n stream_handler.addFilter(TraceIdFilter())\n\n handlers = [stream_handler]\n\n # Add file Handler\n log_path = Path(log_dir)\n log_path.mkdir(parents=True, exist_ok=True)\n log_file = log_path / f\"server_{datetime.now().strftime('%Y%m%d')}.log\"\n\n file_handler = logging.FileHandler(log_file, encoding=\"utf-8\")\n file_handler.setFormatter(formatter)\n file_handler.addFilter(TraceIdFilter())\n handlers.append(file_handler)\n\n # Ensure basicConfig uses UTF-8 encoding\n logging.basicConfig(\n level=log_level,\n handlers=handlers,\n encoding=\"utf-8\", # Add UTF-8 encoding support\n )\n" }, { "path": "cortex/server/log/trace.py", "content": "import logging\nfrom contextvars import ContextVar\n\ntrace_id_var: ContextVar[str | None] = ContextVar(\"trace_id\", default=None)\n\n\ndef set_trace_id(trace_id: str):\n trace_id_var.set(trace_id)\n\n\ndef get_trace_id() -> str | None:\n return trace_id_var.get()\n\n\nclass TraceIdFilter(logging.Filter):\n \"\"\"Automatically inject trace_id from contextvars into LogRecord.\"\"\"\n\n def filter(self, record: logging.LogRecord) -> bool:\n record.trace_id = get_trace_id() or \"-\"\n return True\n" }, { "path": "cortex/tools/__init__.py", "content": "\"\"\"Tool system for Agent components.\"\"\"\n\nfrom .agent_tool import AgentTool\nfrom .base import Tool, ToolSchema\nfrom .channel import Channel\nfrom .client_tool import ClientTool\nfrom .function_tool import FunctionTool\nfrom .mcp_tool import MCPTool\nfrom .session_tool import SessionTool\nfrom .toolset import ToolSet\nfrom .types import ExecutionType, ToolConfig, ToolType\n\n__all__ = [\n \"Tool\",\n \"ToolSchema\",\n \"ExecutionType\",\n \"ToolType\",\n \"ToolConfig\",\n \"MCPTool\",\n \"FunctionTool\",\n \"SessionTool\",\n \"ClientTool\",\n \"AgentTool\",\n \"ToolSet\",\n \"Channel\",\n]\n" }, { "path": "cortex/tools/agent_tool.py", "content": "\"\"\"Agent Tool implementation - calls agents through Orchestrator.\"\"\"\n\nimport json\nimport logging\nfrom typing import Any, Optional\n\nfrom cortex.model.definition import ChatMessage\n\nfrom .base import Tool, ToolSchema\nfrom .channel import Channel\nfrom .types import ToolParameters, ToolType\n\nlogger = logging.getLogger(__name__)\n\n\nclass AgentTool(Tool):\n \"\"\"Agent tool implementation for calling other Agents, communicates with Orchestrator through Channel.\"\"\"\n\n share_context: bool = False\n\n def __init__(\n self,\n name: str,\n description: str = \"\",\n channel: Optional[Channel] = None,\n timeout: Optional[float] = None,\n share_context: bool = False,\n **kwargs,\n ):\n \"\"\"\n Initialize Agent tool.\n\n Args:\n name: Tool name\n description: Tool description\n channel: Channel instance (for async communication)\n timeout: Default timeout in seconds\n share_context: Whether to share context\n **kwargs: Additional parameters\n \"\"\"\n super().__init__(\n name=name, description=description, tool_type=ToolType.AGENT, **kwargs\n )\n if channel is None:\n raise ValueError(\"AgentTool requires a Channel instance\")\n self.channel = channel\n self.timeout = timeout or 30.0\n logger.debug(\n \"AgentTool initialized: name=%s, description=%s, timeout=%s\",\n name,\n description,\n self.timeout,\n )\n\n def _define_schema(self) -> ToolSchema:\n \"\"\"Define Agent tool schema.\"\"\"\n schema = ToolSchema(\n name=self.name,\n description=self.description,\n parameters={\n \"type\": \"object\",\n \"properties\": {\n \"content\": {\n \"type\": \"string\",\n \"description\": \"Specific instruction content to send to the Agent (optional, ignored if messages is provided)\",\n },\n \"messages\": {\n \"type\": \"array\",\n \"description\": \"List of messages to send to the Agent (optional, ChatMessage format with role and content fields)\",\n \"items\": {\n \"type\": \"object\",\n \"properties\": {\n \"role\": {\n \"type\": \"string\",\n \"description\": \"Message role, e.g., 'user', 'assistant', 'system'\",\n \"enum\": [\"user\", \"assistant\", \"system\", \"tool\"],\n },\n \"content\": {\n \"type\": \"string\",\n \"description\": \"Message content (string format)\",\n },\n },\n \"required\": [\"role\", \"content\"],\n },\n },\n \"timeout\": {\n \"type\": \"number\",\n \"description\": \"Timeout in seconds (optional)\",\n },\n },\n \"required\": [],\n },\n return_type=\"agent_response\",\n tool_type=self.tool_type,\n )\n logger.debug(\n \"AgentTool define schema: name=%s, parameters=%s\", self.name, schema.parameters\n )\n return schema\n\n async def _call(self, parameters: str, **kwargs) -> Any:\n \"\"\"\n Call the Agent tool.\n\n Send request to Orchestrator through Channel, which creates a new Agent and executes the task.\n\n Args:\n parameters: Tool parameters (JSON string format), including:\n - content: Specific instruction content (optional, ignored if messages is provided)\n - messages: List of messages (optional, ChatMessage format with role and content fields)\n - timeout: Timeout in seconds (optional)\n **kwargs: Additional parameters (unused)\n\n Returns:\n Any: Agent execution result\n\n Raises:\n ValueError: Invalid parameters format\n TimeoutError: Request timeout\n \"\"\"\n tool_call_id = kwargs.get(\"tool_call_id\")\n agent_name = self.name\n # Process messages: prefer messages, if not provided use content to create ChatMessage\n messages = AgentTool.parse_messages(parameters)\n\n # Build request data in format expected by orchestrator\n request_data = {\n \"agent_name\": agent_name,\n }\n\n if messages:\n request_data[\"messages\"] = messages\n\n request_data.update(kwargs)\n\n logger.debug(\n \"AgentTool._call: tool=%s, agent=%s, messages=%d, tool_call_id=%s\",\n self.name,\n agent_name,\n len(messages) if messages else 0,\n tool_call_id,\n )\n\n # Send request through Channel and wait for response\n try:\n tool_parameters = ToolParameters(parameters=parameters, kwargs=request_data)\n\n _, response = await self.channel.send_request(\n request_id=tool_call_id,\n tool_name=self.name,\n data=tool_parameters,\n tool_schema=self.get_schema(),\n timeout=self.timeout,\n )\n\n logger.debug(\n \"AgentTool._call completed: tool=%s, agent=%s, response_type=%s\",\n self.name,\n agent_name,\n type(response).__name__,\n )\n\n return response\n\n except Exception as e:\n logger.error(\n \"AgentTool._call failed: tool=%s, agent=%s, error=%s\",\n self.name,\n agent_name,\n e,\n exc_info=True,\n )\n raise\n\n @staticmethod\n def parse_messages(parameters: str) -> list[ChatMessage]:\n \"\"\"\n Parse and prepare message list.\n\n Args:\n parameters: Tool parameters (JSON string format)\n\n Returns:\n list[ChatMessage]: Processed message list\n \"\"\"\n # Parse parameters JSON string\n try:\n params_dict = json.loads(parameters) if parameters else {}\n except json.JSONDecodeError as e:\n raise ValueError(f\"Invalid parameters JSON format: {e}\") from e\n\n # Prioritize using messages\n messages = params_dict.get(\"messages\", [])\n if messages:\n processed = []\n for msg in messages:\n if isinstance(msg, dict):\n processed.append(ChatMessage(**msg))\n elif isinstance(msg, ChatMessage):\n processed.append(msg)\n else:\n processed.append(msg)\n return processed\n\n # If no messages, try using content\n content = params_dict.get(\"content\")\n if content:\n return [ChatMessage(role=\"user\", content=content)]\n\n return []\n" }, { "path": "cortex/tools/base.py", "content": "\"\"\"Base Tool class.\"\"\"\n\nfrom abc import ABC, abstractmethod\nfrom dataclasses import dataclass, field\nfrom typing import Any, Dict, Optional\n\nfrom .types import ToolType\n\n\n@dataclass\nclass ToolSchema:\n \"\"\"Tool schema definition.\"\"\"\n\n name: str\n description: str\n parameters: Dict[str, Any] = field(default_factory=dict)\n return_type: Optional[str] = None\n tool_type: Optional[ToolType] = None\n\n\nclass Tool(ABC):\n \"\"\"Base class for tools.\"\"\"\n\n def __init__(\n self,\n name: str,\n description: str = \"\",\n tool_type: Optional[ToolType] = None,\n **kwargs, # noqa: ARG002\n ):\n \"\"\"\n Initialize the tool.\n\n Args:\n name: Tool name\n description: Tool description\n tool_type: Tool type\n **kwargs: Additional parameters (passed to subclasses)\n \"\"\"\n self.name = name\n self.description = description\n self.tool_type = tool_type\n self._schema: Optional[ToolSchema] = None\n\n def get_schema(self) -> ToolSchema:\n \"\"\"\n Get tool schema.\n\n Returns:\n ToolSchema: Tool schema object\n \"\"\"\n if self._schema is None:\n self._schema = self._define_schema()\n return self._schema\n\n @abstractmethod\n def _define_schema(self) -> ToolSchema:\n \"\"\"\n Define tool schema (must be implemented by subclasses).\n\n Returns:\n ToolSchema: Tool schema object\n \"\"\"\n pass # pragma: no cover\n\n async def call(self, parameters: str, **kwargs) -> Any:\n \"\"\"\n Call the tool (async).\n\n Args:\n parameters: Tool parameters (string format)\n **kwargs: Additional parameters\n\n Returns:\n Any: Tool execution result\n \"\"\"\n return await self._call(parameters, **kwargs)\n\n @abstractmethod\n async def _call(self, parameters: str, **kwargs) -> Any:\n \"\"\"\n Tool call implementation (must be implemented by subclasses).\n\n Args:\n parameters: Tool parameters (string format)\n **kwargs: Additional parameters\n\n Returns:\n Any: Tool execution result\n \"\"\"\n" }, { "path": "cortex/tools/channel.py", "content": "\"\"\"Channel for async tool execution communication.\"\"\"\n\nimport asyncio\nfrom dataclasses import dataclass\nfrom enum import Enum\nfrom typing import Any, Awaitable, Callable, Dict, Optional, Tuple\n\nfrom .base import ToolSchema\nfrom .types import ToolParameters\n\n\nclass MessageType(Enum):\n \"\"\"Message type.\"\"\"\n\n REQUEST = \"request\"\n RESPONSE = \"response\"\n ERROR = \"error\"\n\n\n@dataclass\nclass ChannelMessage:\n \"\"\"Channel message.\"\"\"\n\n message_type: MessageType\n tool_name: str\n request_id: str\n data: Any\n error: Optional[str] = None\n\n\nclass Channel:\n \"\"\"Channel for async tool execution communication.\"\"\"\n\n def __init__(\n self,\n on_send: Optional[\n Callable[[str, ToolSchema, ToolParameters], Awaitable[None]]\n ] = None,\n ):\n \"\"\"\n Initialize Channel.\n\n Args:\n on_send: Callback function for sending requests,\n receives (tool_name, tool_schema, data) and sends data asynchronously\n \"\"\"\n self._pending_requests: Dict[str, asyncio.Future] = {}\n self._request_counter = 0\n self._on_send = on_send\n\n def set_on_send(\n self,\n on_send: Optional[Callable[[str, ToolSchema, ToolParameters], Awaitable[None]]],\n ):\n \"\"\"\n Set the send callback function.\n\n Args:\n on_send: Callback function for sending requests,\n receives (tool_name, tool_schema, tool_parameters) and sends data asynchronously\n \"\"\"\n self._on_send = on_send\n\n def create_request_id(self) -> str:\n \"\"\"\n Create a request ID.\n\n Returns:\n str: Request ID\n \"\"\"\n self._request_counter += 1\n return f\"req_{self._request_counter}\"\n\n async def send_request(\n self,\n tool_name: str,\n data: ToolParameters,\n tool_schema: ToolSchema,\n request_id: Optional[str] = None,\n timeout: Optional[float] = None,\n on_send: Optional[\n Callable[[str, ToolSchema, ToolParameters], Awaitable[None]]\n ] = None,\n ) -> Tuple[str, Any]:\n \"\"\"\n Send request and wait for response.\n\n Args:\n tool_name: Tool name\n data: Request data (ToolParameters object)\n tool_schema: Tool schema\n request_id: Request ID (auto-generated if not provided)\n timeout: Timeout in seconds\n on_send: Callback function for sending requests,\n receives (tool_name, tool_schema, tool_parameters) and sends data asynchronously.\n If provided, overrides the on_send set during initialization.\n\n Returns:\n Tuple[str, Any]: (request_id, response data)\n\n Raises:\n TimeoutError: Request timeout\n Exception: Tool execution error\n \"\"\"\n # Auto-generate request_id if not provided\n if request_id is None:\n request_id = self.create_request_id()\n\n future = asyncio.Future()\n self._pending_requests[request_id] = future\n\n try:\n # Send data to external system\n send_handler = on_send or self._on_send\n if send_handler:\n # Create a copy of data to avoid modifying the original object\n tool_parameters = ToolParameters(\n parameters=data.parameters, kwargs={**data.kwargs}\n )\n\n # Add request_id to kwargs\n tool_parameters.kwargs[\"_request_id\"] = request_id\n\n await send_handler(tool_name, tool_schema, tool_parameters)\n\n # Wait for response (response is set via set_response)\n if timeout:\n response = await asyncio.wait_for(future, timeout=timeout)\n else:\n response = await future\n\n return (request_id, response)\n\n except asyncio.TimeoutError as exc:\n self._pending_requests.pop(request_id, None)\n raise TimeoutError(\n f\"Request {request_id} for tool {tool_name} timed out\"\n ) from exc\n except Exception:\n self._pending_requests.pop(request_id, None)\n raise\n finally:\n self._pending_requests.pop(request_id, None)\n\n def set_response(self, request_id: str, data: Any, error: Optional[str] = None):\n \"\"\"\n Set response data.\n\n Args:\n request_id: Request ID\n data: Response data\n error: Error message (if any)\n \"\"\"\n if request_id not in self._pending_requests:\n return\n\n future = self._pending_requests[request_id]\n\n # Check if future has already been set\n if future.done():\n return\n\n try:\n if error:\n future.set_exception(Exception(error))\n else:\n future.set_result(data)\n except Exception:\n # Ignore if future has already been set or cancelled\n # This may happen in concurrent scenarios\n pass\n" }, { "path": "cortex/tools/client_tool.py", "content": "\"\"\"Client Tool implementation - uses Channel for communication.\"\"\"\n\nimport logging\nfrom typing import Any, Optional\n\nfrom .base import Tool, ToolSchema\nfrom .channel import Channel\nfrom .types import ToolParameters, ToolType\n\nlogger = logging.getLogger(__name__)\n\n\nclass ClientTool(Tool):\n \"\"\"Client tool implementation, uses Channel for async communication.\"\"\"\n\n def __init__(\n self,\n name: str,\n description: str = \"\",\n tool_type: Optional[ToolType] = None,\n channel: Optional[Channel] = None,\n timeout: Optional[float] = None,\n **kwargs,\n ):\n \"\"\"\n Initialize client tool.\n\n Args:\n name: Tool name\n description: Tool description\n channel: Channel instance (for async communication)\n timeout: Default timeout in seconds\n **kwargs: Additional parameters\n \"\"\"\n # If tool_type is not specified, default to CLIENT\n # But if name is \"ask_input\", use ASK_INPUT\n if tool_type is None:\n if name == \"ask_input\":\n tool_type = ToolType.ASK_INPUT\n else:\n tool_type = ToolType.CLIENT\n\n super().__init__(\n name=name, description=description, tool_type=tool_type, **kwargs\n )\n if channel is None:\n raise ValueError(\"ClientTool requires a Channel instance\")\n self.channel = channel\n self.timeout = timeout or 30.0\n self._client_params = kwargs.get(\"client_params\", {})\n\n def _define_schema(self) -> ToolSchema:\n \"\"\"Define client tool schema.\"\"\"\n properties = self._client_params.get(\"properties\", {})\n required = self._client_params.get(\"required\", [])\n\n return ToolSchema(\n name=self.name,\n description=self.description,\n parameters={\n \"type\": \"object\",\n \"properties\": properties,\n \"required\": required,\n },\n return_type=\"client_response\",\n tool_type=self.tool_type,\n )\n\n async def _call(self, parameters: str, **kwargs) -> Any:\n \"\"\"\n Call the client tool.\n\n Send request through Channel and wait for response.\n\n Args:\n parameters: Tool parameters (string format)\n **kwargs: Additional parameters\n\n Returns:\n Any: Execution result\n \"\"\"\n # Build request data, convert parameters and kwargs to ToolParameters\n tool_parameters = ToolParameters(parameters=parameters, kwargs=kwargs)\n\n tool_call_id = kwargs.get(\"tool_call_id\")\n if tool_call_id is None:\n tool_call_id = f\"tool_call_{hash(self.name)}_{hash(parameters)}\"\n\n logger.debug(\n \"ClientTool._call: tool=%s, parameters=%s, tool_call_id=%s\",\n self.name,\n parameters,\n tool_call_id,\n )\n\n # Send request through Channel and wait for response\n _, response = await self.channel.send_request(\n tool_name=self.name,\n data=tool_parameters,\n tool_schema=self.get_schema(),\n request_id=kwargs.get(\"tool_call_id\"),\n timeout=kwargs.get(\"timeout\") or self.timeout,\n )\n\n return response\n" }, { "path": "cortex/tools/function_tool.py", "content": "\"\"\"Function Tool implementation.\"\"\"\n\nimport inspect\nimport json\nfrom typing import Any, Callable, Optional\n\nfrom agents.function_schema import function_schema\n\nfrom .base import Tool, ToolSchema\nfrom .types import ToolType\n\n\nclass FunctionTool(Tool):\n \"\"\"Function tool implementation.\"\"\"\n\n def __init__(\n self, name: str, func: Callable, description: Optional[str] = None, **kwargs\n ):\n \"\"\"\n Initialize function tool.\n\n Args:\n name: Tool name\n func: Function to wrap\n description: Tool description\n **kwargs: Additional parameters\n \"\"\"\n self.func = func\n description = description or func.__doc__ or \"\"\n\n super().__init__(\n name=name, description=description, tool_type=ToolType.FUNCTION, **kwargs\n )\n\n def _define_schema(self) -> ToolSchema:\n \"\"\"Define function tool schema.\"\"\"\n # Infer parameters from function signature\n # strict_json_schema=False to preserve default value behavior for optional parameters\n data = function_schema(self.func, strict_json_schema=False)\n return ToolSchema(\n name=data.name,\n description=data.description,\n parameters=data.params_json_schema,\n return_type=\"any\",\n )\n\n async def _call(self, parameters: str, **kwargs) -> Any:\n \"\"\"\n Call the function tool.\n\n Args:\n parameters: Tool parameters (string format)\n **kwargs: Additional parameters\n\n Returns:\n Any: Function execution result\n \"\"\"\n # Convert parameters to dictionary\n parameters_dict = json.loads(parameters)\n # If the function is async\n if inspect.iscoroutinefunction(self.func):\n return await self.func(**parameters_dict)\n else:\n # Call sync function in async context\n import asyncio\n\n return await asyncio.to_thread(self.func, **parameters_dict)\n" }, { "path": "cortex/tools/mcp.py", "content": "import logging\nfrom contextlib import AsyncExitStack\nfrom typing import Any, final\n\nfrom mcp import ClientSession, Tool\nfrom mcp.client.streamable_http import streamablehttp_client\nfrom mcp.types import CallToolResult\n\nlogger = logging.getLogger(__name__)\n\n\nclass MCPNotInitializedError(Exception):\n \"\"\"Exception raised when MCP client is not initialized.\"\"\"\n\n def __init__(self) -> None:\n super().__init__(\"MCP client is not initialized\")\n\n\n@final\nclass MCPClient:\n def __init__(self, server_url: str) -> None:\n self.server_url = server_url\n self.session = None\n self.exit_stack = AsyncExitStack()\n\n async def initialize(self) -> None:\n read_stream, write_stream, _ = await self.exit_stack.enter_async_context(\n streamablehttp_client(self.server_url)\n )\n self.session = await self.exit_stack.enter_async_context(\n ClientSession(read_stream, write_stream)\n )\n logger.info(f\"Connected to server: {self.server_url}\")\n await self.session.initialize()\n logger.info(\"Initialized session\")\n\n async def aclose(self) -> None:\n \"\"\"Close all resources opened via the exit stack.\n\n This must be called from the same task that created the contexts to\n avoid AnyIO cancel scope errors during shutdown.\n \"\"\"\n try:\n await self.exit_stack.aclose()\n finally:\n self.session = None\n\n async def list_tools(self) -> list[Tool]:\n if self.session is None:\n raise MCPNotInitializedError()\n\n result = await self.session.list_tools()\n return result.tools\n\n async def call_tool(\n self, tool_name: str, arguments: dict[str, Any]\n ) -> CallToolResult:\n if self.session is None:\n raise MCPNotInitializedError()\n return await self.session.call_tool(tool_name, arguments)\n" }, { "path": "cortex/tools/mcp_tool.py", "content": "\"\"\"MCP (Model Context Protocol) Tool implementation.\"\"\"\n\nimport json\nfrom typing import Any\n\nfrom .base import Tool, ToolSchema\nfrom .mcp import MCPClient\nfrom .types import ToolType\n\n\nclass MCPTool(Tool):\n \"\"\"MCP tool implementation.\"\"\"\n\n def __init__(\n self, name: str, description: str = \"\", mcp_server: str = None, **kwargs\n ):\n \"\"\"\n Initialize MCP tool.\n\n Args:\n name: Tool name\n description: Tool description\n mcp_server: MCP server\n **kwargs: Additional parameters\n \"\"\"\n super().__init__(\n name=name, description=description, tool_type=ToolType.MCP, **kwargs\n )\n self.mcp_server = mcp_server\n self._mcp_params = kwargs.get(\"mcp_params\", {})\n # MCP tool can handle directly, no Channel needed\n\n def _define_schema(self) -> ToolSchema:\n \"\"\"Define MCP tool schema.\"\"\"\n return ToolSchema(\n name=self.name,\n description=self.description,\n parameters={\n \"type\": \"object\",\n \"properties\": self._mcp_params.get(\"properties\", {}),\n \"required\": self._mcp_params.get(\"required\", []),\n },\n return_type=\"mcp_response\",\n )\n\n async def _call(self, parameters: str, **kwargs) -> Any:\n \"\"\"\n Call MCP tool.\n\n Args:\n parameters: Tool parameters (string format)\n **kwargs: Additional parameters\n\n Returns:\n Any: Execution result\n \"\"\"\n arguments = json.loads(parameters)\n mcp_client = MCPClient(self.mcp_server)\n await mcp_client.initialize()\n try:\n result = await mcp_client.call_tool(self.name, arguments)\n return result\n finally:\n await mcp_client.aclose()\n" }, { "path": "cortex/tools/session_tool.py", "content": "\"\"\"Session Tool implementation.\"\"\"\n\nfrom typing import Any, Dict, Optional\n\nfrom .base import Tool, ToolSchema\nfrom .types import ToolType\n\n\nclass SessionTool(Tool):\n \"\"\"Session tool implementation for maintaining session state.\"\"\"\n\n def __init__(\n self,\n name: str,\n description: str = \"\",\n session_id: Optional[str] = None,\n **kwargs,\n ):\n \"\"\"\n Initialize session tool.\n\n Args:\n name: Tool name\n description: Tool description\n session_id: Session ID\n **kwargs: Additional parameters\n \"\"\"\n super().__init__(\n name=name, description=description, tool_type=ToolType.SESSION, **kwargs\n )\n self.session_id = session_id\n self._session_state: Dict[str, Any] = {}\n\n def _define_schema(self) -> ToolSchema:\n \"\"\"Define session tool schema.\"\"\"\n return ToolSchema(\n name=self.name,\n description=self.description,\n parameters={\n \"type\": \"object\",\n \"properties\": {\n \"action\": {\n \"type\": \"string\",\n \"description\": \"Session action (get, set, update, clear)\",\n \"enum\": [\"get\", \"set\", \"update\", \"clear\"],\n },\n \"key\": {\"type\": \"string\", \"description\": \"Session key\"},\n \"value\": {\"type\": \"any\", \"description\": \"Session value\"},\n },\n \"required\": [\"action\"],\n },\n return_type=\"session_response\",\n )\n\n async def _call(self, parameters: str, **kwargs) -> Any:\n \"\"\"\n Call session tool.\n\n Args:\n parameters: Tool parameters (string format)\n **kwargs: Additional parameters\n\n Returns:\n Any: Execution result\n \"\"\"\n action = kwargs.get(\"action\", \"get\")\n\n if action == \"get\":\n key = kwargs.get(\"key\")\n if key:\n return self._session_state.get(key)\n return self._session_state\n elif action == \"set\":\n key = kwargs.get(\"key\")\n value = kwargs.get(\"value\")\n if key:\n self._session_state[key] = value\n return {\"status\": \"ok\", \"key\": key, \"value\": value}\n elif action == \"update\":\n updates = kwargs.get(\"value\", {})\n self._session_state.update(updates)\n return {\"status\": \"ok\", \"updated\": list(updates.keys())}\n elif action == \"clear\":\n self._session_state.clear()\n return {\"status\": \"ok\", \"message\": \"session cleared\"}\n else:\n raise ValueError(f\"Unknown action: {action}\")\n" }, { "path": "cortex/tools/toolset.py", "content": "\"\"\"ToolSet for managing and executing tools.\"\"\"\n\nimport logging\nfrom typing import Any, Awaitable, Callable, Dict, List, Optional, Type\n\nfrom agentkit.trace import get_current_context\n\nfrom .agent_tool import AgentTool\nfrom .base import Tool, ToolSchema\nfrom .channel import Channel\nfrom .client_tool import ClientTool\nfrom .function_tool import FunctionTool\nfrom .mcp import MCPClient\nfrom .mcp_tool import MCPTool\nfrom .session_tool import SessionTool\nfrom .types import ToolConfig, ToolParameters, ToolType\n\nlogger = logging.getLogger(__name__)\n\n\nclass ToolSet:\n \"\"\"Tool collection that manages tool registration, initialization, and invocation.\"\"\"\n\n _client_tool_results: dict[str, Any] = {}\n\n def __init__(\n self,\n channel: Optional[Channel] = None,\n on_send: Optional[\n Callable[[str, ToolSchema, ToolParameters], Awaitable[None]]\n ] = None,\n ):\n \"\"\"\n Initialize ToolSet.\n\n Args:\n channel: Shared Channel instance (optional)\n on_send: Callback function for sending requests, passed to Channel\n (if channel is not provided, a new Channel will be created with this callback)\n \"\"\"\n self._tools: Dict[str, Tool] = {}\n self._tool_factories: Dict[ToolType, Type[Tool]] = {\n ToolType.MCP: MCPTool,\n ToolType.FUNCTION: FunctionTool,\n ToolType.SESSION: SessionTool,\n ToolType.CLIENT: ClientTool,\n ToolType.AGENT: AgentTool,\n ToolType.ASK_INPUT: ClientTool, # ASK_INPUT uses ClientTool implementation\n }\n\n if channel:\n self.channel = channel\n # If on_send parameter is provided, set it (overrides existing)\n if on_send:\n self.channel.set_on_send(on_send)\n else:\n # If no channel is provided, create a new channel with on_send\n self.channel = Channel(on_send=on_send)\n\n def set_on_send(\n self, on_send: Callable[[str, ToolSchema, ToolParameters], Awaitable[None]]\n ) -> None:\n \"\"\"\n Set the on_send callback function for ToolSet.\n\n Args:\n on_send: Callback function for sending requests,\n receives (tool_name, tool_schema, tool_parameters) and sends data asynchronously\n\n Returns:\n None\n \"\"\"\n self.channel.set_on_send(on_send)\n\n def register(self, tool: Tool, name: Optional[str] = None) -> None:\n \"\"\"\n Register a tool.\n\n Args:\n tool: Tool instance\n name: Tool name (if not provided, uses tool.name)\n \"\"\"\n tool_name = name or tool.name\n\n if tool_name in self._tools:\n raise ValueError(f\"Tool '{tool_name}' is already registered\")\n\n # If it's a ClientTool or AgentTool without a channel, set the shared channel\n if isinstance(tool, (ClientTool, AgentTool)) and tool.channel != self.channel:\n tool.channel = self.channel\n\n self._tools[tool_name] = tool\n logger.info(f\"✓ Registered tool: {tool_name} ({tool.tool_type.value})\")\n\n async def register_from_mcp_server(\n self, mcp_server: str, tool_names: list[str] | None = None\n ) -> None:\n \"\"\"\n Register tools from an MCP server.\n\n Args:\n mcp_server: The MCP server URL\n tool_names: Optional list of tool names to register. If None, all available tools will be registered.\n \"\"\"\n try:\n mcp_client = MCPClient(mcp_server)\n await mcp_client.initialize()\n mcp_tools = await mcp_client.list_tools()\n if tool_names is not None:\n mcp_tools = [tool for tool in mcp_tools if tool.name in tool_names]\n for mcp_tool in mcp_tools:\n # Extract properties and required from inputSchema\n input_schema = mcp_tool.inputSchema or {}\n mcp_params = {\n \"properties\": input_schema.get(\"properties\", {}),\n \"required\": input_schema.get(\"required\", []),\n }\n\n server_tool = MCPTool(\n name=mcp_tool.name,\n description=mcp_tool.description,\n mcp_server=mcp_server,\n mcp_params=mcp_params,\n )\n self.register(server_tool, mcp_tool.name)\n finally:\n await mcp_client.aclose()\n\n def register_from_config(self, config: ToolConfig) -> Tool:\n \"\"\"\n Initialize and register a tool from configuration.\n\n Args:\n config: Tool configuration\n\n Returns:\n Tool: Created tool instance\n \"\"\"\n factory = self._tool_factories.get(config.tool_type)\n if not factory:\n raise ValueError(f\"Unknown tool type: {config.tool_type}\")\n\n # Prepare initialization parameters\n init_params = config.params or {}\n\n if config.tool_type == ToolType.MCP:\n init_params.setdefault(\"endpoint\", init_params.get(\"endpoint\"))\n elif (\n config.tool_type == ToolType.CLIENT\n or config.tool_type == ToolType.ASK_INPUT\n ):\n # ClientTool and ASK_INPUT must use channel, use ToolSet's channel if not provided\n if \"channel\" not in init_params:\n init_params[\"channel\"] = self.channel\n elif config.tool_type == ToolType.AGENT:\n # AgentTool must use channel, use ToolSet's channel if not provided\n if \"channel\" not in init_params:\n init_params[\"channel\"] = self.channel\n elif config.tool_type == ToolType.FUNCTION:\n if \"func\" not in init_params:\n raise ValueError(\"Function tool requires 'func' parameter\")\n\n # Create tool instance\n tool = factory(\n name=config.name,\n description=init_params.get(\"description\", \"\"),\n **init_params,\n )\n\n # Register tool\n self.register(tool, config.name)\n\n return tool\n\n def get_tool(self, name: str) -> Optional[Tool]:\n \"\"\"\n Get a tool by name.\n\n Args:\n name: Tool name\n\n Returns:\n Tool: Tool instance, or None if not found\n \"\"\"\n return self._tools.get(name)\n\n def list_tools(self) -> List[str]:\n \"\"\"\n List all registered tool names.\n\n Returns:\n List[str]: List of tool names\n \"\"\"\n return list(self._tools.keys())\n\n async def call(self, tool_name: str, parameters: str, **kwargs) -> Any:\n \"\"\"\n Call a tool.\n\n Args:\n tool_name: Tool name\n parameters: Tool parameters\n **kwargs: Additional parameters\n\n Returns:\n Any: Tool execution result\n\n Raises:\n ValueError: Tool not found\n \"\"\"\n ctx = get_current_context()\n with ctx.tool_span(name=f\"ToolSet.call {tool_name}\") as span:\n span.update_payload_data(\n request=kwargs,\n )\n tool = self.get_tool(tool_name)\n if not tool:\n raise ValueError(f\"Tool '{tool_name}' is not registered\")\n\n logger.info(\n f\"ToolSet.call {tool_name} parameters: {parameters} kwargs: {kwargs}\"\n )\n resp = await tool.call(parameters, **kwargs)\n\n span.update_payload_data(\n response=resp,\n )\n return resp\n\n def get_schema(self, tool_name: str) -> Any:\n \"\"\"\n Get tool schema.\n\n Args:\n tool_name: Tool name\n\n Returns:\n Any: Tool schema\n\n Raises:\n ValueError: Tool not found\n \"\"\"\n tool = self.get_tool(tool_name)\n if not tool:\n raise ValueError(f\"Tool '{tool_name}' is not registered\")\n\n return tool.get_schema()\n\n def get_all_schemas(self) -> Dict[str, Any]:\n \"\"\"\n Get schemas for all tools.\n\n Returns:\n Dict[str, Any]: Mapping of tool names to schemas\n \"\"\"\n return {name: tool.get_schema() for name, tool in self._tools.items()}\n\n def get_client_tool_call_result(self, tool_call_id: str) -> Any:\n \"\"\"\n Get the result of a ClientTool call.\n\n Args:\n tool_call_id: Tool call ID\n\n Returns:\n Any: Tool call result\n \"\"\"\n result, error = self._client_tool_results.get(tool_call_id, (None, None))\n if error:\n raise Exception(error)\n return result\n\n def set_client_tool_call_result(\n self, tool_call_id: str, result: Any, error: Optional[str] = None\n ):\n \"\"\"\n Set the result of a ClientTool call.\n \"\"\"\n self._client_tool_results[tool_call_id] = (result, error)\n\n def set_response(self, request_id: str, data: Any, error: Optional[str] = None):\n \"\"\"\n Set ClientTool response.\n\n Args:\n request_id: Request ID\n data: Response data\n error: Error message (if any)\n\n Raises:\n ValueError: Tool not found or not a ClientTool\n \"\"\"\n # tool = self.get_tool(tool_name)\n # if not tool:\n # raise ValueError(f\"Tool '{tool_name}' is not registered\")\n\n # if not isinstance(tool, ClientTool):\n # raise ValueError(\n # f\"Tool '{tool_name}' is not a ClientTool, cannot set response\"\n # )\n\n # Set response through Channel\n self.channel.set_response(request_id, data, error)\n self.set_client_tool_call_result(request_id, data, error)\n" }, { "path": "cortex/tools/types.py", "content": "\"\"\"Tool types and enums.\"\"\"\n\nfrom dataclasses import dataclass\nfrom enum import Enum\nfrom typing import Any, Dict, Optional\n\nfrom pydantic import BaseModel, Field\n\n\nclass ToolParameters(BaseModel):\n \"\"\"Tool parameters model.\"\"\"\n\n parameters: str = Field(description=\"Tool parameters\", default=\"\")\n kwargs: Dict[str, Any] = Field(description=\"Additional parameters\", default={})\n\n\nclass ToolType(Enum):\n \"\"\"Tool type enum.\"\"\"\n\n MCP = \"mcp\"\n FUNCTION = \"function\"\n SESSION = \"session\"\n CLIENT = \"client\"\n AGENT = \"agent\"\n ASK_INPUT = \"ask_input\"\n\n\nclass ExecutionType(Enum):\n \"\"\"Execution type enum.\"\"\"\n\n SYNC = \"sync\"\n ASYNC = \"async\"\n\n\n@dataclass\nclass ToolConfig:\n \"\"\"Tool configuration.\"\"\"\n\n name: str\n tool_type: ToolType\n params: Optional[Dict[str, Any]] = None\n schema: Optional[Dict[str, Any]] = None\n" }, { "path": "cortex/tools/ublock_agent_tool.py", "content": "\"\"\"Agent Tool implementation - calls agents through Orchestrator.\"\"\"\n\nimport json\nimport logging\nfrom typing import Any, Optional\n\nfrom cortex.model.definition import ChatMessage\n\nfrom .base import Tool, ToolSchema\nfrom .channel import Channel\nfrom .types import ToolParameters, ToolType\n\nlogger = logging.getLogger(__name__)\n\n\nclass UnblockAgentTool(Tool):\n \"\"\"Agent tool implementation for calling other Agents through Channel communication with Orchestrator.\"\"\"\n\n share_context: bool = False\n\n def __init__(\n self,\n name: str,\n description: str = \"\",\n channel: Optional[Channel] = None,\n timeout: Optional[float] = None,\n share_context: bool = False,\n **kwargs,\n ):\n \"\"\"\n Initialize Agent tool.\n\n Args:\n name: Tool name.\n description: Tool description.\n channel: Channel instance (for async communication).\n timeout: Default timeout in seconds.\n share_context: Whether to share context.\n **kwargs: Other parameters.\n \"\"\"\n super().__init__(\n name=name, description=description, tool_type=ToolType.AGENT, **kwargs\n )\n if channel is None:\n raise ValueError(\"AgentTool requires a Channel instance\")\n self.channel = channel\n self.timeout = timeout or 30.0\n logger.debug(\n \"AgentTool initialized: name=%s, description=%s, timeout=%s\",\n name,\n description,\n self.timeout,\n )\n\n def _define_schema(self) -> ToolSchema:\n \"\"\"Define Agent tool schema.\"\"\"\n schema = ToolSchema(\n name=self.name,\n description=self.description,\n parameters={\n \"type\": \"object\",\n \"properties\": {\n \"content\": {\n \"type\": \"string\",\n \"description\": \"Specific instruction content to send to Agent (optional, ignored if messages is provided)\",\n },\n \"messages\": {\n \"type\": \"array\",\n \"description\": \"List of messages to send to Agent (optional, ChatMessage format with role and content fields)\",\n \"items\": {\n \"type\": \"object\",\n \"properties\": {\n \"role\": {\n \"type\": \"string\",\n \"description\": \"Message role, e.g., 'user', 'assistant', 'system'\",\n \"enum\": [\"user\", \"assistant\", \"system\", \"tool\"],\n },\n \"content\": {\n \"type\": \"string\",\n \"description\": \"Message content (string format)\",\n },\n },\n \"required\": [\"role\", \"content\"],\n },\n },\n \"timeout\": {\n \"type\": \"number\",\n \"description\": \"Timeout in seconds (optional)\",\n },\n },\n \"required\": [],\n },\n return_type=\"agent_response\",\n tool_type=self.tool_type,\n )\n logger.debug(\n \"AgentTool schema defined: name=%s, parameters=%s\", self.name, schema.parameters\n )\n return schema\n\n async def _call(self, parameters: str, **kwargs) -> Any:\n \"\"\"\n Call Agent tool.\n\n Send request to Orchestrator through Channel, which creates new Agent and executes task.\n\n Args:\n parameters: Tool parameters (JSON string format), containing:\n - content: Specific instruction content (optional, ignored if messages is provided)\n - messages: Message list (optional, ChatMessage format with role and content fields)\n - timeout: Timeout in seconds (optional)\n **kwargs: Other parameters (unused)\n\n Returns:\n Any: Agent execution result\n\n Raises:\n ValueError: Invalid parameters format\n TimeoutError: Request timeout\n \"\"\"\n tool_call_id = kwargs.get(\"tool_call_id\")\n agent_name = self.name\n # Process messages: prefer messages, use content to create ChatMessage if not available\n messages = UnblockAgentTool.parse_messages(parameters)\n\n # Build request data in orchestrator expected format\n request_data = {\n \"agent_name\": agent_name,\n }\n\n if messages:\n request_data[\"messages\"] = messages\n\n request_data.update(kwargs)\n\n logger.debug(\n \"UnblockAgentTool._call: tool=%s, agent=%s, messages=%d, tool_call_id=%s\",\n self.name,\n agent_name,\n len(messages) if messages else 0,\n tool_call_id,\n )\n\n # Send request through Channel and wait for response\n tool_parameters = ToolParameters(parameters=parameters, kwargs=request_data)\n await self.channel._on_send(self.name, self.get_schema(), tool_parameters)\n return None\n\n @staticmethod\n def parse_messages(parameters: str) -> list[ChatMessage]:\n \"\"\"\n Prepare message list.\n\n Args:\n parameters: Tool parameters (JSON string format)\n\n Returns:\n list[ChatMessage]: Processed message list\n \"\"\"\n # Parse parameters JSON string\n try:\n params_dict = json.loads(parameters) if parameters else {}\n except json.JSONDecodeError as e:\n raise ValueError(f\"Invalid parameters JSON format: {e}\") from e\n\n # Prefer messages\n messages = params_dict.get(\"messages\", [])\n if messages:\n processed = []\n for msg in messages:\n if isinstance(msg, dict):\n processed.append(ChatMessage(**msg))\n elif isinstance(msg, ChatMessage):\n processed.append(msg)\n else:\n processed.append(msg)\n return processed\n\n # If no messages, try to use content\n content = params_dict.get(\"content\")\n if content:\n return [ChatMessage(role=\"user\", content=content)]\n\n return []\n" }, { "path": "cortex/tools/unblock_client_tool.py", "content": "\"\"\"Client Tool implementation - uses Channel for communication.\"\"\"\n\nimport logging\nfrom typing import Any, Optional\n\nfrom .base import Tool, ToolSchema\nfrom .channel import Channel\nfrom .types import ToolParameters, ToolType\n\nlogger = logging.getLogger(__name__)\n\n\nclass UnblockClientTool(Tool):\n \"\"\"Client tool implementation using Channel for asynchronous communication.\"\"\"\n\n def __init__(\n self,\n name: str,\n description: str = \"\",\n tool_type: Optional[ToolType] = None,\n channel: Optional[Channel] = None,\n timeout: Optional[float] = None,\n **kwargs,\n ):\n \"\"\"\n Initialize Client tool.\n\n Args:\n name: Tool name\n description: Tool description\n channel: Channel instance (for async communication)\n timeout: Default timeout in seconds\n **kwargs: Additional parameters\n \"\"\"\n # If tool_type is not specified, default to CLIENTd, default to CLIENT\n # But if name is \"ask_input\", use ASK_INPUT\n if tool_type is None:\n if name == \"ask_input\":\n tool_type = ToolType.ASK_INPUT\n else:\n tool_type = ToolType.CLIENT\n\n super().__init__(\n name=name, description=description, tool_type=tool_type, **kwargs\n )\n\n self.timeout = timeout or 30.0\n self._client_params = kwargs.get(\"client_params\", {})\n self.channel = channel\n\n def _define_schema(self) -> ToolSchema:\n \"\"\"Define Client tool schema.\"\"\"\n properties = self._client_params.get(\"properties\", {})\n required = self._client_params.get(\"required\", [])\n\n return ToolSchema(\n name=self.name,\n description=self.description,\n parameters={\n \"type\": \"object\",\n \"properties\": properties,\n \"required\": required,\n },\n return_type=\"client_response\",\n tool_type=self.tool_type,\n )\n\n async def _call(self, parameters: str, **kwargs) -> Any:\n \"\"\"\n Call the Client tool.\n\n Send request through Channel and wait for response.\n\n Args:\n parameters: Tool parameters (string format)\n **kwargs: Additional parameters\n\n Returns:\n Any: Execution result\n \"\"\"\n # Build request data, convert parameters and kwargs to ToolParameters\n tool_parameters = ToolParameters(parameters=parameters, kwargs=kwargs)\n\n tool_call_id = kwargs.get(\"tool_call_id\")\n if tool_call_id is None:\n tool_call_id = f\"tool_call_{hash(self.name)}_{hash(parameters)}\"\n\n logger.debug(\n \"ClientTool._call: tool=%s, parameters=%s, tool_call_id=%s\",\n self.name,\n parameters,\n tool_call_id,\n )\n\n await self.channel._on_send(self.name, self.get_schema(), tool_parameters)\n return None\n" }, { "path": "cortex/tui/__init__.py", "content": "\"\"\"Agent TUI module - Provides Textual-based TUI interface\"\"\"\n\nfrom cortex.tui.tui import AgentTUIApp\n\n__all__ = [\"AgentTUIApp\"]\n" }, { "path": "cortex/tui/tui.py", "content": "\"\"\"Agent TUI - TUI interface built with Textual and Rich, displaying AgentEvent\"\"\"\n\nimport asyncio\nimport json\nimport logging\nfrom pathlib import Path\nfrom typing import Callable, Optional\n\nfrom agentkit.trace import SpanContext, Tracer\nfrom cortex.model.definition import ChatMessage\nfrom rich.console import Group, RenderableType\nfrom rich.markdown import Markdown\nfrom rich.panel import Panel\nfrom rich.text import Text\nfrom textual import on, work\nfrom textual.app import App, ComposeResult\nfrom textual.binding import Binding\nfrom textual.containers import Container, Horizontal, Vertical, VerticalScroll\nfrom textual.events import Key\nfrom textual.widgets import Input, Label, ListItem, ListView, Static\n\nfrom cortex.agents.types import AgentConfig, AgentResponse, AgentRunningStatus\nfrom cortex.orchestrator import AgentEvent\nfrom cortex.orchestrator.orchestrator import Orchestrator\nfrom cortex.orchestrator.types import AgentEventType, AgentRequest, ClientToolCallType\n\nlogger = logging.getLogger(__name__)\n\n\ndef _content_to_string(content) -> str:\n \"\"\"Convert content to string\n\n Args:\n content: Can be string, list, or dict\n\n Returns:\n str: Converted string\n \"\"\"\n if content is None:\n return \"\"\n\n if isinstance(content, str):\n return content\n\n if isinstance(content, list):\n # Handle list format content, e.g., [{\"type\": \"text\", \"text\": \"...\"}]\n text_parts = []\n for item in content:\n if isinstance(item, dict):\n # Try various possible key names to get text content\n text = None\n # Common key names\n for key in [\"text\", \"content\", \"value\"]:\n if key in item:\n text = item[key]\n break\n # If not found, try to find fields with string values\n if not text:\n for key, value in item.items():\n if key != \"type\" and isinstance(value, str):\n text = value\n break\n if text:\n text_parts.append(str(text))\n elif isinstance(item, str):\n text_parts.append(item)\n return \"\\n\".join(text_parts)\n\n if isinstance(content, dict):\n # Handle dict format\n # Try various possible key names\n for key in [\"text\", \"content\", \"value\"]:\n if key in content:\n return str(content[key])\n # If not found, return string representation of entire dict\n return str(content)\n\n return str(content)\n\n\nclass EventItem(ListItem):\n \"\"\"Event list item\"\"\"\n\n def __init__(self, event: AgentEvent, index: int):\n self.event = event\n self.index = index\n super().__init__()\n\n def render(self) -> RenderableType:\n \"\"\"Render event item\"\"\"\n event = self.event\n event_type = event.type.value\n agent_name = event.agent_name or \"N/A\"\n task_id = (\n event.task_id[:8] + \"...\"\n if event.task_id and len(event.task_id) > 8\n else (event.task_id or \"N/A\")\n )\n\n # Set color based on event type\n type_colors = {\n \"request\": \"blue\",\n \"response\": \"green\",\n \"error\": \"red\",\n \"signal\": \"yellow\",\n \"client_tool_call\": \"cyan\",\n \"client_tool_result\": \"magenta\",\n }\n color = type_colors.get(event_type, \"white\")\n\n # Build display text\n lines = [\n Text(f\"[{self.index}] \", style=\"dim\"),\n Text(f\"[{event_type}]\", style=f\"bold {color}\"),\n Text(f\" Agent: {agent_name}\", style=\"cyan\"),\n Text(f\" Task: {task_id}\", style=\"dim\"),\n ]\n\n # Add response content\n if event.response and event.response.message:\n # Check if there are tool calls\n tool_calls = getattr(event.response.message, \"tool_calls\", None)\n has_tool_calls = False\n if tool_calls:\n if isinstance(tool_calls, (list, tuple)) and len(tool_calls) > 0:\n has_tool_calls = True\n elif not isinstance(tool_calls, (list, tuple)):\n has_tool_calls = True\n\n raw_content = event.response.message.content or \"\"\n content = _content_to_string(raw_content)\n\n # If no tool calls, show full markdown content\n if not has_tool_calls and content:\n # Use markdown to render full content\n header = Text.assemble(\n Text(f\"[{self.index}] \", style=\"dim\"),\n Text(f\"[{event_type}]\", style=f\"bold {color}\"),\n Text(f\" Agent: {agent_name}\", style=\"cyan\"),\n Text(f\" Task: {task_id}\", style=\"dim\"),\n )\n markdown_content = Markdown(content)\n return Panel(\n Group(header, markdown_content),\n border_style=color,\n title=f\"Event #{self.index}\",\n )\n else:\n # With tool calls, only show preview\n if content:\n preview = content[:50] + \"...\" if len(content) > 50 else content\n lines.append(Text(f\"\\n {preview}\", style=\"dim\"))\n\n # Add tool call info\n if tool_calls:\n if isinstance(tool_calls, (list, tuple)):\n tool_names = []\n for tool_call in tool_calls:\n # Try to get tool name\n if hasattr(tool_call, \"function\"):\n if hasattr(tool_call.function, \"name\"):\n tool_names.append(tool_call.function.name)\n elif isinstance(tool_call.function, dict):\n tool_names.append(\n tool_call.function.get(\"name\", \"unknown\")\n )\n elif isinstance(tool_call, dict):\n func = tool_call.get(\"function\", {})\n tool_names.append(\n func.get(\"name\", \"unknown\")\n if isinstance(func, dict)\n else \"unknown\"\n )\n if tool_names:\n tools_str = \", \".join(tool_names)\n lines.append(Text(f\"\\n 🔧 Tools: {tools_str}\", style=\"yellow\"))\n else:\n # Single tool_call\n if hasattr(tool_calls, \"function\"):\n if hasattr(tool_calls.function, \"name\"):\n lines.append(\n Text(\n f\"\\n 🔧 Tool: {tool_calls.function.name}\",\n style=\"yellow\",\n )\n )\n\n # Add client_tool_call info\n if event.client_tool_call:\n tool_name = \"unknown\"\n tool_call_id = (\n getattr(event.client_tool_call, \"tool_call_id\", None) or \"N/A\"\n )\n\n if hasattr(event.client_tool_call, \"function\"):\n if hasattr(event.client_tool_call.function, \"name\"):\n tool_name = event.client_tool_call.function.name\n elif isinstance(event.client_tool_call.function, dict):\n tool_name = event.client_tool_call.function.get(\"name\", \"unknown\")\n\n # Check if it's ask_input type\n tool_type = getattr(event.client_tool_call, \"type\", None)\n\n # Get function.arguments\n args_dict = {}\n context_content = None\n prompt_content = None\n if hasattr(event.client_tool_call, \"function\"):\n func = event.client_tool_call.function\n if hasattr(func, \"arguments\"):\n args = func.arguments\n if isinstance(args, str):\n try:\n args_dict = json.loads(args)\n except json.JSONDecodeError:\n args_dict = {\"raw\": args}\n elif isinstance(args, dict):\n args_dict = args\n else:\n args_dict = {}\n\n # Check if context field exists\n if isinstance(args_dict, dict) and \"context\" in args_dict:\n context_content = args_dict.get(\"context\")\n # Remove context from args_dict, handle separately\n args_dict = {\n k: v for k, v in args_dict.items() if k != \"context\"\n }\n\n # Check if prompt field exists\n if isinstance(args_dict, dict) and \"prompt\" in args_dict:\n prompt_content = args_dict.get(\"prompt\")\n # Remove prompt from args_dict, handle separately\n args_dict = {\n k: v for k, v in args_dict.items() if k != \"prompt\"\n }\n\n # Set icon and style based on type\n if tool_type == ClientToolCallType.ASK_INPUT:\n icon = \"❓\"\n style_color = \"bold yellow\"\n type_label = \"Ask Input\"\n elif tool_type == ClientToolCallType.AGENT:\n icon = \"🤖\"\n style_color = \"bold cyan\"\n type_label = \"Agent Tool\"\n else:\n icon = \"🔧\"\n style_color = \"bold cyan\"\n type_label = \"Client Tool\"\n\n # Show tool call basic info\n lines.append(\n Text(f\"\\n {icon} {type_label}: {tool_name}\", style=style_color)\n )\n lines.append(Text(f\" Tool Call ID: {tool_call_id}\", style=\"dim\"))\n\n # Show all parameters (except context)\n if args_dict:\n lines.append(\n Text(\"\\n Parameters:\", style=f\"bold {style_color.split()[-1]}\")\n )\n for key, value in args_dict.items():\n value_str = (\n json.dumps(value, ensure_ascii=False, indent=2)\n if isinstance(value, (dict, list))\n else str(value)\n )\n # If parameter value is too long, display on multiple lines\n if len(value_str) > 200:\n lines.append(Text(f\" {key}:\", style=style_color))\n # Display long content on multiple lines\n for line in value_str.split(\"\\n\"):\n if line.strip():\n lines.append(Text(f\" {line}\", style=\"dim\"))\n else:\n lines.append(Text(f\" {key}: {value_str}\", style=\"dim\"))\n\n # If context exists, display with markdown\n if context_content is not None:\n context_str = _content_to_string(context_content)\n if context_str:\n lines.append(\n Text(\"\\n Context:\", style=f\"bold {style_color.split()[-1]}\")\n )\n # Add context content to lines, will be rendered with Markdown later\n # Add a marker here to indicate markdown rendering is needed\n lines.append((\"markdown_context\", context_str))\n\n # If prompt exists, display with markdown\n if prompt_content is not None:\n prompt_str = _content_to_string(prompt_content)\n if prompt_str:\n lines.append(\n Text(\"\\n Prompt:\", style=f\"bold {style_color.split()[-1]}\")\n )\n # Add prompt content to lines, will be rendered with Markdown later\n lines.append((\"markdown_prompt\", prompt_str))\n\n # Show extra info\n if (\n hasattr(event.client_tool_call, \"extra\")\n and event.client_tool_call.extra\n ):\n lines.append(\n Text(\"\\n Extra Info:\", style=f\"bold {style_color.split()[-1]}\")\n )\n for key, value in event.client_tool_call.extra.items():\n value_str = (\n json.dumps(value, ensure_ascii=False, indent=2)\n if isinstance(value, (dict, list))\n else str(value)\n )\n if len(value_str) > 200:\n lines.append(Text(f\" {key}:\", style=style_color))\n for line in value_str.split(\"\\n\"):\n if line.strip():\n lines.append(Text(f\" {line}\", style=\"dim\"))\n else:\n lines.append(Text(f\" {key}: {value_str}\", style=\"dim\"))\n\n # Add client_tool_result info\n if event.client_tool_result and event.client_tool_result.message:\n result_content = None\n tool_call_id = None\n\n # Get tool_call_id\n if hasattr(event.client_tool_result.message, \"tool_call_id\"):\n tool_call_id = event.client_tool_result.message.tool_call_id\n\n # Get user input content\n raw_content = event.client_tool_result.message.content or \"\"\n result_content = _content_to_string(raw_content)\n\n # Show tool call result\n lines.append(Text(\"\\n 📥 Tool Call Result:\", style=\"bold magenta\"))\n if tool_call_id:\n lines.append(Text(f\" Tool Call ID: {tool_call_id}\", style=\"dim\"))\n\n # If content exists, display with markdown\n if result_content:\n lines.append(Text(\"\\n User Input:\", style=\"bold magenta\"))\n # Add user input content to lines, will be rendered with Markdown later\n lines.append((\"markdown_result\", result_content))\n\n # Add error info\n if event.error:\n lines.append(Text(f\"\\n Error: {event.error}\", style=\"red\"))\n\n # Add completion signal info\n if event.type == AgentEventType.SIGNAL and event.metadata:\n status = event.metadata.get(\"status\", \"\")\n message = event.metadata.get(\"message\", \"\")\n if status == \"completed\":\n lines.append(Text(f\"\\n ✅ {message}\", style=\"green\"))\n\n # Check if there's markdown content (context, prompt or result)\n markdown_context = None\n markdown_prompt = None\n markdown_result = None\n text_lines = []\n for item in lines:\n if isinstance(item, tuple) and len(item) == 2:\n if item[0] == \"markdown_context\":\n markdown_context = item[1]\n elif item[0] == \"markdown_prompt\":\n markdown_prompt = item[1]\n elif item[0] == \"markdown_result\":\n markdown_result = item[1]\n else:\n text_lines.append(item)\n else:\n text_lines.append(item)\n\n # If there's markdown content, combine Text and Markdown\n markdown_parts = []\n if markdown_context:\n markdown_parts.append(Markdown(markdown_context))\n if markdown_prompt:\n markdown_parts.append(Markdown(markdown_prompt))\n if markdown_result:\n markdown_parts.append(Markdown(markdown_result))\n\n if markdown_parts:\n header = Text.assemble(*text_lines)\n # If there are multiple markdown contents, combine with Group\n if len(markdown_parts) == 1:\n return Panel(\n Group(header, markdown_parts[0]),\n border_style=color,\n title=f\"Event #{self.index}\",\n )\n else:\n return Panel(\n Group(header, *markdown_parts),\n border_style=color,\n title=f\"Event #{self.index}\",\n )\n else:\n return Panel(\n Text.assemble(*text_lines),\n border_style=color,\n title=f\"Event #{self.index}\",\n )\n\n\nclass CommandItem(ListItem):\n \"\"\"Command list item\"\"\"\n\n def __init__(self, command: str):\n self.command = command\n super().__init__()\n\n def render(self) -> RenderableType:\n \"\"\"Render command item\"\"\"\n return Text(f\"/{self.command}\", style=\"cyan\")\n\n\nclass AgentItem(ListItem):\n \"\"\"Agent list item\"\"\"\n\n def __init__(self, agent_config: AgentConfig):\n self.agent_config = agent_config\n super().__init__()\n\n def render(self) -> RenderableType:\n \"\"\"Render Agent item\"\"\"\n name = self.agent_config.name or \"N/A\"\n description = self.agent_config.description or \"No description\"\n agent_type = self.agent_config.agent_type or \"N/A\"\n\n lines = [\n Text(f\"Name: {name}\", style=\"bold cyan\"),\n Text(f\" Type: {agent_type}\", style=\"dim\"),\n Text(f\" Description: {description}\", style=\"dim\"),\n ]\n\n return Panel(\n Text.assemble(*lines),\n border_style=\"blue\",\n title=f\"Agent: {name}\",\n )\n\n\nclass ProcessView(Container):\n \"\"\"Process view - displays event list\"\"\"\n\n def __init__(self, *args, **kwargs):\n super().__init__(*args, **kwargs)\n self.events: list[AgentEvent] = []\n self.list_view = ListView(id=\"event-list\")\n self.json_view = Static(\"\", id=\"process-json-view\")\n\n def compose(self) -> ComposeResult:\n with Horizontal():\n with Container(id=\"process-list-container\"):\n yield Label(\"Agent Event Stream\", id=\"process-title\")\n yield self.list_view\n with VerticalScroll(id=\"process-json-container\"):\n yield Label(\"Event Details (JSON)\", id=\"process-json-title\")\n yield self.json_view\n\n def add_event(self, event: AgentEvent) -> None:\n \"\"\"Add event to list\"\"\"\n self.events.append(event)\n index = len(self.events)\n item = EventItem(event, index)\n self.list_view.append(item)\n # Auto scroll to bottom\n self.list_view.scroll_end(animate=False)\n\n @on(ListView.Selected, \"#event-list\")\n def on_list_item_selected(self, event: ListView.Selected) -> None:\n \"\"\"Handle list item selection\"\"\"\n if hasattr(event, \"item\") and isinstance(event.item, EventItem):\n event_obj = event.item.event\n json_text = event_obj.model_dump_json(indent=2, ensure_ascii=False)\n self.json_view.update(\n Panel(json_text, title=\"Event JSON\", border_style=\"blue\")\n )\n # Notify main app to update placeholder\n app = self.app\n if app and hasattr(app, \"_update_placeholder_by_focus\"):\n app._update_placeholder_by_focus()\n\n def clear_events(self) -> None:\n \"\"\"Clear event list\"\"\"\n self.events.clear()\n self.list_view.clear()\n\n\nclass AgentsListView(Container):\n \"\"\"Agents list view - displays registered Agents\"\"\"\n\n def __init__(self, *args, **kwargs):\n super().__init__(*args, **kwargs)\n self.list_view = ListView(id=\"agents-list\")\n self.json_view = Static(\"\", id=\"agents-json-view\")\n self.selection_callback: Optional[Callable[[str], None]] = None\n\n def set_selection_callback(self, callback: Callable[[str], None]) -> None:\n \"\"\"Set selection callback function\"\"\"\n self.selection_callback = callback\n\n def compose(self) -> ComposeResult:\n with Horizontal():\n with Container(id=\"agents-list-container\"):\n yield Label(\"Registered Agents\", id=\"agents-title\")\n yield self.list_view\n with VerticalScroll(id=\"agents-json-container\"):\n yield Label(\"Agent Details (JSON)\", id=\"agents-json-title\")\n yield self.json_view\n\n def update_agents(self, agents: list[AgentConfig]) -> None:\n \"\"\"Update Agents list\"\"\"\n self.list_view.clear()\n for agent_config in agents:\n item = AgentItem(agent_config)\n self.list_view.append(item)\n\n @on(ListView.Selected, \"#agents-list\")\n def on_list_item_selected(self, event: ListView.Selected) -> None:\n \"\"\"Handle list item selection\"\"\"\n if hasattr(event, \"item\") and isinstance(event.item, AgentItem):\n agent_config = event.item.agent_config\n json_text = agent_config.model_dump_json(indent=2, ensure_ascii=False)\n self.json_view.update(\n Panel(json_text, title=\"Agent Config JSON\", border_style=\"green\")\n )\n # Call callback to notify main app\n if self.selection_callback:\n self.selection_callback(agent_config.name)\n # Notify main app to update placeholder\n app = self.app\n if app and hasattr(app, \"_update_placeholder_by_focus\"):\n app._update_placeholder_by_focus()\n\n\nclass TaskItem(ListItem):\n \"\"\"Task list item\"\"\"\n\n def __init__(self, task_id: str, request_data: dict):\n self.task_id = task_id\n self.request_data = request_data\n super().__init__()\n\n def render(self) -> RenderableType:\n \"\"\"Render Task item\"\"\"\n task_id_short = (\n self.task_id[:16] + \"...\" if len(self.task_id) > 16 else self.task_id\n )\n messages = self.request_data.get(\"messages\", [])\n content_preview = \"\"\n if messages and len(messages) > 0:\n first_msg = messages[0] if isinstance(messages, list) else messages\n if isinstance(first_msg, dict):\n content = first_msg.get(\"content\", \"\")\n else:\n content = getattr(first_msg, \"content\", \"\")\n content_str = _content_to_string(content)\n content_preview = (\n content_str[:50] + \"...\" if len(content_str) > 50 else content_str\n )\n\n lines = [\n Text(f\"Task ID: {task_id_short}\", style=\"bold cyan\"),\n Text(f\" Content: {content_preview}\", style=\"dim\"),\n ]\n\n return Panel(\n Text.assemble(*lines),\n border_style=\"green\",\n title=f\"Task: {task_id_short}\",\n )\n\n\nclass TasksListView(Container):\n \"\"\"Tasks list view - displays saved Tasks\"\"\"\n\n def __init__(self, *args, **kwargs):\n super().__init__(*args, **kwargs)\n self.list_view = ListView(id=\"tasks-list\")\n self.json_view = Static(\"\", id=\"tasks-json-view\")\n self.selection_callback: Optional[Callable[[str], None]] = None\n\n def set_selection_callback(self, callback: Callable[[str], None]) -> None:\n \"\"\"Set selection callback function\"\"\"\n self.selection_callback = callback\n\n def compose(self) -> ComposeResult:\n with Horizontal():\n with Container(id=\"tasks-list-container\"):\n yield Label(\"Saved Tasks\", id=\"tasks-title\")\n yield self.list_view\n with VerticalScroll(id=\"tasks-json-container\"):\n yield Label(\"Task Details (JSON)\", id=\"tasks-json-title\")\n yield self.json_view\n\n def update_tasks(self, tasks: list[tuple[str, dict]]) -> None:\n \"\"\"Update Tasks list\n\n Args:\n tasks: [(task_id, request_data), ...] list\n \"\"\"\n self.list_view.clear()\n for task_id, request_data in tasks:\n item = TaskItem(task_id, request_data)\n self.list_view.append(item)\n\n @on(ListView.Selected, \"#tasks-list\")\n def on_list_item_selected(self, event: ListView.Selected) -> None:\n \"\"\"Handle list item selection\"\"\"\n if hasattr(event, \"item\") and isinstance(event.item, TaskItem):\n task_id = event.item.task_id\n request_data = event.item.request_data\n json_text = json.dumps(request_data, indent=2, ensure_ascii=False)\n self.json_view.update(\n Panel(json_text, title=\"Request JSON\", border_style=\"green\")\n )\n # Call callback to notify main app\n if self.selection_callback:\n self.selection_callback(task_id)\n\n\nclass CommandSelector(Container):\n \"\"\"Command selector - displays command list and supports up/down key selection\"\"\"\n\n def __init__(self, *args, **kwargs):\n super().__init__(*args, **kwargs)\n self.command_list_view = ListView(id=\"command-list\")\n self.matched_commands: list[str] = []\n\n def compose(self) -> ComposeResult:\n yield self.command_list_view\n\n def update_commands(self, commands: list[str]) -> None:\n \"\"\"Update command list\"\"\"\n self.matched_commands = commands\n self.command_list_view.clear()\n for cmd in commands:\n item = CommandItem(cmd)\n self.command_list_view.append(item)\n # If there are commands, select the first one\n if commands:\n self.command_list_view.index = 0\n\n def get_selected_command(self) -> Optional[str]:\n \"\"\"Get currently selected command\"\"\"\n if not self.matched_commands:\n return None\n # Get current selected index\n try:\n index = self.command_list_view.index\n if 0 <= index < len(self.matched_commands):\n return self.matched_commands[index]\n except Exception:\n pass\n # If no selected item, return the first one\n if self.matched_commands:\n return self.matched_commands[0]\n return None\n\n\nclass AgentTUIApp(App):\n \"\"\"Agent TUI Application\"\"\"\n\n CSS = \"\"\"\n #process-title, #process-json-title, #agents-title, #agents-json-title {\n text-style: dim;\n padding: 0 1;\n margin-bottom: 1;\n height: 1;\n text-align: left;\n }\n\n #process-view, #agents-view {\n height: 1fr;\n }\n \n #process-list-container, #agents-list-container {\n width: 50%;\n border-right: wide $primary;\n height: 1fr;\n }\n \n #process-json-container, #agents-json-container {\n width: 50%;\n height: 1fr;\n }\n \n #process-json-view, #agents-json-view {\n padding: 1;\n }\n \n #event-list, #agents-list {\n height: 1fr;\n }\n \n #input-container {\n height: auto;\n min-height: 3;\n border-top: wide $primary;\n margin-top: 1;\n padding-top: 1;\n }\n \n #selected-agent-label {\n height: 1;\n padding: 0 1;\n margin-bottom: 1;\n text-style: bold;\n background: $surface;\n }\n \n #input-box {\n width: 1fr;\n }\n \n #command-selector {\n height: auto;\n max-height: 10;\n border-top: wide $primary;\n background: $surface;\n margin-top: 1;\n margin-bottom: 1;\n }\n \n #command-list {\n height: auto;\n max-height: 10;\n }\n \n .hidden {\n display: none;\n }\n \n #main-container {\n height: 1fr;\n overflow: hidden;\n }\n \n Screen {\n layout: vertical;\n }\n \"\"\"\n\n BINDINGS = [\n Binding(\"q\", \"quit\", \"Quit\"),\n Binding(\"ctrl+c\", \"quit\", \"Quit\"),\n ]\n\n def __init__(\n self,\n orchestrator: Orchestrator,\n workdir: Optional[str | Path] = None,\n tracer: Optional[Tracer] = None,\n ):\n super().__init__()\n self.process_view: Optional[ProcessView] = None\n self.agents_view: Optional[AgentsListView] = None\n self.tasks_view: Optional[TasksListView] = None\n self.command_selector: Optional[CommandSelector] = None\n self.input_box: Optional[Input] = None\n self.selected_agent_label: Optional[Label] = None\n self.orchestrator: Orchestrator = orchestrator\n self.running_task: Optional[asyncio.Task] = None\n self.commands = [\n \"quit\",\n \"clear\",\n \"help\",\n \"ls agents\",\n \"ls tasks\",\n \"view process\",\n ]\n self.show_command_selector = False\n self.current_view = (\n \"agents\" # \"process\" or \"agents\" or \"tasks\", default to agents view\n )\n self.selected_agent_name: Optional[str] = None # Currently selected agent name\n self.events: list[AgentEvent] = [] # Store all events in memory\n self.default_placeholder = \"/command or input...\" # Default placeholder\n self.workdir: Optional[Path] = Path(workdir) if workdir else None\n if self.workdir:\n self.workdir.mkdir(parents=True, exist_ok=True)\n self.pending_ask_input_event: Optional[AgentEvent] = (\n None # Pending ask_input event\n )\n self.tracer = tracer\n\n def compose(self) -> ComposeResult:\n \"\"\"Compose application interface\"\"\"\n with Vertical():\n with Container(id=\"main-container\"):\n process_view = ProcessView(id=\"process-view\")\n process_view.add_class(\"hidden\")\n yield process_view\n agents_view = AgentsListView(id=\"agents-view\")\n yield agents_view\n tasks_view = TasksListView(id=\"tasks-view\")\n tasks_view.add_class(\"hidden\")\n yield tasks_view\n command_selector = CommandSelector(id=\"command-selector\")\n command_selector.add_class(\"hidden\")\n yield command_selector\n with Container(id=\"input-container\"):\n yield Label(\"Current Agent: Not selected\", id=\"selected-agent-label\")\n yield Input(placeholder=self.default_placeholder, id=\"input-box\")\n\n def on_mount(self) -> None:\n \"\"\"Initialize on app mount\"\"\"\n self.process_view = self.query_one(\"#process-view\", ProcessView)\n self.agents_view = self.query_one(\"#agents-view\", AgentsListView)\n self.tasks_view = self.query_one(\"#tasks-view\", TasksListView)\n self.command_selector = self.query_one(\"#command-selector\", CommandSelector)\n self.input_box = self.query_one(\"#input-box\", Input)\n self.selected_agent_label = self.query_one(\"#selected-agent-label\", Label)\n\n # Set callbacks\n self.agents_view.set_selection_callback(self._on_agent_selected)\n self.tasks_view.set_selection_callback(self._on_task_selected)\n\n # Initialize agents list (default to agents view)\n agents = self.orchestrator.list_agents()\n self.agents_view.update_agents(agents)\n\n # Select first agent by default\n if agents:\n first_agent = agents[0]\n if first_agent.name:\n self._on_agent_selected(first_agent.name)\n # Select first item in list\n self.agents_view.list_view.index = 0\n\n # Focus input box\n self._focus_input()\n\n @on(Input.Changed, \"#input-box\")\n def on_input_changed(self, event: Input.Changed) -> None:\n \"\"\"Handle input changes\"\"\"\n value = event.value\n if value == \"/\":\n # Show all commands\n self.show_command_selector = True\n self.command_selector.update_commands(self.commands)\n self.command_selector.remove_class(\"hidden\")\n # Keep input focus, user can continue typing to filter\n elif value.startswith(\"/\"):\n # Filter matching commands\n query = value[1:].lower()\n matched = [cmd for cmd in self.commands if cmd.startswith(query)]\n if matched:\n self.show_command_selector = True\n self.command_selector.update_commands(matched)\n self.command_selector.remove_class(\"hidden\")\n else:\n # No matching commands, hide selector\n self.show_command_selector = False\n self.command_selector.add_class(\"hidden\")\n else:\n # Hide command selector\n self.show_command_selector = False\n self.command_selector.add_class(\"hidden\")\n\n @on(ListView.Selected, \"#command-list\")\n def on_command_selected(self, event: ListView.Selected) -> None:\n \"\"\"Handle command selection\"\"\"\n if hasattr(event, \"item\") and isinstance(event.item, CommandItem):\n command = event.item.command\n self._handle_command(command)\n self.input_box.value = \"\"\n self.command_selector.add_class(\"hidden\")\n self.show_command_selector = False\n self._focus_input()\n\n def _focus_input(self) -> None:\n \"\"\"Focus input box and update placeholder\"\"\"\n self.set_focus(self.input_box)\n if self.input_box:\n self.input_box.placeholder = self.default_placeholder\n\n def _update_placeholder_by_focus(self) -> None:\n \"\"\"Update placeholder based on focus state\"\"\"\n if not self.input_box:\n return\n focused = self.screen.focused\n if focused == self.input_box:\n self.input_box.placeholder = self.default_placeholder\n else:\n self.input_box.placeholder = \"ESC to return to input box\"\n\n @on(Key)\n def on_key(self, event: Key) -> None:\n \"\"\"Handle keyboard events\"\"\"\n # Global ESC handling: if focus is not on input box, press ESC to return to input\n if event.key == \"escape\":\n focused = self.screen.focused\n if focused != self.input_box:\n self._focus_input()\n event.prevent_default()\n return\n\n # If command selector is shown\n if self.show_command_selector and not self.command_selector.has_class(\"hidden\"):\n focused = self.screen.focused\n\n # If focus is on input box, transfer to command list on up/down keys\n if focused == self.input_box:\n if event.key == \"up\" or event.key == \"down\":\n # Transfer focus to command list\n self.set_focus(self.command_selector.command_list_view)\n self._update_placeholder_by_focus()\n event.prevent_default()\n return\n\n # If focus is on command list\n elif focused == self.command_selector.command_list_view:\n # Press enter in command list to execute selected command\n if event.key == \"enter\":\n selected = self.command_selector.get_selected_command()\n if selected:\n self._handle_command(selected)\n self.input_box.value = \"\"\n self.command_selector.add_class(\"hidden\")\n self.show_command_selector = False\n self._focus_input()\n event.prevent_default()\n return\n # Press Esc to return to input box\n elif event.key == \"escape\":\n self.command_selector.add_class(\"hidden\")\n self.show_command_selector = False\n self._focus_input()\n event.prevent_default()\n return\n\n @on(Input.Submitted, \"#input-box\")\n def on_input_submitted(self, event: Input.Submitted) -> None:\n \"\"\"Handle input submission\"\"\"\n value = event.value.strip()\n if not value:\n return\n\n # Check if there's a pending ask_input event\n if self.pending_ask_input_event:\n # Send user input as tool_call result\n self._send_ask_input_result(value)\n self.input_box.value = \"\"\n self.pending_ask_input_event = None\n # Restore default placeholder\n if self.input_box:\n self.input_box.placeholder = self.default_placeholder\n return\n\n # Handle commands\n if value.startswith(\"/\"):\n # If command selector is shown, use selected command\n if self.show_command_selector and not self.command_selector.has_class(\n \"hidden\"\n ):\n selected = self.command_selector.get_selected_command()\n if selected:\n command = selected\n else:\n command = value[1:].lower()\n else:\n command = value[1:].lower()\n\n self._handle_command(command)\n self.input_box.value = \"\"\n self.command_selector.add_class(\"hidden\")\n self.show_command_selector = False\n return\n\n # Execute user instruction\n self._run_agent(value)\n self.input_box.value = \"\"\n return\n\n def _handle_command(self, command: str) -> None:\n \"\"\"Handle command\"\"\"\n if command == \"quit\" or command == \"q\":\n self.exit()\n elif command == \"clear\":\n self.events.clear()\n self.process_view.clear_events()\n elif command == \"help\":\n # Show help information\n help_text = \"Available commands:\\n\"\n for cmd in self.commands:\n help_text += f\" /{cmd}\\n\"\n self.notify(help_text, title=\"Help\", timeout=5)\n elif command == \"ls agents\" or command == \"ls\":\n # Switch to agents view\n self._switch_view(\"agents\")\n elif command == \"ls tasks\":\n # Switch to tasks view\n self._switch_view(\"tasks\")\n elif command == \"view process\" or command == \"process\":\n # Switch to process view\n self._switch_view(\"process\")\n else:\n # Try fuzzy match\n if \"quit\" in command or \"q\" in command:\n self.exit()\n elif \"clear\" in command:\n self.events.clear()\n self.process_view.clear_events()\n elif \"help\" in command:\n help_text = \"Available commands:\\n\"\n for cmd in self.commands:\n help_text += f\" /{cmd}\\n\"\n self.notify(help_text, title=\"Help\", timeout=5)\n elif \"ls\" in command and \"agent\" in command:\n self._switch_view(\"agents\")\n elif \"ls\" in command and \"task\" in command:\n self._switch_view(\"tasks\")\n elif (\"view\" in command and \"process\" in command) or command == \"process\":\n self._switch_view(\"process\")\n\n def _switch_view(self, view_name: str) -> None:\n \"\"\"Switch view\"\"\"\n self.current_view = view_name\n if view_name == \"agents\":\n # Show agents view\n self.process_view.add_class(\"hidden\")\n self.tasks_view.add_class(\"hidden\")\n self.agents_view.remove_class(\"hidden\")\n # Update agents list\n agents = self.orchestrator.list_agents()\n self.agents_view.update_agents(agents)\n elif view_name == \"tasks\":\n # Show tasks view\n self.process_view.add_class(\"hidden\")\n self.agents_view.add_class(\"hidden\")\n self.tasks_view.remove_class(\"hidden\")\n # Update tasks list\n self._load_tasks()\n # Auto focus to task list\n self.set_focus(self.tasks_view.list_view)\n self._update_placeholder_by_focus()\n else:\n # Show process view\n self.process_view.remove_class(\"hidden\")\n self.agents_view.add_class(\"hidden\")\n self.tasks_view.add_class(\"hidden\")\n\n def _on_agent_selected(self, agent_name: str) -> None:\n \"\"\"Handle Agent selection\"\"\"\n self.selected_agent_name = agent_name\n if self.selected_agent_label:\n self.selected_agent_label.update(f\"Current Agent: {agent_name}\")\n\n def _send_ask_input_result(self, user_input: str) -> None:\n \"\"\"Send ask_input result\"\"\"\n if not self.pending_ask_input_event or not self.orchestrator:\n return\n\n ask_input_event = self.pending_ask_input_event\n\n # Get tool_call_id from ask_input_event\n tool_call_id = None\n if ask_input_event.client_tool_call:\n tool_call_id = ask_input_event.client_tool_call.tool_call_id\n\n # Create AgentResponse as tool_call result\n # role should be \"tool\" instead of \"user\", because this is the result of a tool call\n result_response = AgentResponse(\n agent_name=ask_input_event.agent_name,\n message=ChatMessage(\n role=\"tool\", content=user_input, tool_call_id=tool_call_id\n ),\n status=AgentRunningStatus.FINISHED,\n )\n\n # Create CLIENT_TOOL_RESULT event\n result_event = AgentEvent(\n task_id=ask_input_event.task_id,\n parent_task_id=ask_input_event.parent_task_id,\n root_task_id=ask_input_event.root_task_id,\n type=AgentEventType.CLIENT_TOOL_RESULT,\n client_tool_result=result_response,\n )\n\n # Send event through orchestrator\n try:\n # Use send_event method to send event\n asyncio.create_task(self.orchestrator.send_event(result_event))\n\n # Save to memory and view\n self.events.append(result_event)\n self.process_view.add_event(result_event)\n\n # Save to file\n if ask_input_event.root_task_id and self.workdir:\n self._save_event(ask_input_event.root_task_id, result_event)\n except Exception as e:\n logger.error(\"Failed to send ask_input result: %s\", e)\n self.notify(f\"Failed to send ask_input result: {e}\", title=\"Error\", timeout=3)\n\n def _on_task_selected(self, task_id: str) -> None:\n \"\"\"Handle Task selection\"\"\"\n # Switch to process view\n self._switch_view(\"process\")\n # Load all events for this task\n self._load_task_events(task_id)\n\n def _load_tasks(self) -> None:\n \"\"\"Load all tasks\"\"\"\n if not self.workdir:\n return\n\n tasks = []\n # Find all {task_id}_request.json files\n for request_file in self.workdir.glob(\"*_request.json\"):\n task_id = request_file.stem.replace(\"_request\", \"\")\n try:\n with open(request_file, \"r\", encoding=\"utf-8\") as f:\n request_data = json.load(f)\n # Get file modification time\n mtime = request_file.stat().st_mtime\n tasks.append((task_id, request_data, mtime))\n except Exception as e:\n logger.error(f\"Failed to load task {task_id}: {e}\")\n\n # Sort by file modification time (newest first)\n tasks.sort(key=lambda x: x[2], reverse=True)\n # Remove mtime, keep only (task_id, request_data)\n tasks = [(task_id, request_data) for task_id, request_data, _ in tasks]\n self.tasks_view.update_tasks(tasks)\n\n def _load_task_events(self, task_id: str) -> None:\n \"\"\"Load all events for a task from file\"\"\"\n if not self.workdir:\n return\n\n jsonl_file = self.workdir / f\"{task_id}.jsonl\"\n if not jsonl_file.exists():\n self.notify(f\"Event file not found for task {task_id}\", title=\"Error\", timeout=3)\n return\n\n # Clear current events\n self.events.clear()\n self.process_view.clear_events()\n\n # Load events from jsonl file\n try:\n with open(jsonl_file, \"r\", encoding=\"utf-8\") as f:\n for line in f:\n strip_line = line.strip()\n if not strip_line:\n continue\n try:\n event_data = json.loads(strip_line)\n event = AgentEvent(**event_data)\n self.events.append(event)\n self.process_view.add_event(event)\n except Exception as e:\n logger.error(f\"Failed to parse event: {e}\")\n except Exception as e:\n logger.error(f\"Failed to load task events: {e}\")\n self.notify(f\"Failed to load task events: {e}\", title=\"Error\", timeout=3)\n\n def _save_request(self, root_task_id: str, request: AgentRequest) -> None:\n \"\"\"Save request to file\"\"\"\n if not self.workdir:\n return\n\n request_file = self.workdir / f\"{root_task_id}_request.json\"\n try:\n request_data = request.model_dump()\n with open(request_file, \"w\", encoding=\"utf-8\") as f:\n json.dump(request_data, f, indent=2, ensure_ascii=False)\n except Exception as e:\n logger.error(f\"Failed to save request: {e}\")\n\n def _save_event(self, root_task_id: str, event: AgentEvent) -> None:\n \"\"\"Save event to jsonl file\"\"\"\n if not self.workdir:\n return\n\n jsonl_file = self.workdir / f\"{root_task_id}.jsonl\"\n try:\n event_data = event.model_dump()\n with open(jsonl_file, \"a\", encoding=\"utf-8\") as f:\n f.write(json.dumps(event_data, ensure_ascii=False) + \"\\n\")\n except Exception as e:\n logger.error(f\"Failed to save event: {e}\")\n\n @work(exclusive=True)\n async def _run_agent(self, user_input: str) -> None:\n \"\"\"Run Agent\"\"\"\n if not self.orchestrator:\n return\n\n # Check if agent is selected\n if not self.selected_agent_name:\n self.notify(\"Please select an Agent first (use /ls agents)\", title=\"Error\", timeout=3)\n return\n\n # Switch to process view to show execution\n self._switch_view(\"process\")\n\n # Clear previous events\n self.events.clear()\n self.process_view.clear_events()\n\n messages = [ChatMessage(role=\"user\", content=user_input)]\n request = AgentRequest(\n agent_name=self.selected_agent_name,\n messages=messages,\n )\n event = AgentEvent(\n type=AgentEventType.REQUEST,\n request=request,\n )\n\n # Get root_task_id (from the first event)\n root_task_id = None\n ctx = SpanContext(tracer=self.tracer, app_name=self.selected_agent_name)\n with ctx.span(name=f\"{self.selected_agent_name}_run_agent\"):\n try:\n async for agent_event in self.orchestrator.run(\n agent_name=self.selected_agent_name,\n event=event,\n agent_config=None,\n ):\n # Get root_task_id (from the first event)\n if root_task_id is None:\n root_task_id = agent_event.root_task_id or agent_event.task_id\n # Save request to file\n if root_task_id and self.workdir:\n self._save_request(root_task_id, request)\n\n # Check if it's ask_input type client_tool_call\n if (\n agent_event.client_tool_call\n and getattr(agent_event.client_tool_call, \"type\", None)\n == ClientToolCallType.ASK_INPUT\n ):\n # Save as pending ask_input event\n self.pending_ask_input_event = agent_event\n # Update placeholder to prompt user input\n if self.input_box:\n self.input_box.placeholder = (\n \"Please enter content (as ask_input response)...\"\n )\n\n # Save to memory\n self.events.append(agent_event)\n # Add to view\n self.process_view.add_event(agent_event)\n\n # Save to file\n if root_task_id and self.workdir:\n self._save_event(root_task_id, agent_event)\n\n # After execution completes, add completion event\n completion_event = AgentEvent(\n type=AgentEventType.SIGNAL,\n agent_name=self.selected_agent_name,\n metadata={\"status\": \"completed\", \"message\": \"Execution completed\"},\n )\n self.events.append(completion_event)\n self.process_view.add_event(completion_event)\n\n # Save completion event to file\n if root_task_id and self.workdir:\n self._save_event(root_task_id, completion_event)\n except Exception as e:\n logger.error(\"Error running Agent: %s\", e, exc_info=True)\n error_event = AgentEvent(\n type=AgentEventType.ERROR,\n error=str(e),\n )\n # Save to memory\n self.events.append(error_event)\n # Add to view\n self.process_view.add_event(error_event)\n\n async def action_quit(self) -> None:\n \"\"\"Quit application\"\"\"\n self.exit()\n" }, { "path": "cortex/utils/__init__.py", "content": "from .generator_merger import GeneratorMerger\n\n__all__ = [\"GeneratorMerger\"]\n" }, { "path": "cortex/utils/generator_merger.py", "content": "\"\"\"Generator merger implementation using asyncio.\"\"\"\n\nimport asyncio\nfrom collections.abc import Generator\nfrom typing import Any, AsyncGenerator, Awaitable, Callable, Optional\n\n\nclass GeneratorMerger:\n \"\"\"A merger that combines multiple generator functions using asyncio for concurrent execution.\"\"\"\n\n def __init__(\n self,\n on_generator_complete: Optional[\n Callable[[str, str, Optional[Exception]], Awaitable[None]]\n ] = None,\n ):\n \"\"\"\n Initialize the merger.\n\n Args:\n on_generator_complete: Optional callback function called when a sub-generator completes.\n Parameters: (generator_id, generator_type, error)\n If not provided, completion events will be yielded.\n \"\"\"\n self._generators: dict[str, Callable[[], Generator[Any, None, None]]] = {}\n self._async_generators: dict[str, Callable[[], AsyncGenerator[Any, None]]] = {}\n self._running_tasks: list[asyncio.Task] = []\n self._queue: asyncio.Queue = asyncio.Queue()\n self._active_count: int = 0\n self._lock: asyncio.Lock = asyncio.Lock()\n self._processed_generators: set[str] = set() # Processed generator IDs\n self._on_generator_complete = on_generator_complete\n self._generator_id_counter: int = 0 # Counter for generating unique generator IDs\n\n def add_generator(\n self,\n generator_func: Callable[[], Generator[Any, None, None]],\n generator_id: Optional[str] = None,\n ):\n \"\"\"\n Dynamically add a synchronous generator function.\n\n Note: This method is synchronous and does not use locks when adding.\n Lock protection is primarily used in the async merge() and delete operations.\n\n Args:\n generator_func: A function that returns a synchronous generator\n generator_id: Optional generator identifier used in completion events\n \"\"\"\n if generator_id is None:\n generator_id = f\"sync_gen_{len(self._generators)}\"\n\n if generator_id in self._generators or generator_id in self._async_generators:\n raise ValueError(f\"Generator {generator_id} already exists\")\n\n self._generators[generator_id] = generator_func\n\n def add_async_generator(\n self,\n async_generator_func: Callable[[], AsyncGenerator[Any, None]],\n generator_id: Optional[str] = None,\n ):\n \"\"\"\n Dynamically add an async generator function.\n\n Args:\n async_generator_func: A function that returns an async generator\n generator_id: Optional generator identifier used in completion events\n \"\"\"\n if generator_id is None:\n generator_id = f\"async_gen_{len(self._async_generators)}\"\n\n # Check if generator_id already exists\n if generator_id in self._generators or generator_id in self._async_generators:\n raise ValueError(f\"Generator {generator_id} already exists\")\n\n self._async_generators[generator_id] = async_generator_func\n\n def _get_next_item(self, generator: Generator[Any, None, None]) -> tuple[bool, Any]:\n \"\"\"\n Get the next value from a generator (runs in thread).\n\n Args:\n generator: generator object\n\n Returns:\n (has_more, item) tuple, where has_more indicates if there are more values\n \"\"\"\n try:\n item = next(generator)\n return (True, item)\n except StopIteration:\n return (False, None)\n\n async def _run_generator_with_wrapper(\n self,\n generator_id: str,\n generator_type: str,\n generator_executor: Callable[[], AsyncGenerator[Any, None]],\n ):\n \"\"\"\n Generic generator execution wrapper.\n\n Args:\n generator_id: generator identifier\n generator_type: generator type (\"sync\" or \"async\")\n generator_executor: async function that executes the generator and yields values\n \"\"\"\n async with self._lock:\n self._active_count += 1\n\n error = None\n\n try:\n # Execute specific generator logic\n async for item in generator_executor():\n # Put generator-produced values into queue\n await self._queue.put(item)\n\n except Exception as e:\n error = e\n # If generator errors, also put into queue for processing\n await self._queue.put((\"__error__\", e))\n finally:\n async with self._lock:\n self._active_count -= 1\n\n # Notify generator completion\n await self._notify_generator_complete(generator_id, generator_type, error)\n\n async def _run_generator(\n self,\n generator_id: str,\n generator_func: Callable[[], Generator[Any, None, None]],\n ):\n \"\"\"\n Run a synchronous generator in the event loop.\n\n Args:\n generator_id: generator identifier\n generator_func: function that returns a synchronous generator\n \"\"\"\n\n async def sync_generator_executor():\n # Create generator (this step is fast, doesn't need to run in thread)\n generator = generator_func()\n loop = asyncio.get_event_loop()\n\n # Get values one by one, run in thread pool to avoid blocking event loop\n while True:\n has_more, item = await loop.run_in_executor(\n None, self._get_next_item, generator\n )\n\n if not has_more:\n break\n\n yield item\n\n await self._run_generator_with_wrapper(\n generator_id, \"sync\", sync_generator_executor\n )\n\n async def _run_async_generator(\n self,\n generator_id: str,\n async_generator_func: Callable[[], AsyncGenerator[Any, None]],\n ):\n \"\"\"\n Run an async generator in the event loop.\n\n Args:\n generator_id: generator identifier\n async_generator_func: function that returns an async generator\n \"\"\"\n\n async def async_generator_executor():\n # Create async generator\n async_generator = async_generator_func()\n\n # Get values directly from async generator\n async for item in async_generator:\n yield item\n\n await self._run_generator_with_wrapper(\n generator_id, \"async\", async_generator_executor\n )\n\n async def _notify_generator_complete(\n self, generator_id: str, generator_type: str, error: Optional[Exception]\n ):\n \"\"\"\n Notify generator completion.\n\n Args:\n generator_id: Generator identifier\n generator_type: Generator type (\"sync\" or \"async\")\n error: Error object if any; otherwise None\n \"\"\"\n # Use lock to protect dictionary modification\n async with self._lock:\n # Mark generator as completed\n self._processed_generators.add(generator_id)\n\n # Auto-delete completed generator\n if generator_type == \"sync\" and generator_id in self._generators:\n del self._generators[generator_id]\n elif generator_type == \"async\" and generator_id in self._async_generators:\n del self._async_generators[generator_id]\n\n event = {\n \"type\": \"generator_complete\",\n \"generator_id\": generator_id,\n \"generator_type\": generator_type,\n \"status\": \"error\" if error else \"completed\",\n \"error\": str(error) if error else None,\n }\n\n if self._on_generator_complete:\n # If there's a callback, call it\n if asyncio.iscoroutinefunction(self._on_generator_complete):\n await self._on_generator_complete(generator_id, generator_type, error)\n else:\n self._on_generator_complete(generator_id, generator_type, error)\n else:\n # If no callback, put event into queue, will be yielded\n await self._queue.put((\"__event__\", event))\n\n async def merge(self) -> AsyncGenerator[Any, None]:\n \"\"\"\n Merge all added generators (sync and async), returning an async generator.\n Supports dynamically adding new generators during iteration.\n\n Yields:\n Values produced by the various generators\n \"\"\"\n if not self._generators and not self._async_generators:\n return\n\n # Reset state\n self._queue = asyncio.Queue()\n self._active_count = 0\n self._running_tasks = []\n self._processed_generators = set()\n\n # Use lock to get initial generator snapshot\n async with self._lock:\n sync_generators_snapshot = list(self._generators.items())\n async_generators_snapshot = list(self._async_generators.items())\n\n # Create tasks for initial sync generators\n for generator_id, generator_func in sync_generators_snapshot:\n task = asyncio.create_task(\n self._run_generator(generator_id, generator_func),\n name=f\"sync_{generator_id}\",\n )\n self._running_tasks.append(task)\n\n # Create tasks for initial async generators\n for generator_id, async_generator_func in async_generators_snapshot:\n task = asyncio.create_task(\n self._run_async_generator(generator_id, async_generator_func),\n name=f\"async_{generator_id}\",\n )\n self._running_tasks.append(task)\n\n # Get values from queue and yield\n while True:\n # Use lock to get current generator snapshot\n async with self._lock:\n sync_generators_snapshot = list(self._generators.items())\n async_generators_snapshot = list(self._async_generators.items())\n processed_generators_snapshot = self._processed_generators.copy()\n\n # Check for newly added sync generators\n for generator_id, generator_func in sync_generators_snapshot:\n if generator_id not in processed_generators_snapshot:\n # Check if a task is already running this generator\n task_exists = any(\n task.get_name() == f\"sync_{generator_id}\" and not task.done()\n for task in self._running_tasks\n )\n if not task_exists:\n task = asyncio.create_task(\n self._run_generator(generator_id, generator_func),\n name=f\"sync_{generator_id}\",\n )\n self._running_tasks.append(task)\n\n # Check for newly added async generators\n for generator_id, async_generator_func in async_generators_snapshot:\n if generator_id not in processed_generators_snapshot:\n # Check if a task is already running this generator\n task_exists = any(\n task.get_name() == f\"async_{generator_id}\" and not task.done()\n for task in self._running_tasks\n )\n if not task_exists:\n task = asyncio.create_task(\n self._run_async_generator(\n generator_id, async_generator_func\n ),\n name=f\"async_{generator_id}\",\n )\n self._running_tasks.append(task)\n\n # Check if all generators have been processed and no active tasks\n async with self._lock:\n all_generators = set(self._generators.keys()) | set(\n self._async_generators.keys()\n )\n all_done = self._active_count == 0 and all_generators.issubset(\n self._processed_generators\n )\n\n if all_done:\n # Re-check for newly added generators (avoid race condition)\n await asyncio.sleep(0) # Yield control to allow other tasks to run\n async with self._lock:\n all_generators = set(self._generators.keys()) | set(\n self._async_generators.keys()\n )\n if self._active_count == 0 and all_generators.issubset(\n self._processed_generators\n ):\n break\n\n # Get value from queue (with timeout to avoid infinite wait while allowing new generator checks)\n try:\n item = await asyncio.wait_for(self._queue.get(), timeout=0.01)\n except asyncio.TimeoutError:\n # Continue checking for new generators after timeout\n continue\n\n # Check for special markers\n if isinstance(item, tuple):\n if item[0] == \"__error__\":\n raise item[1]\n elif item[0] == \"__event__\":\n # If it's an event, yield the event object\n yield item[1]\n continue\n\n yield item\n\n # Wait for all tasks to complete\n await asyncio.gather(*self._running_tasks, return_exceptions=True)\n\n async def __aiter__(self):\n \"\"\"Allow the merger to be used as an async iterator.\"\"\"\n async for item in self.merge():\n yield item\n\n\nasync def example_usage():\n \"\"\"Example usage - demonstrates generator completion events and deletion functionality.\"\"\"\n print(\"=\" * 60)\n print(\"Example 1: Using callback\")\n print(\"=\" * 60)\n\n # Define callback\n async def on_complete(\n generator_id: str, generator_type: str, error: Optional[Exception]\n ):\n if error:\n print(\n f\" [CALLBACK] Generator {generator_id} ({generator_type}) completed with error: {error}\"\n )\n else:\n print(f\" [CALLBACK] Generator {generator_id} ({generator_type}) completed\")\n\n # Create merger with callback\n merger = GeneratorMerger(on_generator_complete=on_complete)\n\n def generator1():\n for i in range(1, 5):\n yield f\"sync-gen1-{i}\"\n\n async def async_generator1():\n for i in range(20, 25):\n await asyncio.sleep(0.01)\n yield f\"async-gen1-{i}\"\n\n def generator2():\n for i in range(1, 5):\n yield f\"sync-gen2-{i}\"\n\n async def async_generator2():\n for i in range(40, 45):\n await asyncio.sleep(0.01)\n yield f\"async-gen2-{i}\"\n\n # Add generators\n merger.add_generator(generator1, generator_id=\"gen1\")\n merger.add_async_generator(async_generator1, generator_id=\"async_gen1\")\n merger.add_generator(generator2, generator_id=\"gen2\")\n merger.add_async_generator(async_generator2, generator_id=\"async_gen2\")\n\n print(\"Merged generator output:\")\n async for item in merger:\n if isinstance(item, dict) and item.get(\"type\") == \"generator_complete\":\n # If callback exists, events won't be yielded, so this won't execute\n print(f\" Event: {item}\")\n else:\n print(f\" Data: {item}\")\n\n print(\"\\n\" + \"=\" * 60)\n print(\"Example 2: Dynamic addition and deletion of generators\")\n print(\"=\" * 60)\n\n\nif __name__ == \"__main__\":\n asyncio.run(example_usage())\n" }, { "path": "cortex/utils/generator_merger_examples.py", "content": "\"\"\"GeneratorMerger usage examples collection.\"\"\"\n\nimport asyncio\nimport time\n\nfrom .generator_merger import GeneratorMerger\n\n\nasync def example_1_basic_usage():\n \"\"\"Example 1: Basic usage - merging multiple simple generators.\"\"\"\n print(\"=\" * 60)\n print(\"Example 1: Basic Usage\")\n print(\"=\" * 60)\n\n merger = GeneratorMerger()\n\n def number_generator(start: int, end: int, prefix: str):\n \"\"\"Generator that produces number sequences.\"\"\"\n for i in range(start, end):\n yield f\"{prefix}-{i}\"\n\n # Add multiple generators\n merger.add_generator(lambda: number_generator(1, 4, \"A\"))\n merger.add_generator(lambda: number_generator(10, 13, \"B\"))\n merger.add_generator(lambda: number_generator(20, 23, \"C\"))\n\n print(\"Merged output (order may vary due to concurrent execution):\")\n async for item in merger:\n print(f\" Received: {item}\")\n\n print()\n\n\nasync def example_2_dynamic_addition():\n \"\"\"Example 2: Dynamic generator addition - adding new generators during iteration.\"\"\"\n print(\"=\" * 60)\n print(\"Example 2: Dynamic Generator Addition\")\n print(\"=\" * 60)\n\n merger = GeneratorMerger()\n\n def slow_generator(name: str, count: int, delay: float):\n \"\"\"Generator with delay, simulating slow data source.\"\"\"\n for i in range(count):\n time.sleep(delay) # Simulate I/O operation\n yield f\"{name}-{i}\"\n\n # Initially add two generators\n merger.add_generator(lambda: slow_generator(\"Fast\", 3, 0.1))\n merger.add_generator(lambda: slow_generator(\"Medium\", 3, 0.15))\n\n print(\"Starting iteration, will dynamically add new generator later...\")\n count = 0\n async for item in merger:\n print(f\" Received: {item}\")\n count += 1\n\n # Dynamically add new generator during iteration\n if count == 3:\n print(\" Dynamically adding new generator...\")\n merger.add_generator(lambda: slow_generator(\"Slow\", 3, 0.2))\n\n print()\n\n\nasync def example_3_different_data_types():\n \"\"\"Example 3: Handling different data types.\"\"\"\n print(\"=\" * 60)\n print(\"Example 3: Handling Different Data Types\")\n print(\"=\" * 60)\n\n merger = GeneratorMerger()\n\n def string_generator():\n \"\"\"Generate strings.\"\"\"\n for s in [\"hello\", \"world\", \"python\"]:\n yield s\n\n def number_generator():\n \"\"\"Generate numbers.\"\"\"\n for n in [1, 2, 3, 4, 5]:\n yield n\n\n def dict_generator():\n \"\"\"Generate dictionaries.\"\"\"\n for i in range(3):\n yield {\"id\": i, \"name\": f\"item_{i}\", \"value\": i * 10}\n\n merger.add_generator(string_generator)\n merger.add_generator(number_generator)\n merger.add_generator(dict_generator)\n\n print(\"Merging different data types:\")\n async for item in merger:\n print(f\" Type: {type(item).__name__}, Value: {item}\")\n\n print()\n\n\nasync def example_4_data_streams():\n \"\"\"Example 4: Simulating multiple data stream scenarios.\"\"\"\n print(\"=\" * 60)\n print(\"Example 4: Simulating Multiple Data Stream Scenarios\")\n print(\"=\" * 60)\n\n merger = GeneratorMerger()\n\n def log_stream(source: str):\n \"\"\"Simulate log stream.\"\"\"\n for i in range(5):\n time.sleep(0.05)\n yield {\"source\": source, \"level\": \"INFO\", \"message\": f\"Log entry {i}\"}\n\n def metric_stream(metric_name: str):\n \"\"\"Simulate metric stream.\"\"\"\n for i in range(4):\n time.sleep(0.08)\n yield {\"metric\": metric_name, \"value\": i * 10, \"timestamp\": time.time()}\n\n def event_stream(event_type: str):\n \"\"\"Simulate event stream.\"\"\"\n for i in range(3):\n time.sleep(0.06)\n yield {\"event\": event_type, \"id\": i, \"data\": f\"event_data_{i}\"}\n\n # Add multiple data streams\n merger.add_generator(lambda: log_stream(\"server1\"))\n merger.add_generator(lambda: log_stream(\"server2\"))\n merger.add_generator(lambda: metric_stream(\"cpu_usage\"))\n merger.add_generator(lambda: metric_stream(\"memory_usage\"))\n merger.add_generator(lambda: event_stream(\"user_action\"))\n\n print(\"Merging multiple data streams (real-time output):\")\n async for item in merger:\n if \"source\" in item:\n print(f\" [LOG] {item['source']}: {item['message']}\")\n elif \"metric\" in item:\n print(f\" [METRIC] {item['metric']}: {item['value']}\")\n elif \"event\" in item:\n print(f\" [EVENT] {item['event']}: {item['data']}\")\n\n print()\n\n\nasync def example_5_file_processing():\n \"\"\"Example 5: Simulating file processing scenario - merging content from multiple files.\"\"\"\n print(\"=\" * 60)\n print(\"Example 5: Simulating File Processing Scenario\")\n print(\"=\" * 60)\n\n merger = GeneratorMerger()\n\n def file_reader(filename: str, lines: list[str]):\n \"\"\"Simulate file reader.\"\"\"\n for line_num, line in enumerate(lines, 1):\n time.sleep(0.02) # Simulate read delay\n yield {\"file\": filename, \"line\": line_num, \"content\": line}\n\n # Simulate content of three files\n file1_content = [\"Line 1\", \"Line 2\", \"Line 3\"]\n file2_content = [\"A\", \"B\", \"C\", \"D\"]\n file3_content = [\"Data 1\", \"Data 2\"]\n\n merger.add_generator(lambda: file_reader(\"file1.txt\", file1_content))\n merger.add_generator(lambda: file_reader(\"file2.txt\", file2_content))\n merger.add_generator(lambda: file_reader(\"file3.txt\", file3_content))\n\n print(\"Merging and processing multiple files:\")\n async for item in merger:\n print(f\" [{item['file']}] Line {item['line']}: {item['content']}\")\n\n print()\n\n\nasync def example_6_batch_processing():\n \"\"\"Example 6: Batch processing scenario - merging results from multiple tasks.\"\"\"\n print(\"=\" * 60)\n print(\"Example 6: Batch Processing Scenario\")\n print(\"=\" * 60)\n\n merger = GeneratorMerger()\n\n def task_processor(task_id: int, items: list[str]):\n \"\"\"Simulate task processor.\"\"\"\n for item in items:\n time.sleep(0.03) # Simulate processing time\n yield {\n \"task_id\": task_id,\n \"item\": item,\n \"status\": \"processed\",\n \"result\": f\"result_{item}\",\n }\n\n # Add multiple tasks\n merger.add_generator(lambda: task_processor(1, [\"item1\", \"item2\", \"item3\"]))\n merger.add_generator(lambda: task_processor(2, [\"itemA\", \"itemB\"]))\n merger.add_generator(\n lambda: task_processor(3, [\"data1\", \"data2\", \"data3\", \"data4\"])\n )\n\n print(\"Concurrently processing multiple tasks:\")\n results = []\n async for item in merger:\n results.append(item)\n print(f\" Task {item['task_id']} completed: {item['item']} -> {item['result']}\")\n\n print(f\"\\nTotal processed: {len(results)} items\")\n print()\n\n\nasync def example_7_error_handling():\n \"\"\"Example 7: Error handling - demonstrating behavior when generator errors occur.\"\"\"\n print(\"=\" * 60)\n print(\"Example 7: Error Handling\")\n print(\"=\" * 60)\n\n merger = GeneratorMerger()\n\n def normal_generator():\n \"\"\"Normal generator.\"\"\"\n for i in range(3):\n yield f\"normal-{i}\"\n\n def error_generator():\n \"\"\"Generator that will error.\"\"\"\n yield \"error-1\"\n yield \"error-2\"\n raise ValueError(\"Simulated error\")\n\n def another_normal_generator():\n \"\"\"Another normal generator.\"\"\"\n for i in range(2):\n yield f\"another-{i}\"\n\n merger.add_generator(normal_generator)\n merger.add_generator(error_generator)\n merger.add_generator(another_normal_generator)\n\n print(\"Processing generators with errors:\")\n try:\n async for item in merger:\n print(f\" Received: {item}\")\n except ValueError as e:\n print(f\" Caught error: {e}\")\n\n print()\n\n\nasync def example_8_callback_usage():\n \"\"\"Example 8: Using callback to monitor generator completion events.\"\"\"\n print(\"=\" * 60)\n print(\"Example 8: Using Callback to Monitor Generator Completion Events\")\n print(\"=\" * 60)\n\n # Track completed generators\n completed_generators = []\n\n # Define callback function\n async def on_generator_complete(\n generator_id: str, generator_type: str, error: Exception | None\n ):\n \"\"\"Called when generator completes.\"\"\"\n status = \"Success\" if error is None else f\"Failed: {error}\"\n completed_generators.append(\n {\"id\": generator_id, \"type\": generator_type, \"status\": status}\n )\n print(\n f\" [CALLBACK] Generator '{generator_id}' ({generator_type}) completed: {status}\"\n )\n\n # Create merger with callback\n merger = GeneratorMerger(on_generator_complete=on_generator_complete)\n\n def fast_generator():\n \"\"\"Fast completing generator.\"\"\"\n for i in range(1, 4):\n yield f\"fast-{i}\"\n\n def slow_generator():\n \"\"\"Slow completing generator.\"\"\"\n for i in range(10, 13):\n time.sleep(0.05) # Simulate slow operation\n yield f\"slow-{i}\"\n\n async def async_generator():\n \"\"\"Async generator.\"\"\"\n for i in range(20, 23):\n await asyncio.sleep(0.03)\n yield f\"async-{i}\"\n\n def error_generator():\n \"\"\"Generator that will error.\"\"\"\n yield \"error-1\"\n yield \"error-2\"\n raise ValueError(\"Test error\")\n\n # Add generators with generator_id specified\n merger.add_generator(fast_generator, generator_id=\"fast_gen\")\n merger.add_generator(slow_generator, generator_id=\"slow_gen\")\n merger.add_async_generator(async_generator, generator_id=\"async_gen\")\n merger.add_generator(error_generator, generator_id=\"error_gen\")\n\n print(\"Starting generator processing (callback will be called on completion):\")\n print()\n\n try:\n async for item in merger:\n # Note: If there's a callback, completion events won't be yielded, only callback is called\n print(f\" [DATA] Received data: {item}\")\n except ValueError as e:\n print(f\" [ERROR] Caught error: {e}\")\n\n print()\n print(\"Completed generators statistics:\")\n for gen_info in completed_generators:\n print(f\" - {gen_info['id']} ({gen_info['type']}): {gen_info['status']}\")\n\n print()\n\n\nasync def example_9_callback_vs_event():\n \"\"\"Example 9: Comparing callback vs event yield differences.\"\"\"\n print(\"=\" * 60)\n print(\"Example 9: Comparing Callback vs Event Yield Differences\")\n print(\"=\" * 60)\n\n def simple_generator(name: str):\n \"\"\"Simple generator.\"\"\"\n for i in range(1, 3):\n yield f\"{name}-{i}\"\n\n print(\"Method 1: Using callback (events won't appear in iteration)\")\n print(\"-\" * 60)\n\n async def callback(\n generator_id: str, _generator_type: str, _error: Exception | None\n ):\n print(f\" [CALLBACK] {generator_id} completed\")\n\n merger1 = GeneratorMerger(on_generator_complete=callback)\n merger1.add_generator(lambda: simple_generator(\"A\"), generator_id=\"gen_A\")\n merger1.add_generator(lambda: simple_generator(\"B\"), generator_id=\"gen_B\")\n\n print(\" Iteration output:\")\n async for item in merger1:\n print(f\" {item}\")\n\n print()\n print(\"Method 2: Without callback (events will be yielded)\")\n print(\"-\" * 60)\n\n merger2 = GeneratorMerger()\n merger2.add_generator(lambda: simple_generator(\"X\"), generator_id=\"gen_X\")\n merger2.add_generator(lambda: simple_generator(\"Y\"), generator_id=\"gen_Y\")\n\n print(\" Iteration output (includes completion events):\")\n async for item in merger2:\n if isinstance(item, dict) and item.get(\"type\") == \"generator_complete\":\n print(\n f\" [EVENT] {item['generator_id']} ({item['generator_type']}) \"\n f\"Status: {item['status']}\"\n )\n else:\n print(f\" [DATA] {item}\")\n\n print()\n\n\nasync def run_all_examples():\n \"\"\"Run all examples.\"\"\"\n examples = [\n example_1_basic_usage,\n example_2_dynamic_addition,\n example_3_different_data_types,\n example_4_data_streams,\n example_5_file_processing,\n example_6_batch_processing,\n example_7_error_handling,\n example_8_callback_usage,\n example_9_callback_vs_event,\n ]\n\n for example in examples:\n try:\n await example()\n await asyncio.sleep(0.5) # Interval between examples\n except Exception as e:\n print(f\"Example execution error: {e}\\n\")\n\n\nif __name__ == \"__main__\":\n print(\"\\n\" + \"=\" * 60)\n print(\"GeneratorMerger Usage Examples\")\n print(\"=\" * 60 + \"\\n\")\n\n # Run all examples\n asyncio.run(run_all_examples())\n\n print(\"=\" * 60)\n print(\"All examples completed!\")\n print(\"=\" * 60)\n" }, { "path": "cortex-ui/.gitignore", "content": "# Logs\nlogs\n*.log\nnpm-debug.log*\nyarn-debug.log*\nyarn-error.log*\npnpm-debug.log*\nlerna-debug.log*\n\nnode_modules\ndist\ndist-ssr\n*.local\n\n# Editor directories and files\n.vscode/*\n!.vscode/extensions.json\n.idea\n.DS_Store\n*.suo\n*.ntvs*\n*.njsproj\n*.sln\n*.sw?\n\n" }, { "path": "cortex-ui/.gitlab-ci.yml", "content": "variables:\n APP: $CI_PROJECT_NAMESPACE\n SERVER: $CI_PROJECT_NAME\n VERSION: $CI_COMMIT_SHA\n IMAGE_TAG_PROD: ${CI_COMMIT_TAG}-${CI_COMMIT_SHORT_SHA}\n IMAGE_TAG_DEV: dev-${CI_COMMIT_BRANCH}${CI_MERGE_REQUEST_SOURCE_BRANCH_NAME}-${CI_COMMIT_SHORT_SHA}\n IMAGE: catalyst-cn-shanghai.cr.volces.com/capy/${CI_PROJECT_PATH}\n\nworkflow:\n rules:\n - if: $CI_PIPELINE_SOURCE == \"push\" && $CI_COMMIT_TITLE == \"edit by devops\"\n when: never\n - if: $CI_COMMIT_BRANCH && $CI_OPEN_MERGE_REQUESTS\n when: never\n - if: $CI_PIPELINE_SOURCE == \"merge_request_event\"\n - if: $CI_PIPELINE_SOURCE == \"schedule\"\n - if: $CI_COMMIT_BRANCH\n - if: $CI_COMMIT_TAG\n\ndefault:\n tags: [shell]\n before_script:\n - export PATH=$PWD/node_modules/.bin:$PATH\n # - command -v moon || pnpm install\n - pnpm install\n - git fetch origin main\n\nstages:\n # - lint\n - image\n\n\nimage-push-dev:\n stage: image\n script:\n - docker build -f \"Dockerfile\" --platform=linux/amd64 -t ${IMAGE}:${IMAGE_TAG_DEV} ${CI_PROJECT_DIR}\n - docker push ${IMAGE}:${IMAGE_TAG_DEV}\n rules:\n - when: manual\n allow_failure: true\n\n" }, { "path": "cortex-ui/index.html", "content": "\n\n \n \n \n \n Agent Cortex\n \n \n
\n \n \n\n\n" }, { "path": "cortex-ui/package.json", "content": "{\n \"name\": \"agent-cortex-frontend\",\n \"version\": \"1.0.0\",\n \"type\": \"module\",\n \"scripts\": {\n \"dev\": \"vite\",\n \"build\": \"tsc && vite build\",\n \"preview\": \"vite preview\"\n },\n \"dependencies\": {\n \"@ant-design/icons\": \"^5.2.6\",\n \"antd\": \"^5.12.0\",\n \"axios\": \"^1.6.2\",\n \"react\": \"^18.2.0\",\n \"react-dom\": \"^18.2.0\",\n \"react-markdown\": \"^9.0.1\",\n \"react-router-dom\": \"^6.20.0\",\n \"rehype-raw\": \"^7.0.0\",\n \"remark-gfm\": \"^4.0.0\"\n },\n \"devDependencies\": {\n \"@types/react\": \"^18.2.43\",\n \"@types/react-dom\": \"^18.2.17\",\n \"@vitejs/plugin-react\": \"^4.2.1\",\n \"typescript\": \"^5.3.3\",\n \"vite\": \"^5.0.8\"\n }\n}\n" }, { "path": "cortex-ui/src/App.tsx", "content": "import React from 'react';\nimport { BrowserRouter as Router, Routes, Route } from 'react-router-dom';\nimport { ConfigProvider } from 'antd';\nimport zhCN from 'antd/locale/zh_CN';\nimport AgentList from './pages/AgentList';\nimport ChatPage from './pages/ChatPage';\nimport { ErrorBoundary } from './components/ErrorBoundary';\n\nconst App: React.FC = () => {\n return (\n \n \n \n \n } />\n } />\n \n \n \n \n );\n};\n\nexport default App;\n\n" }, { "path": "cortex-ui/src/components/EndpointConfig.tsx", "content": "import React, { useState } from 'react';\nimport { Input, Button, Typography, message } from 'antd';\nimport { EditOutlined, SaveOutlined, CloseOutlined } from '@ant-design/icons';\nimport { getStoredEndpoint } from '../services/api';\n\nconst { Text } = Typography;\n\ninterface EndpointConfigProps {\n onSave?: (endpoint: string) => void;\n}\n\nexport const EndpointConfig: React.FC = ({ onSave }) => {\n const [endpoint, setEndpoint] = useState(() => {\n return getStoredEndpoint() || '';\n });\n const [isEditing, setIsEditing] = useState(false);\n\n const handleSave = () => {\n localStorage.setItem('api_endpoint', endpoint);\n setIsEditing(false);\n message.success('Endpoint saved');\n if (onSave) {\n onSave(endpoint);\n }\n };\n\n const handleCancel = () => {\n setEndpoint(getStoredEndpoint() || '');\n setIsEditing(false);\n };\n\n return (\n
\n {isEditing ? (\n <>\n setEndpoint(e.target.value)}\n placeholder=\"http://localhost:8001\"\n style={{ width: 250 }}\n size=\"small\"\n />\n }\n onClick={handleSave}\n >\n Save\n \n }\n onClick={handleCancel}\n >\n Cancel\n \n \n ) : (\n <>\n {endpoint && (\n \n {endpoint}\n \n )}\n }\n onClick={() => setIsEditing(true)}\n >\n Configure Endpoint\n \n \n )}\n
\n );\n};\n\n" }, { "path": "cortex-ui/src/components/ErrorBoundary.tsx", "content": "import React, { Component, ReactNode } from 'react';\nimport { Result, Button } from 'antd';\n\ninterface Props {\n children: ReactNode;\n}\n\ninterface State {\n hasError: boolean;\n error?: Error;\n}\n\nexport class ErrorBoundary extends Component {\n constructor(props: Props) {\n super(props);\n this.state = { hasError: false };\n }\n\n static getDerivedStateFromError(error: Error): State {\n return { hasError: true, error };\n }\n\n componentDidCatch(error: Error, errorInfo: React.ErrorInfo) {\n console.error('Error caught by boundary:', error, errorInfo);\n }\n\n render() {\n if (this.state.hasError) {\n return (\n window.location.reload()}>\n Refresh Page\n ,\n ,\n ]}\n />\n );\n }\n\n return this.props.children;\n }\n}\n\n" }, { "path": "cortex-ui/src/components/FilePanel.tsx", "content": "/**\n * FilePanel Component\n * Display file tool execution results with elegant UI\n */\n\nimport React, { useState } from 'react';\nimport { Card, Typography, Space, Tag, Tooltip } from 'antd';\nimport {\n FileOutlined,\n FileTextOutlined,\n FolderOutlined,\n EditOutlined,\n PlusOutlined,\n InfoCircleOutlined,\n CheckCircleOutlined,\n CloseCircleOutlined,\n CopyOutlined,\n ExpandOutlined,\n CompressOutlined,\n ClockCircleOutlined,\n} from '@ant-design/icons';\n\nconst { Text } = Typography;\n\nexport interface FileResult {\n action: string;\n path: string;\n size?: string;\n content?: string;\n error?: string;\n note?: string;\n entries?: string[]; // For list action\n modified?: string;\n created?: string;\n}\n\ninterface FilePanelProps {\n result: FileResult;\n title?: string;\n agentName?: string;\n timestamp?: Date;\n}\n\n/**\n * Parse file result from XML-like string\n */\nexport const parseFileResult = (resultStr: string): FileResult | null => {\n try {\n // Try to extract the actual file_result content from nested JSON structure\n let actualContent = resultStr;\n \n // Handle nested JSON structure like {\"type\":\"tool_result\",\"content\":[{\"type\":\"text\",\"text\":\"...\"}]}\n if (resultStr.includes('\"type\":\"tool_result\"') || resultStr.includes('\"tool_result\"')) {\n try {\n const parsed = JSON.parse(resultStr);\n if (parsed.content && Array.isArray(parsed.content)) {\n for (const item of parsed.content) {\n if (item.type === 'text' && item.text) {\n actualContent = item.text;\n break;\n }\n }\n }\n } catch {\n // If JSON parse fails, try regex extraction\n const textMatch = resultStr.match(/\"text\"\\s*:\\s*\"([\\s\\S]*?)(?:\"\\s*}|\\\"\\s*,)/);\n if (textMatch) {\n // Unescape the JSON string\n actualContent = textMatch[1].replace(/\\\\n/g, '\\n').replace(/\\\\\"/g, '\"').replace(/\\\\\\\\/g, '\\\\');\n }\n }\n }\n\n // Check if it's a file_result\n if (!actualContent.includes('')) return null;\n\n const actionMatch = actualContent.match(/action:\\s*(\\w+)/);\n const pathMatch = actualContent.match(/path:\\s*(.+?)(?:\\n|$)/);\n const sizeMatch = actualContent.match(/size:\\s*(.+?)(?:\\n|$)/);\n const errorMatch = actualContent.match(/error:\\s*(.+?)(?:\\n|<\\/file_result>)/s);\n const noteMatch = actualContent.match(/note:\\s*(.+?)(?:\\n|$)/);\n const modifiedMatch = actualContent.match(/modified:\\s*(.+?)(?:\\n|$)/);\n const createdMatch = actualContent.match(/created:\\s*(.+?)(?:\\n|$)/);\n\n if (!actionMatch || !pathMatch) return null;\n\n const result: FileResult = {\n action: actionMatch[1],\n path: pathMatch[1].trim(),\n };\n\n if (sizeMatch) result.size = sizeMatch[1].trim();\n if (errorMatch) result.error = errorMatch[1].trim();\n if (noteMatch) result.note = noteMatch[1].trim();\n if (modifiedMatch) result.modified = modifiedMatch[1].trim();\n if (createdMatch) result.created = createdMatch[1].trim();\n\n // Extract content for read action (between ``` markers)\n const contentMatch = actualContent.match(/```\\n?([\\s\\S]*?)\\n?```/);\n if (contentMatch) {\n result.content = contentMatch[1];\n }\n\n // Extract entries for list action\n if (actionMatch[1] === 'list') {\n const entriesSection = actualContent.match(/-{20,}\\n([\\s\\S]*?)<\\/file_result>/);\n if (entriesSection) {\n const lines = entriesSection[1].trim().split('\\n').filter(line => line.trim());\n result.entries = lines;\n }\n }\n\n return result;\n } catch (e) {\n return null;\n }\n};\n\n/**\n * Check if content is file result\n */\nexport const isFileResult = (content: string): boolean => {\n // Check for direct file_result tag\n if (content.includes('') && content.includes('action:')) {\n return true;\n }\n // Check for nested JSON structure containing file_result\n if (content.includes('\"tool_result\"') && content.includes('file_result')) {\n return true;\n }\n return false;\n};\n\n/**\n * Get icon based on action type\n */\nconst getActionIcon = (action: string) => {\n switch (action.toLowerCase()) {\n case 'read':\n return ;\n case 'write':\n return ;\n case 'append':\n return ;\n case 'list':\n return ;\n case 'stat':\n return ;\n default:\n return ;\n }\n};\n\n/**\n * Get action display text\n */\nconst getActionText = (action: string): string => {\n switch (action.toLowerCase()) {\n case 'read':\n return 'Read File';\n case 'write':\n return 'Write File';\n case 'append':\n return 'Append File';\n case 'list':\n return 'List Directory';\n case 'stat':\n return 'File Info';\n default:\n return action;\n }\n};\n\n/**\n * Get action color\n */\nconst getActionColor = (action: string): string => {\n switch (action.toLowerCase()) {\n case 'read':\n return '#3b82f6';\n case 'write':\n return '#22c55e';\n case 'append':\n return '#f59e0b';\n case 'list':\n return '#8b5cf6';\n case 'stat':\n return '#6366f1';\n default:\n return '#64748b';\n }\n};\n\n/**\n * FilePanel main component\n */\nconst FilePanel: React.FC = ({\n result,\n title = 'File Tool',\n agentName,\n timestamp,\n}) => {\n const [expanded, setExpanded] = useState(false);\n const [copied, setCopied] = useState(false);\n\n const hasError = !!result.error;\n const hasContent = !!result.content;\n const hasEntries = result.entries && result.entries.length > 0;\n const actionColor = getActionColor(result.action);\n\n const copyPath = () => {\n navigator.clipboard.writeText(result.path);\n setCopied(true);\n setTimeout(() => setCopied(false), 2000);\n };\n\n const copyContent = () => {\n if (result.content) {\n navigator.clipboard.writeText(result.content);\n setCopied(true);\n setTimeout(() => setCopied(false), 2000);\n }\n };\n\n // Extract filename from path\n const filename = result.path.split('/').pop() || result.path;\n\n return (\n \n {/* Header */}\n \n \n \n {React.cloneElement(getActionIcon(result.action), { style: { color: 'white', fontSize: '16px' } })}\n \n
\n
\n \n {agentName ? `${agentName}` : title}\n \n \n {result.action.toUpperCase()}\n \n
\n
\n {getActionText(result.action)} • {result.size || 'N/A'}\n
\n
\n \n \n : }\n color={hasError ? 'error' : 'success'}\n style={{\n margin: 0,\n borderRadius: 12,\n fontSize: 11,\n padding: '0 8px',\n }}\n >\n {hasError ? 'Error' : 'Success'}\n \n {timestamp && (\n \n \n \n )}\n \n \n\n {/* Path section */}\n \n
\n \n \n \n {filename}\n \n \n \n \n \n
\n \n {result.path}\n \n \n\n {/* Error section */}\n {hasError && (\n \n \n \n {result.error}\n \n \n )}\n\n {/* Note section */}\n {result.note && !hasError && (\n \n \n \n {result.note}\n \n \n )}\n\n {/* Stat info section */}\n {result.action === 'stat' && (result.modified || result.created) && (\n
\n \n {result.modified && (\n
\n Modified:\n {result.modified}\n
\n )}\n {result.created && (\n
\n Created:\n {result.created}\n
\n )}\n \n
\n )}\n\n {/* List entries section */}\n {hasEntries && (\n
\n
\n \n Entries ({result.entries!.length})\n \n setExpanded(!expanded)}\n >\n {expanded ? (\n <>\n Collapse\n \n ) : (\n <>\n Expand\n \n )}\n \n
\n \n \n {result.entries!.join('\\n')}\n \n
\n \n )}\n\n {/* Content section */}\n {hasContent && (\n
\n
\n \n Content\n \n \n \n \n \n setExpanded(!expanded)}\n >\n {expanded ? (\n <>\n Collapse\n \n ) : (\n <>\n Expand\n \n )}\n \n \n
\n \n \n {result.content}\n \n
\n \n )}\n\n {/* Timestamp footer */}\n {timestamp && (\n \n \n {timestamp.toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })}\n \n \n )}\n \n );\n};\n\nexport default FilePanel;\n" }, { "path": "cortex-ui/src/components/FinalAnswer.tsx", "content": "/**\n * FinalAnswer Component\n * Renders final answer with citations, supports Markdown and \\cite{} format\n */\n\nimport React, { useMemo } from 'react';\nimport { Card, Typography, Tooltip, Tag } from 'antd';\nimport {\n CheckCircleOutlined,\n LinkOutlined,\n BookOutlined,\n} from '@ant-design/icons';\nimport ReactMarkdown from 'react-markdown';\nimport remarkGfm from 'remark-gfm';\nimport rehypeRaw from 'rehype-raw';\nimport { CitationStore, CitationInfo } from '../types/citation';\n\nconst { Text, Link } = Typography;\n\ninterface FinalAnswerProps {\n content: string; // Content with ... tags or plain answer\n citations: CitationStore; // Citation info storage\n}\n\n/**\n * Extract content between ... tags\n */\nfunction extractAnswerContent(content: string): string {\n const answerMatch = content.match(/([\\s\\S]*?)<\\/answer>/);\n if (answerMatch) {\n return answerMatch[1].trim();\n }\n // If no answer tag, return original content\n return content;\n}\n\n/**\n * Convert \\cite{web_xxx} to HTML span tags before Markdown rendering\n * This allows rehype-raw to render them properly\n */\nfunction preprocessCitations(content: string, citations: CitationStore): string {\n // Match \\cite{web_xxxxx} format\n const citeRegex = /\\\\cite\\{(web_[a-zA-Z0-9]+)\\}/g;\n \n return content.replace(citeRegex, (_, citeIndex: string) => {\n const citation = citations.get(citeIndex);\n const shortId = citeIndex.replace('web_', '').substring(0, 4);\n \n if (citation) {\n // Create a clickable citation tag with data attributes\n return `${shortId}`;\n } else {\n // Unknown citation\n return `${shortId}`;\n }\n });\n}\n\n/**\n * Custom Markdown rendering components\n */\nconst createMarkdownComponents = (_citations: CitationStore) => ({\n // Custom cite-tag rendering\n 'cite-tag': ({ node, children, ...props }: any) => {\n const citeIndex = props['data-cite-index'];\n const isUnknown = props['data-unknown'] === 'true';\n \n if (isUnknown) {\n return (\n \n \n {children}\n \n \n );\n }\n \n const url = decodeURIComponent(props['data-url'] || '');\n const title = decodeURIComponent(props['data-title'] || '');\n const snippet = decodeURIComponent(props['data-snippet'] || '');\n \n return (\n \n
\n {title || 'No Title'}\n
\n {snippet && (\n
\n {snippet.substring(0, 150)}\n {snippet.length > 150 ? '...' : ''}\n
\n )}\n
\n Click to visit source →\n
\n \n }\n placement=\"top\"\n >\n window.open(url, '_blank')}\n >\n \n {children}\n \n \n );\n },\n // Custom heading rendering\n h1: ({ children, ...props }: any) => (\n

\n {children}\n

\n ),\n h2: ({ children, ...props }: any) => (\n

\n {children}\n

\n ),\n h3: ({ children, ...props }: any) => (\n

\n {children}\n

\n ),\n // Custom table styles\n table: ({ children, ...props }: any) => (\n
\n \n {children}\n
\n
\n ),\n th: ({ children, ...props }: any) => (\n \n {children}\n \n ),\n td: ({ children, ...props }: any) => (\n \n {children}\n \n ),\n // Custom code block\n code: ({ inline, className, children, ...props }: any) => {\n if (inline) {\n return (\n \n {children}\n \n );\n }\n return (\n \n \n {children}\n \n \n );\n },\n // Custom blockquote\n blockquote: ({ children, ...props }: any) => (\n \n {children}\n \n ),\n // Horizontal rule\n hr: ({ ...props }: any) => (\n \n ),\n});\n\n/**\n * FinalAnswer main component\n */\nconst FinalAnswer: React.FC = ({ content, citations }) => {\n // Extract answer content\n const answerContent = useMemo(() => extractAnswerContent(content), [content]);\n \n // Preprocess citations - convert \\cite{web_xxx} to HTML tags\n const processedContent = useMemo(\n () => preprocessCitations(answerContent, citations),\n [answerContent, citations]\n );\n \n // Create Markdown components\n const components = useMemo(() => createMarkdownComponents(citations), [citations]);\n \n // Collect all used citations\n const usedCitations = useMemo(() => {\n const used: CitationInfo[] = [];\n const citeRegex = /\\\\cite\\{(web_[a-zA-Z0-9]+)\\}/g;\n let match;\n const seen = new Set();\n \n while ((match = citeRegex.exec(answerContent)) !== null) {\n const citeIndex = match[1];\n if (!seen.has(citeIndex)) {\n seen.add(citeIndex);\n const citation = citations.get(citeIndex);\n if (citation) {\n used.push(citation);\n }\n }\n }\n \n return used;\n }, [answerContent, citations]);\n \n return (\n \n
\n \n
\n
\n
\n Final Answer\n
\n
\n Research completed\n
\n
\n {usedCitations.length > 0 && (\n }\n style={{ \n borderRadius: '12px',\n background: 'linear-gradient(135deg, #3b82f6 0%, #1d4ed8 100%)',\n border: 'none',\n color: 'white',\n fontSize: '11px',\n padding: '2px 10px',\n }}\n >\n {usedCitations.length} citations\n \n )}\n \n }\n >\n
\n \n {processedContent}\n \n
\n \n {/* Citation list */}\n {usedCitations.length > 0 && (\n \n
\n
\n \n
\n \n References\n \n
\n
\n {usedCitations.map((citation) => (\n {\n e.currentTarget.style.backgroundColor = '#f0fdf4';\n e.currentTarget.style.borderColor = '#bbf7d0';\n }}\n onMouseLeave={(e) => {\n e.currentTarget.style.backgroundColor = '#f8fafc';\n e.currentTarget.style.borderColor = '#e2e8f0';\n }}\n >\n \n {citation.cite_index.replace('web_', '').substring(0, 6)}\n \n
\n \n {citation.title}\n \n
\n
\n ))}\n \n \n )}\n \n );\n};\n\nexport default FinalAnswer;\n" }, { "path": "cortex-ui/src/components/SearchResultsPanel.tsx", "content": "/**\n * SearchResultsPanel Component\n * Display search results from batch_search tool\n */\n\nimport React, { useState } from 'react';\nimport { Card, Collapse, Tag, Typography, Space, Tooltip, Empty } from 'antd';\nimport {\n SearchOutlined,\n ClockCircleOutlined,\n GlobalOutlined,\n ExpandOutlined,\n CompressOutlined,\n} from '@ant-design/icons';\nimport { QueryResult, SearchResultItem } from '../types/citation';\n\nconst { Text, Paragraph, Link } = Typography;\n\ninterface SearchResultsPanelProps {\n queryResults: QueryResult[];\n title?: string;\n agentName?: string;\n timestamp?: Date;\n}\n\n/**\n * Single search result item component\n */\nconst SearchResultItemCard: React.FC<{ item: SearchResultItem; index: number }> = ({ item }) => {\n const [expanded, setExpanded] = useState(false);\n \n return (\n {\n e.currentTarget.style.borderColor = '#1890ff';\n e.currentTarget.style.boxShadow = '0 2px 8px rgba(24, 144, 255, 0.15)';\n }}\n onMouseLeave={(e) => {\n e.currentTarget.style.borderColor = '#f0f0f0';\n e.currentTarget.style.boxShadow = 'none';\n }}\n >\n \n {/* Title and citation index */}\n
\n
\n \n {item.title}\n \n
\n \n \n {item.cite_index}\n \n \n
\n \n {/* URL and metadata */}\n \n {item.site && (\n \n \n {item.site}\n \n )}\n {item.published_time && (\n \n \n {item.published_time.split('T')[0]}\n \n )}\n \n \n {/* Summary */}\n {item.snippet && (\n \n {item.snippet}\n \n )}\n \n {/* Full content (expandable) */}\n {item.content && item.content !== item.snippet && (\n
\n setExpanded(!expanded)}\n >\n {expanded ? (\n <>\n \n Collapse\n \n ) : (\n <>\n \n Show more\n \n )}\n \n {expanded && (\n \n \n {item.content}\n \n
\n )}\n \n )}\n \n \n );\n};\n\n/**\n * SearchResultsPanel main component\n */\nconst SearchResultsPanel: React.FC = ({ \n queryResults, \n title = 'Search Results',\n agentName,\n timestamp,\n}) => {\n if (!queryResults || queryResults.length === 0) {\n return (\n \n \n \n );\n }\n \n // Calculate total results count\n const totalItems = queryResults.reduce((sum, qr) => sum + qr.items.length, 0);\n \n // Prepare Collapse items\n const collapseItems = queryResults.map((queryResult, index) => ({\n key: `query_${index}`,\n label: (\n
\n
\n \n
\n \n {queryResult.query}\n \n \n {queryResult.items.length} results\n \n
\n ),\n children: (\n
\n {queryResult.items.length === 0 ? (\n No results for this query\n ) : (\n queryResult.items.map((item, itemIndex) => (\n \n ))\n )}\n
\n ),\n }));\n \n return (\n \n
\n \n
\n
\n
\n {agentName ? `${agentName}` : title}\n
\n
\n {queryResults.length} queries • {totalItems} results\n
\n
\n \n }\n >\n \n {timestamp && (\n
\n {timestamp.toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })}\n
\n )}\n \n );\n};\n\nexport default SearchResultsPanel;\n" }, { "path": "cortex-ui/src/components/ShellPanel.tsx", "content": "/**\n * ShellPanel Component\n * Display shell command execution results with elegant terminal-style UI\n */\n\nimport React, { useState } from 'react';\nimport { Card, Typography, Space, Tag, Tooltip } from 'antd';\nimport {\n CodeOutlined,\n CheckCircleOutlined,\n CloseCircleOutlined,\n CopyOutlined,\n ExpandOutlined,\n CompressOutlined,\n ClockCircleOutlined,\n} from '@ant-design/icons';\n\nconst { Text } = Typography;\n\nexport interface ShellResult {\n cmd: string;\n exitCode: number;\n stdout: string;\n stderr: string;\n}\n\ninterface ShellPanelProps {\n result: ShellResult;\n title?: string;\n agentName?: string;\n timestamp?: Date;\n}\n\n/**\n * Parse shell result from XML string\n */\nexport const parseShellResult = (resultStr: string): ShellResult | null => {\n try {\n const cmdMatch = resultStr.match(/([\\s\\S]*?)<\\/cmd>/);\n const exitCodeMatch = resultStr.match(/(\\d+)<\\/exit_code>/);\n const stdoutMatch = resultStr.match(/([\\s\\S]*?)<\\/stdout>/);\n const stderrMatch = resultStr.match(/([\\s\\S]*?)<\\/stderr>/);\n\n if (!cmdMatch) return null;\n\n return {\n cmd: cmdMatch[1] || '',\n exitCode: exitCodeMatch ? parseInt(exitCodeMatch[1], 10) : 0,\n stdout: stdoutMatch ? stdoutMatch[1] : '',\n stderr: stderrMatch ? stderrMatch[1] : '',\n };\n } catch (e) {\n return null;\n }\n};\n\n/**\n * Check if content is shell result\n */\nexport const isShellResult = (content: string): boolean => {\n return content.includes('') && content.includes('');\n};\n\n/**\n * ShellPanel main component\n */\nconst ShellPanel: React.FC = ({\n result,\n title = 'Terminal',\n agentName,\n timestamp,\n}) => {\n const [expanded, setExpanded] = useState(false);\n const [copied, setCopied] = useState(false);\n\n const isSuccess = result.exitCode === 0;\n const hasOutput = result.stdout.trim() || result.stderr.trim();\n\n const copyCommand = () => {\n navigator.clipboard.writeText(result.cmd);\n setCopied(true);\n setTimeout(() => setCopied(false), 2000);\n };\n\n const maxOutputHeight = expanded ? 'none' : '200px';\n\n return (\n \n {/* Terminal header */}\n \n \n {/* Traffic light buttons */}\n \n
\n
\n
\n \n \n \n \n {agentName ? `${agentName} • ${title}` : title}\n \n \n \n \n : }\n color={isSuccess ? 'success' : 'error'}\n style={{ \n margin: 0, \n borderRadius: 12,\n fontSize: 11,\n padding: '0 8px',\n }}\n >\n {isSuccess ? 'Success' : `Exit: ${result.exitCode}`}\n \n {timestamp && (\n \n \n \n )}\n \n
\n\n {/* Command section */}\n \n
\n \n $\n \n
\n \n {result.cmd}\n \n
\n \n {\n if (!copied) e.currentTarget.style.color = '#a8a8c0';\n }}\n onMouseLeave={(e) => {\n if (!copied) e.currentTarget.style.color = '#7c7c9a';\n }}\n />\n \n
\n
\n\n {/* Output section */}\n {hasOutput && (\n
\n {/* stdout */}\n {result.stdout.trim() && (\n \n \n {result.stdout}\n \n
\n )}\n\n {/* stderr */}\n {result.stderr.trim() && (\n \n
\n \n stderr\n \n
\n \n {result.stderr}\n \n
\n )}\n\n {/* Expand/Collapse toggle */}\n {(result.stdout.length > 300 || result.stderr.length > 300) && (\n \n \n setExpanded(!expanded)}\n style={{\n display: 'flex',\n alignItems: 'center',\n gap: 4,\n padding: '4px 12px',\n borderRadius: 12,\n backgroundColor: 'rgba(255, 255, 255, 0.05)',\n cursor: 'pointer',\n transition: 'all 0.2s',\n }}\n onMouseEnter={(e) => {\n e.currentTarget.style.backgroundColor = 'rgba(255, 255, 255, 0.1)';\n }}\n onMouseLeave={(e) => {\n e.currentTarget.style.backgroundColor = 'rgba(255, 255, 255, 0.05)';\n }}\n >\n {expanded ? (\n \n ) : (\n \n )}\n \n {expanded ? 'Collapse' : 'Expand'}\n \n \n \n \n )}\n \n )}\n\n {/* Empty output hint */}\n {!hasOutput && (\n \n \n Command completed with no output\n \n \n )}\n \n );\n};\n\nexport default ShellPanel;\n" }, { "path": "cortex-ui/src/components/TodoPanel.tsx", "content": "/**\n * TodoPanel Component\n * Display task list results from todo tool\n */\n\nimport React from 'react';\nimport { Card, Typography, Space, Tag, Progress, Tooltip, Timeline } from 'antd';\nimport {\n CheckCircleOutlined,\n ClockCircleOutlined,\n SyncOutlined,\n UnorderedListOutlined,\n CalendarOutlined,\n LinkOutlined,\n} from '@ant-design/icons';\n\nconst { Text, Paragraph } = Typography;\n\n// Task status type\ntype TaskStatus = 'completed' | 'in_progress' | 'pending' | 'not_started';\n\n// Single task item\nexport interface TodoItem {\n step: number;\n title: string;\n status: TaskStatus;\n description?: string;\n dependsOn?: number[];\n updatedAt?: string;\n createdAt?: string;\n priority?: 'high' | 'medium' | 'low';\n}\n\n// TodoPanel props\ninterface TodoPanelProps {\n items: TodoItem[];\n title?: string;\n totalTasks?: number;\n completedTasks?: number;\n agentName?: string;\n timestamp?: Date;\n}\n\n/**\n * Parse task list from todo_result XML string\n */\nexport function parseTodoResult(content: string): TodoItem[] {\n const items: TodoItem[] = [];\n \n // Normalize content - handle escaped newlines from JSON serialization\n const normalizedContent = content\n .replace(/\\\\n/g, '\\n')\n .replace(/\\\\t/g, '\\t');\n \n // Match each task line\n // Format: ✅ Step 1: Title 🟠\n // Status icons: ✅(completed) ⏳(pending) 🔄(in_progress) 🚫(blocked)\n // Priority icons: 🔴(critical) 🟠(high) 🟡(medium) 🔵(low)\n const taskRegex = /([✅⏳🔄🚫]) Step ([^\\s:]+): ([^\\n🔴🟠🟡🔵]+)([🔴🟠🟡🔵])?/g;\n let match;\n \n while ((match = taskRegex.exec(normalizedContent)) !== null) {\n const statusIcon = match[1];\n const stepId = match[2].trim();\n const step = parseInt(stepId, 10) || 0; // Support non-numeric step IDs\n const title = match[3].trim();\n const priorityIcon = match[4];\n \n // Parse status\n let status: TaskStatus = 'not_started';\n if (statusIcon === '✅') status = 'completed';\n else if (statusIcon === '⏳') status = 'pending';\n else if (statusIcon === '🔄') status = 'in_progress';\n else if (statusIcon === '🚫') status = 'pending'; // Map blocked to pending for display\n \n // Parse priority - match todo.py output format\n let priority: 'high' | 'medium' | 'low' | undefined;\n if (priorityIcon === '🔴') priority = 'high'; // critical -> high\n else if (priorityIcon === '🟠') priority = 'high'; // high\n else if (priorityIcon === '🟡') priority = 'medium'; // medium \n else if (priorityIcon === '🔵') priority = 'low'; // low\n \n // Try to extract description and dependencies\n const taskBlock = normalizedContent.substring(match.index);\n const nextTaskMatch = taskBlock.substring(match[0].length).match(/[✅⏳🔄🚫] Step [^\\s:]+:/);\n const taskContent = nextTaskMatch \n ? taskBlock.substring(0, match[0].length + (nextTaskMatch.index || 0))\n : taskBlock;\n \n // Extract description (lines starting with 📝)\n const descMatch = taskContent.match(/📝\\s*([^\\n]+)/);\n const description = descMatch ? descMatch[1].trim() : undefined;\n \n // Extract dependencies (🔗 Depends on: [x, y])\n const depsMatch = taskContent.match(/🔗\\s*Depends on:\\s*\\[([^\\]]+)\\]/);\n const dependsOn = depsMatch \n ? depsMatch[1].split(',').map(d => parseInt(d.trim(), 10)).filter(n => !isNaN(n))\n : undefined;\n \n // Extract time - handle ISO format timestamps\n const updatedMatch = taskContent.match(/updated_at:\\s*([\\d\\-T:]+)/);\n const createdMatch = taskContent.match(/created_at:\\s*([\\d\\-T:]+)/);\n \n items.push({\n step: step || items.length + 1, // Use parsed step or fallback to sequential number\n title,\n status,\n description,\n dependsOn,\n priority,\n updatedAt: updatedMatch ? updatedMatch[1].trim() : undefined,\n createdAt: createdMatch ? createdMatch[1].trim() : undefined,\n });\n }\n \n return items.sort((a, b) => a.step - b.step);\n}\n\n/**\n * Check if content is todo result\n */\nexport function isTodoResult(content: string): boolean {\n // Normalize content for checking\n const normalizedContent = content.replace(/\\\\n/g, '\\n');\n // Check for todo_result tags or characteristic patterns\n return normalizedContent.includes('') || \n (normalizedContent.includes('Step') && normalizedContent.includes('📊') && normalizedContent.includes('tasks'));\n}\n\n/**\n * Get icon and color for status\n */\nfunction getStatusConfig(status: TaskStatus) {\n switch (status) {\n case 'completed':\n return { icon: , color: '#52c41a', text: 'Completed' };\n case 'in_progress':\n return { icon: , color: '#1890ff', text: 'In Progress' };\n case 'pending':\n return { icon: , color: '#faad14', text: 'Pending' };\n default:\n return { icon: , color: '#d9d9d9', text: 'Not Started' };\n }\n}\n\n/**\n * Get priority tag\n */\nfunction getPriorityTag(priority?: 'high' | 'medium' | 'low') {\n if (!priority) return null;\n const config = {\n high: { color: 'red', text: 'High' },\n medium: { color: 'orange', text: 'Medium' },\n low: { color: 'green', text: 'Low' },\n };\n return {config[priority].text};\n}\n\n/**\n * Format time\n */\nfunction formatTime(timeStr?: string) {\n if (!timeStr) return null;\n try {\n // Ensure the time string is valid ISO format\n const cleanedTimeStr = timeStr.trim();\n const date = new Date(cleanedTimeStr);\n // Check if date is valid\n if (isNaN(date.getTime())) {\n return cleanedTimeStr; // Return original string if parsing fails\n }\n return date.toLocaleString('en-US', { \n month: '2-digit', \n day: '2-digit', \n hour: '2-digit', \n minute: '2-digit' \n });\n } catch {\n return timeStr;\n }\n}\n\n/**\n * TodoPanel main component\n */\nconst TodoPanel: React.FC = ({ \n items, \n title = 'Task Manager',\n agentName,\n timestamp,\n}) => {\n if (!items || items.length === 0) {\n return null;\n }\n \n // Task status statistics\n const completedCount = items.filter(i => i.status === 'completed').length;\n const inProgressCount = items.filter(i => i.status === 'in_progress').length;\n const totalCount = items.length;\n const progressPercent = Math.round((completedCount / totalCount) * 100);\n \n // Prepare Timeline items\n const timelineItems = items.map((item) => {\n const statusConfig = getStatusConfig(item.status);\n \n return {\n key: item.step,\n dot: statusConfig.icon,\n color: statusConfig.color,\n children: (\n
\n \n {/* Title row */}\n \n \n Step {item.step}: {item.title}\n \n \n {statusConfig.text}\n \n {getPriorityTag(item.priority)}\n \n \n {/* Description */}\n {item.description && (\n \n {item.description}\n \n )}\n \n {/* Meta info */}\n \n {item.dependsOn && item.dependsOn.length > 0 && (\n \n \n \n Deps: {item.dependsOn.join(', ')}\n \n \n )}\n {item.updatedAt && (\n \n \n \n {formatTime(item.updatedAt)}\n \n \n )}\n \n \n
\n ),\n };\n });\n\n return (\n \n
\n \n
\n
\n
\n {agentName || title}\n
\n
\n {completedCount}/{totalCount} tasks completed\n
\n
\n {inProgressCount > 0 && (\n }\n style={{ \n borderRadius: '12px',\n background: 'linear-gradient(135deg, #3b82f6 0%, #1d4ed8 100%)',\n border: 'none',\n color: 'white',\n fontSize: '11px',\n }}\n >\n {inProgressCount} active\n \n )}\n \n }\n >\n {/* Progress bar */}\n
\n (\n \n {progressPercent}%\n \n )}\n style={{ marginBottom: 0 }}\n />\n
\n \n {/* Task timeline */}\n \n {timestamp && (\n
\n {timestamp.toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })}\n
\n )}\n \n );\n};\n\nexport default TodoPanel;\n" }, { "path": "cortex-ui/src/components/WebPagePanel.tsx", "content": "/**\n * WebPagePanel Component\n * Display web page content from batch_open tool\n */\n\nimport React, { useState } from 'react';\nimport { Card, Collapse, Tag, Typography, Space, Tooltip, Button } from 'antd';\nimport {\n GlobalOutlined,\n LinkOutlined,\n FileTextOutlined,\n ExpandOutlined,\n CompressOutlined,\n CheckCircleOutlined,\n CloseCircleOutlined,\n CopyOutlined,\n ExportOutlined,\n} from '@ant-design/icons';\nimport { message } from 'antd';\nimport ReactMarkdown from 'react-markdown';\nimport remarkGfm from 'remark-gfm';\n\nconst { Text, Link } = Typography;\n\n// Web page result interface\nexport interface WebPageResult {\n url: string;\n title: string;\n snippet?: string;\n content?: string;\n host?: string;\n site?: string;\n success: boolean;\n error?: string;\n}\n\ninterface WebPagePanelProps {\n pages: WebPageResult[];\n title?: string;\n agentName?: string;\n timestamp?: Date;\n}\n\n/**\n * Single web page card component\n */\nconst WebPageCard: React.FC<{ page: WebPageResult; index: number }> = ({ page, index }) => {\n const [expanded, setExpanded] = useState(false);\n \n const handleCopyContent = () => {\n if (page.content) {\n navigator.clipboard.writeText(page.content);\n message.success('Content copied to clipboard');\n }\n };\n\n const handleOpenUrl = () => {\n window.open(page.url, '_blank');\n };\n\n // Process markdown content - unescape \\n to actual newlines\n const processedContent = page.content \n ? page.content\n .replace(/\\\\n/g, '\\n') // Convert \\n to actual newlines\n .replace(/\\\\t/g, '\\t') // Convert \\t to actual tabs\n .replace(/\\\\\\\\/g, '\\\\') // Convert \\\\ to \\\n : '';\n\n // Truncate content for preview\n const previewContent = processedContent \n ? (processedContent.length > 500 ? processedContent.substring(0, 500) + '...' : processedContent)\n : '';\n \n return (\n {\n e.currentTarget.style.borderColor = '#a78bfa';\n e.currentTarget.style.boxShadow = '0 4px 12px rgba(139, 92, 246, 0.1)';\n }}\n onMouseLeave={(e) => {\n e.currentTarget.style.borderColor = '#e5e7eb';\n e.currentTarget.style.boxShadow = 'none';\n }}\n >\n \n {/* Header: Title and status */}\n
\n
\n
\n \n #{index + 1}\n \n {page.success ? (\n }\n style={{ \n fontSize: '10px',\n borderRadius: '8px',\n background: 'linear-gradient(135deg, #22c55e 0%, #16a34a 100%)',\n border: 'none',\n color: 'white',\n }}\n >\n Loaded\n \n ) : (\n }\n style={{ \n fontSize: '10px',\n borderRadius: '8px',\n background: 'linear-gradient(135deg, #ef4444 0%, #dc2626 100%)',\n border: 'none',\n color: 'white',\n }}\n >\n Failed\n \n )}\n
\n \n {page.title || 'Untitled Page'}\n \n
\n \n {/* Action buttons */}\n \n {page.content && (\n \n }\n onClick={handleCopyContent}\n style={{ color: '#64748b' }}\n />\n \n )}\n \n }\n onClick={handleOpenUrl}\n style={{ color: '#64748b' }}\n />\n \n \n
\n \n {/* URL and site info */}\n
\n {page.site && (\n \n \n {page.site}\n \n )}\n \n \n {page.host || new URL(page.url).hostname}\n \n
\n \n {/* Snippet */}\n {page.snippet && (\n
\n {page.snippet}\n
\n )}\n \n {/* Content preview */}\n {page.content && (\n
\n
\n \n \n Page Content ({page.content.length.toLocaleString()} chars)\n \n : }\n onClick={() => setExpanded(!expanded)}\n style={{ fontSize: '11px', color: '#8b5cf6', padding: '0 4px' }}\n >\n {expanded ? 'Collapse' : 'Expand'}\n \n
\n \n

,\n h2: ({node, ...props}) =>

,\n h3: ({node, ...props}) =>

,\n h4: ({node, ...props}) =>

,\n p: ({node, ...props}) =>

,\n ul: ({node, ...props}) =>