Initial commit: Add QognyX front-end development code (optimized)
This commit is contained in:
commit
5c93561b3a
11
.gitignore
vendored
Normal file
11
.gitignore
vendored
Normal file
@ -0,0 +1,11 @@
|
||||
node_modules/
|
||||
.git/
|
||||
dist/
|
||||
build/
|
||||
.env.local
|
||||
.env.*.local
|
||||
logs/
|
||||
*.log
|
||||
.DS_Store
|
||||
.vscode/
|
||||
.idea/
|
||||
236
CHANGELOG.md
Normal file
236
CHANGELOG.md
Normal file
@ -0,0 +1,236 @@
|
||||
# Changelog
|
||||
|
||||
All notable changes to this project will be documented in this file.
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### ✨ New Features
|
||||
|
||||
- ✨ feat: implement search parameter updates by **@mawburn** in [#7151](https://github.com/danny-avila/QognyX/pull/7151)
|
||||
- 🎏 feat: Add MCP support for Streamable HTTP Transport by **@benverhees** in [#7353](https://github.com/danny-avila/QognyX/pull/7353)
|
||||
- 🔒 feat: Add Content Security Policy using Helmet middleware by **@rubentalstra** in [#7377](https://github.com/danny-avila/QognyX/pull/7377)
|
||||
- ✨ feat: Add Normalization for MCP Server Names by **@danny-avila** in [#7421](https://github.com/danny-avila/QognyX/pull/7421)
|
||||
- 📊 feat: Improve Helm Chart by **@hofq** in [#3638](https://github.com/danny-avila/QognyX/pull/3638)
|
||||
- 🦾 feat: Claude-4 Support by **@danny-avila** in [#7509](https://github.com/danny-avila/QognyX/pull/7509)
|
||||
- 🪨 feat: Bedrock Support for Claude-4 Reasoning by **@danny-avila** in [#7517](https://github.com/danny-avila/QognyX/pull/7517)
|
||||
|
||||
### 🌍 Internationalization
|
||||
|
||||
- 🌍 i18n: Add `Danish` and `Czech` and `Catalan` localization support by **@rubentalstra** in [#7373](https://github.com/danny-avila/QognyX/pull/7373)
|
||||
- 🌍 i18n: Update translation.json with latest translations by **@github-actions[bot]** in [#7375](https://github.com/danny-avila/QognyX/pull/7375)
|
||||
- 🌍 i18n: Update translation.json with latest translations by **@github-actions[bot]** in [#7468](https://github.com/danny-avila/QognyX/pull/7468)
|
||||
|
||||
### 🔧 Fixes
|
||||
|
||||
- 💬 fix: update aria-label for accessibility in ConvoLink component by **@berry-13** in [#7320](https://github.com/danny-avila/QognyX/pull/7320)
|
||||
- 🔑 fix: use `apiKey` instead of `openAIApiKey` in OpenAI-like Config by **@danny-avila** in [#7337](https://github.com/danny-avila/QognyX/pull/7337)
|
||||
- 🔄 fix: update navigation logic in `useFocusChatEffect` to ensure correct search parameters are used by **@mawburn** in [#7340](https://github.com/danny-avila/QognyX/pull/7340)
|
||||
- 🔄 fix: Improve MCP Connection Cleanup by **@danny-avila** in [#7400](https://github.com/danny-avila/QognyX/pull/7400)
|
||||
- 🛡️ fix: Preset and Validation Logic for URL Query Params by **@danny-avila** in [#7407](https://github.com/danny-avila/QognyX/pull/7407)
|
||||
- 🌘 fix: artifact of preview text is illegible in dark mode by **@nhtruong** in [#7405](https://github.com/danny-avila/QognyX/pull/7405)
|
||||
- 🛡️ fix: Temporarily Remove CSP until Configurable by **@danny-avila** in [#7419](https://github.com/danny-avila/QognyX/pull/7419)
|
||||
- 💽 fix: Exclude index page `/` from static cache settings by **@sbruel** in [#7382](https://github.com/danny-avila/QognyX/pull/7382)
|
||||
|
||||
### ⚙️ Other Changes
|
||||
|
||||
- 📜 docs: CHANGELOG for release v0.7.8 by **@github-actions[bot]** in [#7290](https://github.com/danny-avila/QognyX/pull/7290)
|
||||
- 📦 chore: Update API Package Dependencies by **@danny-avila** in [#7359](https://github.com/danny-avila/QognyX/pull/7359)
|
||||
- 📜 docs: Unreleased Changelog by **@github-actions[bot]** in [#7321](https://github.com/danny-avila/QognyX/pull/7321)
|
||||
- 📜 docs: Unreleased Changelog by **@github-actions[bot]** in [#7434](https://github.com/danny-avila/QognyX/pull/7434)
|
||||
- 🛡️ chore: `multer` v2.0.0 for CVE-2025-47935 and CVE-2025-47944 by **@danny-avila** in [#7454](https://github.com/danny-avila/QognyX/pull/7454)
|
||||
- 📂 refactor: Improve `FileAttachment` & File Form Deletion by **@danny-avila** in [#7471](https://github.com/danny-avila/QognyX/pull/7471)
|
||||
- 📊 chore: Remove Old Helm Chart by **@hofq** in [#7512](https://github.com/danny-avila/QognyX/pull/7512)
|
||||
- 🪖 chore: bump helm app version to v0.7.8 by **@austin-barrington** in [#7524](https://github.com/danny-avila/QognyX/pull/7524)
|
||||
|
||||
|
||||
|
||||
---
|
||||
## [v0.7.8] -
|
||||
|
||||
Changes from v0.7.8-rc1 to v0.7.8.
|
||||
|
||||
### ✨ New Features
|
||||
|
||||
- ✨ feat: Enhance form submission for touch screens by **@berry-13** in [#7198](https://github.com/danny-avila/QognyX/pull/7198)
|
||||
- 🔍 feat: Additional Tavily API Tool Parameters by **@glowforge-opensource** in [#7232](https://github.com/danny-avila/QognyX/pull/7232)
|
||||
- 🐋 feat: Add python to Dockerfile for increased MCP compatibility by **@technicalpickles** in [#7270](https://github.com/danny-avila/QognyX/pull/7270)
|
||||
|
||||
### 🔧 Fixes
|
||||
|
||||
- 🔧 fix: Google Gemma Support & OpenAI Reasoning Instructions by **@danny-avila** in [#7196](https://github.com/danny-avila/QognyX/pull/7196)
|
||||
- 🛠️ fix: Conversation Navigation State by **@danny-avila** in [#7210](https://github.com/danny-avila/QognyX/pull/7210)
|
||||
- 🔄 fix: o-Series Model Regex for System Messages by **@danny-avila** in [#7245](https://github.com/danny-avila/QognyX/pull/7245)
|
||||
- 🔖 fix: Custom Headers for Initial MCP SSE Connection by **@danny-avila** in [#7246](https://github.com/danny-avila/QognyX/pull/7246)
|
||||
- 🛡️ fix: Deep Clone `MCPOptions` for User MCP Connections by **@danny-avila** in [#7247](https://github.com/danny-avila/QognyX/pull/7247)
|
||||
- 🔄 fix: URL Param Race Condition and File Draft Persistence by **@danny-avila** in [#7257](https://github.com/danny-avila/QognyX/pull/7257)
|
||||
- 🔄 fix: Assistants Endpoint & Minor Issues by **@danny-avila** in [#7274](https://github.com/danny-avila/QognyX/pull/7274)
|
||||
- 🔄 fix: Ollama Think Tag Edge Case with Tools by **@danny-avila** in [#7275](https://github.com/danny-avila/QognyX/pull/7275)
|
||||
|
||||
### ⚙️ Other Changes
|
||||
|
||||
- 📜 docs: CHANGELOG for release v0.7.8-rc1 by **@github-actions[bot]** in [#7153](https://github.com/danny-avila/QognyX/pull/7153)
|
||||
- 🔄 refactor: Artifact Visibility Management by **@danny-avila** in [#7181](https://github.com/danny-avila/QognyX/pull/7181)
|
||||
- 📦 chore: Bump Package Security by **@danny-avila** in [#7183](https://github.com/danny-avila/QognyX/pull/7183)
|
||||
- 🌿 refactor: Unmount Fork Popover on Hide for Better Performance by **@danny-avila** in [#7189](https://github.com/danny-avila/QognyX/pull/7189)
|
||||
- 🧰 chore: ESLint configuration to enforce Prettier formatting rules by **@mawburn** in [#7186](https://github.com/danny-avila/QognyX/pull/7186)
|
||||
- 🎨 style: Improve KaTeX Rendering for LaTeX Equations by **@andresgit** in [#7223](https://github.com/danny-avila/QognyX/pull/7223)
|
||||
- 📝 docs: Update `.env.example` Google models by **@marlonka** in [#7254](https://github.com/danny-avila/QognyX/pull/7254)
|
||||
- 💬 refactor: MCP Chat Visibility Option, Google Rates, Remove OpenAPI Plugins by **@danny-avila** in [#7286](https://github.com/danny-avila/QognyX/pull/7286)
|
||||
- 📜 docs: Unreleased Changelog by **@github-actions[bot]** in [#7214](https://github.com/danny-avila/QognyX/pull/7214)
|
||||
|
||||
|
||||
|
||||
[See full release details][release-v0.7.8]
|
||||
|
||||
[release-v0.7.8]: https://github.com/danny-avila/QognyX/releases/tag/v0.7.8
|
||||
|
||||
---
|
||||
## [v0.7.8-rc1] -
|
||||
|
||||
Changes from v0.7.7 to v0.7.8-rc1.
|
||||
|
||||
### ✨ New Features
|
||||
|
||||
- 🔍 feat: Mistral OCR API / Upload Files as Text by **@danny-avila** in [#6274](https://github.com/danny-avila/QognyX/pull/6274)
|
||||
- 🤖 feat: Support OpenAI Web Search models by **@danny-avila** in [#6313](https://github.com/danny-avila/QognyX/pull/6313)
|
||||
- 🔗 feat: Agent Chain (Mixture-of-Agents) by **@danny-avila** in [#6374](https://github.com/danny-avila/QognyX/pull/6374)
|
||||
- ⌛ feat: `initTimeout` for Slow Starting MCP Servers by **@perweij** in [#6383](https://github.com/danny-avila/QognyX/pull/6383)
|
||||
- 🚀 feat: `S3` Integration for File handling and Image uploads by **@rubentalstra** in [#6142](https://github.com/danny-avila/QognyX/pull/6142)
|
||||
- 🔒feat: Enable OpenID Auto-Redirect by **@leondape** in [#6066](https://github.com/danny-avila/QognyX/pull/6066)
|
||||
- 🚀 feat: Integrate `Azure Blob Storage` for file handling and image uploads by **@rubentalstra** in [#6153](https://github.com/danny-avila/QognyX/pull/6153)
|
||||
- 🚀 feat: Add support for custom `AWS` endpoint in `S3` by **@rubentalstra** in [#6431](https://github.com/danny-avila/QognyX/pull/6431)
|
||||
- 🚀 feat: Add support for LDAP STARTTLS in LDAP authentication by **@rubentalstra** in [#6438](https://github.com/danny-avila/QognyX/pull/6438)
|
||||
- 🚀 feat: Refactor schema exports and update package version to 0.0.4 by **@rubentalstra** in [#6455](https://github.com/danny-avila/QognyX/pull/6455)
|
||||
- 🔼 feat: Add Auto Submit For URL Query Params by **@mjaverto** in [#6440](https://github.com/danny-avila/QognyX/pull/6440)
|
||||
- 🛠 feat: Enhance Redis Integration, Rate Limiters & Log Headers by **@danny-avila** in [#6462](https://github.com/danny-avila/QognyX/pull/6462)
|
||||
- 💵 feat: Add Automatic Balance Refill by **@rubentalstra** in [#6452](https://github.com/danny-avila/QognyX/pull/6452)
|
||||
- 🗣️ feat: add support for gpt-4o-transcribe models by **@berry-13** in [#6483](https://github.com/danny-avila/QognyX/pull/6483)
|
||||
- 🎨 feat: UI Refresh for Enhanced UX by **@berry-13** in [#6346](https://github.com/danny-avila/QognyX/pull/6346)
|
||||
- 🌍 feat: Add support for Hungarian language localization by **@rubentalstra** in [#6508](https://github.com/danny-avila/QognyX/pull/6508)
|
||||
- 🚀 feat: Add Gemini 2.5 Token/Context Values, Increase Max Possible Output to 64k by **@danny-avila** in [#6563](https://github.com/danny-avila/QognyX/pull/6563)
|
||||
- 🚀 feat: Enhance MCP Connections For Multi-User Support by **@danny-avila** in [#6610](https://github.com/danny-avila/QognyX/pull/6610)
|
||||
- 🚀 feat: Enhance S3 URL Expiry with Refresh; fix: S3 File Deletion by **@danny-avila** in [#6647](https://github.com/danny-avila/QognyX/pull/6647)
|
||||
- 🚀 feat: enhance UI components and refactor settings by **@berry-13** in [#6625](https://github.com/danny-avila/QognyX/pull/6625)
|
||||
- 💬 feat: move TemporaryChat to the Header by **@berry-13** in [#6646](https://github.com/danny-avila/QognyX/pull/6646)
|
||||
- 🚀 feat: Use Model Specs + Specific Endpoints, Limit Providers for Agents by **@danny-avila** in [#6650](https://github.com/danny-avila/QognyX/pull/6650)
|
||||
- 🪙 feat: Sync Balance Config on Login by **@danny-avila** in [#6671](https://github.com/danny-avila/QognyX/pull/6671)
|
||||
- 🔦 feat: MCP Support for Non-Agent Endpoints by **@danny-avila** in [#6775](https://github.com/danny-avila/QognyX/pull/6775)
|
||||
- 🗃️ feat: Code Interpreter File Persistence between Sessions by **@danny-avila** in [#6790](https://github.com/danny-avila/QognyX/pull/6790)
|
||||
- 🖥️ feat: Code Interpreter API for Non-Agent Endpoints by **@danny-avila** in [#6803](https://github.com/danny-avila/QognyX/pull/6803)
|
||||
- ⚡ feat: Self-hosted Artifacts Static Bundler URL by **@danny-avila** in [#6827](https://github.com/danny-avila/QognyX/pull/6827)
|
||||
- 🐳 feat: Add Jemalloc and UV to Docker Builds by **@danny-avila** in [#6836](https://github.com/danny-avila/QognyX/pull/6836)
|
||||
- 🤖 feat: GPT-4.1 by **@danny-avila** in [#6880](https://github.com/danny-avila/QognyX/pull/6880)
|
||||
- 👋 feat: remove Edge TTS by **@berry-13** in [#6885](https://github.com/danny-avila/QognyX/pull/6885)
|
||||
- feat: nav optimization by **@berry-13** in [#5785](https://github.com/danny-avila/QognyX/pull/5785)
|
||||
- 🗺️ feat: Add Parameter Location Mapping for OpenAPI actions by **@peeeteeer** in [#6858](https://github.com/danny-avila/QognyX/pull/6858)
|
||||
- 🤖 feat: Support `o4-mini` and `o3` Models by **@danny-avila** in [#6928](https://github.com/danny-avila/QognyX/pull/6928)
|
||||
- 🎨 feat: OpenAI Image Tools (GPT-Image-1) by **@danny-avila** in [#7079](https://github.com/danny-avila/QognyX/pull/7079)
|
||||
- 🗓️ feat: Add Special Variables for Prompts & Agents, Prompt UI Improvements by **@danny-avila** in [#7123](https://github.com/danny-avila/QognyX/pull/7123)
|
||||
|
||||
### 🌍 Internationalization
|
||||
|
||||
- 🌍 i18n: Add Thai Language Support and Update Translations by **@rubentalstra** in [#6219](https://github.com/danny-avila/QognyX/pull/6219)
|
||||
- 🌍 i18n: Update translation.json with latest translations by **@github-actions[bot]** in [#6220](https://github.com/danny-avila/QognyX/pull/6220)
|
||||
- 🌍 i18n: Update translation.json with latest translations by **@github-actions[bot]** in [#6240](https://github.com/danny-avila/QognyX/pull/6240)
|
||||
- 🌍 i18n: Update translation.json with latest translations by **@github-actions[bot]** in [#6241](https://github.com/danny-avila/QognyX/pull/6241)
|
||||
- 🌍 i18n: Update translation.json with latest translations by **@github-actions[bot]** in [#6277](https://github.com/danny-avila/QognyX/pull/6277)
|
||||
- 🌍 i18n: Update translation.json with latest translations by **@github-actions[bot]** in [#6414](https://github.com/danny-avila/QognyX/pull/6414)
|
||||
- 🌍 i18n: Update translation.json with latest translations by **@github-actions[bot]** in [#6505](https://github.com/danny-avila/QognyX/pull/6505)
|
||||
- 🌍 i18n: Update translation.json with latest translations by **@github-actions[bot]** in [#6530](https://github.com/danny-avila/QognyX/pull/6530)
|
||||
- 🌍 i18n: Add Persian Localization Support by **@rubentalstra** in [#6669](https://github.com/danny-avila/QognyX/pull/6669)
|
||||
- 🌍 i18n: Update translation.json with latest translations by **@github-actions[bot]** in [#6667](https://github.com/danny-avila/QognyX/pull/6667)
|
||||
- 🌍 i18n: Update translation.json with latest translations by **@github-actions[bot]** in [#7126](https://github.com/danny-avila/QognyX/pull/7126)
|
||||
- 🌍 i18n: Update translation.json with latest translations by **@github-actions[bot]** in [#7148](https://github.com/danny-avila/QognyX/pull/7148)
|
||||
|
||||
### 👐 Accessibility
|
||||
|
||||
- 🎨 a11y: Update Model Spec Description Text by **@berry-13** in [#6294](https://github.com/danny-avila/QognyX/pull/6294)
|
||||
- 🗑️ a11y: Add Accessible Name to Button for File Attachment Removal by **@kangabell** in [#6709](https://github.com/danny-avila/QognyX/pull/6709)
|
||||
- ⌨️ a11y: enhance accessibility & visual consistency by **@berry-13** in [#6866](https://github.com/danny-avila/QognyX/pull/6866)
|
||||
- 🙌 a11y: Searchbar/Conversations List Focus by **@danny-avila** in [#7096](https://github.com/danny-avila/QognyX/pull/7096)
|
||||
- 👐 a11y: Improve Fork and SplitText Accessibility by **@danny-avila** in [#7147](https://github.com/danny-avila/QognyX/pull/7147)
|
||||
|
||||
### 🔧 Fixes
|
||||
|
||||
- 🐛 fix: Avatar Type Definitions in Agent/Assistant Schemas by **@danny-avila** in [#6235](https://github.com/danny-avila/QognyX/pull/6235)
|
||||
- 🔧 fix: MeiliSearch Field Error and Patch Incorrect Import by #6210 by **@rubentalstra** in [#6245](https://github.com/danny-avila/QognyX/pull/6245)
|
||||
- 🔏 fix: Enhance Two-Factor Authentication by **@rubentalstra** in [#6247](https://github.com/danny-avila/QognyX/pull/6247)
|
||||
- 🐛 fix: Await saveMessage in abortMiddleware to ensure proper execution by **@sh4shii** in [#6248](https://github.com/danny-avila/QognyX/pull/6248)
|
||||
- 🔧 fix: Axios Proxy Usage And Bump `mongoose` by **@danny-avila** in [#6298](https://github.com/danny-avila/QognyX/pull/6298)
|
||||
- 🔧 fix: comment out MCP servers to resolve service run issues by **@KunalScriptz** in [#6316](https://github.com/danny-avila/QognyX/pull/6316)
|
||||
- 🔧 fix: Update Token Calculations and Mapping, MCP `env` Initialization by **@danny-avila** in [#6406](https://github.com/danny-avila/QognyX/pull/6406)
|
||||
- 🐞 fix: Agent "Resend" Message Attachments + Source Icon Styling by **@danny-avila** in [#6408](https://github.com/danny-avila/QognyX/pull/6408)
|
||||
- 🐛 fix: Prevent Crash on Duplicate Message ID by **@Odrec** in [#6392](https://github.com/danny-avila/QognyX/pull/6392)
|
||||
- 🔐 fix: Invalid Key Length in 2FA Encryption by **@rubentalstra** in [#6432](https://github.com/danny-avila/QognyX/pull/6432)
|
||||
- 🏗️ fix: Fix Agents Token Spend Race Conditions, Expand Test Coverage by **@danny-avila** in [#6480](https://github.com/danny-avila/QognyX/pull/6480)
|
||||
- 🔃 fix: Draft Clearing, Claude Titles, Remove Default Vision Max Tokens by **@danny-avila** in [#6501](https://github.com/danny-avila/QognyX/pull/6501)
|
||||
- 🔧 fix: Update username reference to use user.name in greeting display by **@rubentalstra** in [#6534](https://github.com/danny-avila/QognyX/pull/6534)
|
||||
- 🔧 fix: S3 Download Stream with Key Extraction and Blob Storage Encoding for Vision by **@danny-avila** in [#6557](https://github.com/danny-avila/QognyX/pull/6557)
|
||||
- 🔧 fix: Mistral type strictness for `usage` & update token values/windows by **@danny-avila** in [#6562](https://github.com/danny-avila/QognyX/pull/6562)
|
||||
- 🔧 fix: Consolidate Text Parsing and TTS Edge Initialization by **@danny-avila** in [#6582](https://github.com/danny-avila/QognyX/pull/6582)
|
||||
- 🔧 fix: Ensure continuation in image processing on base64 encoding from Blob Storage by **@danny-avila** in [#6619](https://github.com/danny-avila/QognyX/pull/6619)
|
||||
- ✉️ fix: Fallback For User Name In Email Templates by **@danny-avila** in [#6620](https://github.com/danny-avila/QognyX/pull/6620)
|
||||
- 🔧 fix: Azure Blob Integration and File Source References by **@rubentalstra** in [#6575](https://github.com/danny-avila/QognyX/pull/6575)
|
||||
- 🐛 fix: Safeguard against undefined addedEndpoints by **@wipash** in [#6654](https://github.com/danny-avila/QognyX/pull/6654)
|
||||
- 🤖 fix: Gemini 2.5 Vision Support by **@danny-avila** in [#6663](https://github.com/danny-avila/QognyX/pull/6663)
|
||||
- 🔄 fix: Avatar & Error Handling Enhancements by **@danny-avila** in [#6687](https://github.com/danny-avila/QognyX/pull/6687)
|
||||
- 🔧 fix: Chat Middleware, Zod Conversion, Auto-Save and S3 URL Refresh by **@danny-avila** in [#6720](https://github.com/danny-avila/QognyX/pull/6720)
|
||||
- 🔧 fix: Agent Capability Checks & DocumentDB Compatibility for Agent Resource Removal by **@danny-avila** in [#6726](https://github.com/danny-avila/QognyX/pull/6726)
|
||||
- 🔄 fix: Improve audio MIME type detection and handling by **@berry-13** in [#6707](https://github.com/danny-avila/QognyX/pull/6707)
|
||||
- 🪺 fix: Update Role Handling due to New Schema Shape by **@danny-avila** in [#6774](https://github.com/danny-avila/QognyX/pull/6774)
|
||||
- 🗨️ fix: Show ModelSpec Greeting by **@berry-13** in [#6770](https://github.com/danny-avila/QognyX/pull/6770)
|
||||
- 🔧 fix: Keyv and Proxy Issues, and More Memory Optimizations by **@danny-avila** in [#6867](https://github.com/danny-avila/QognyX/pull/6867)
|
||||
- ✨ fix: Implement dynamic text sizing for greeting and name display by **@berry-13** in [#6833](https://github.com/danny-avila/QognyX/pull/6833)
|
||||
- 📝 fix: Mistral OCR Image Support and Azure Agent Titles by **@danny-avila** in [#6901](https://github.com/danny-avila/QognyX/pull/6901)
|
||||
- 📢 fix: Invalid `engineTTS` and Conversation State on Navigation by **@berry-13** in [#6904](https://github.com/danny-avila/QognyX/pull/6904)
|
||||
- 🛠️ fix: Improve Accessibility and Display of Conversation Menu by **@danny-avila** in [#6913](https://github.com/danny-avila/QognyX/pull/6913)
|
||||
- 🔧 fix: Agent Resource Form, Convo Menu Style, Ensure Draft Clears on Submission by **@danny-avila** in [#6925](https://github.com/danny-avila/QognyX/pull/6925)
|
||||
- 🔀 fix: MCP Improvements, Auto-Save Drafts, Artifact Markup by **@danny-avila** in [#7040](https://github.com/danny-avila/QognyX/pull/7040)
|
||||
- 🐋 fix: Improve Deepseek Compatbility by **@danny-avila** in [#7132](https://github.com/danny-avila/QognyX/pull/7132)
|
||||
- 🐙 fix: Add Redis Ping Interval to Prevent Connection Drops by **@peeeteeer** in [#7127](https://github.com/danny-avila/QognyX/pull/7127)
|
||||
|
||||
### ⚙️ Other Changes
|
||||
|
||||
- 📦 refactor: Move DB Models to `@librechat/data-schemas` by **@rubentalstra** in [#6210](https://github.com/danny-avila/QognyX/pull/6210)
|
||||
- 📦 chore: Patch `axios` to address CVE-2025-27152 by **@danny-avila** in [#6222](https://github.com/danny-avila/QognyX/pull/6222)
|
||||
- ⚠️ refactor: Use Error Content Part Instead Of Throwing Error for Agents by **@danny-avila** in [#6262](https://github.com/danny-avila/QognyX/pull/6262)
|
||||
- 🏃♂️ refactor: Improve Agent Run Context & Misc. Changes by **@danny-avila** in [#6448](https://github.com/danny-avila/QognyX/pull/6448)
|
||||
- 📝 docs: librechat.example.yaml by **@ineiti** in [#6442](https://github.com/danny-avila/QognyX/pull/6442)
|
||||
- 🏃♂️ refactor: More Agent Context Improvements during Run by **@danny-avila** in [#6477](https://github.com/danny-avila/QognyX/pull/6477)
|
||||
- 🔃 refactor: Allow streaming for `o1` models by **@danny-avila** in [#6509](https://github.com/danny-avila/QognyX/pull/6509)
|
||||
- 🔧 chore: `Vite` Plugin Upgrades & Config Optimizations by **@rubentalstra** in [#6547](https://github.com/danny-avila/QognyX/pull/6547)
|
||||
- 🔧 refactor: Consolidate Logging, Model Selection & Actions Optimizations, Minor Fixes by **@danny-avila** in [#6553](https://github.com/danny-avila/QognyX/pull/6553)
|
||||
- 🎨 style: Address Minor UI Refresh Issues by **@berry-13** in [#6552](https://github.com/danny-avila/QognyX/pull/6552)
|
||||
- 🔧 refactor: Enhance Model & Endpoint Configurations with Global Indicators 🌍 by **@berry-13** in [#6578](https://github.com/danny-avila/QognyX/pull/6578)
|
||||
- 💬 style: Chat UI, Greeting, and Message adjustments by **@berry-13** in [#6612](https://github.com/danny-avila/QognyX/pull/6612)
|
||||
- ⚡ refactor: DocumentDB Compatibility for Balance Updates by **@danny-avila** in [#6673](https://github.com/danny-avila/QognyX/pull/6673)
|
||||
- 🧹 chore: Update ESLint rules for React hooks by **@rubentalstra** in [#6685](https://github.com/danny-avila/QognyX/pull/6685)
|
||||
- 🪙 chore: Update Gemini Pricing by **@RedwindA** in [#6731](https://github.com/danny-avila/QognyX/pull/6731)
|
||||
- 🪺 refactor: Nest Permission fields for Roles by **@rubentalstra** in [#6487](https://github.com/danny-avila/QognyX/pull/6487)
|
||||
- 📦 chore: Update `caniuse-lite` dependency to version 1.0.30001706 by **@rubentalstra** in [#6482](https://github.com/danny-avila/QognyX/pull/6482)
|
||||
- ⚙️ refactor: OAuth Flow Signal, Type Safety, Tool Progress & Updated Packages by **@danny-avila** in [#6752](https://github.com/danny-avila/QognyX/pull/6752)
|
||||
- 📦 chore: bump vite from 6.2.3 to 6.2.5 by **@dependabot[bot]** in [#6745](https://github.com/danny-avila/QognyX/pull/6745)
|
||||
- 💾 chore: Enhance Local Storage Handling and Update MCP SDK by **@danny-avila** in [#6809](https://github.com/danny-avila/QognyX/pull/6809)
|
||||
- 🤖 refactor: Improve Agents Memory Usage, Bump Keyv, Grok 3 by **@danny-avila** in [#6850](https://github.com/danny-avila/QognyX/pull/6850)
|
||||
- 💾 refactor: Enhance Memory In Image Encodings & Client Disposal by **@danny-avila** in [#6852](https://github.com/danny-avila/QognyX/pull/6852)
|
||||
- 🔁 refactor: Token Event Handler and Standardize `maxTokens` Key by **@danny-avila** in [#6886](https://github.com/danny-avila/QognyX/pull/6886)
|
||||
- 🔍 refactor: Search & Message Retrieval by **@berry-13** in [#6903](https://github.com/danny-avila/QognyX/pull/6903)
|
||||
- 🎨 style: standardize dropdown styling & fix z-Index layering by **@berry-13** in [#6939](https://github.com/danny-avila/QognyX/pull/6939)
|
||||
- 📙 docs: CONTRIBUTING.md by **@dblock** in [#6831](https://github.com/danny-avila/QognyX/pull/6831)
|
||||
- 🧭 refactor: Modernize Nav/Header by **@danny-avila** in [#7094](https://github.com/danny-avila/QognyX/pull/7094)
|
||||
- 🪶 refactor: Chat Input Focus for Conversation Navigations & ChatForm Optimizations by **@danny-avila** in [#7100](https://github.com/danny-avila/QognyX/pull/7100)
|
||||
- 🔃 refactor: Streamline Navigation, Message Loading UX by **@danny-avila** in [#7118](https://github.com/danny-avila/QognyX/pull/7118)
|
||||
- 📜 docs: Unreleased changelog by **@github-actions[bot]** in [#6265](https://github.com/danny-avila/QognyX/pull/6265)
|
||||
|
||||
|
||||
|
||||
[See full release details][release-v0.7.8-rc1]
|
||||
|
||||
[release-v0.7.8-rc1]: https://github.com/danny-avila/QognyX/releases/tag/v0.7.8-rc1
|
||||
|
||||
---
|
||||
60
Dockerfile
Normal file
60
Dockerfile
Normal file
@ -0,0 +1,60 @@
|
||||
# v0.8.2-rc3
|
||||
|
||||
# Base node image
|
||||
FROM node:20-alpine AS node
|
||||
|
||||
# Install jemalloc
|
||||
RUN apk add --no-cache jemalloc
|
||||
RUN apk add --no-cache python3 py3-pip uv
|
||||
|
||||
# Set environment variable to use jemalloc
|
||||
ENV LD_PRELOAD=/usr/lib/libjemalloc.so.2
|
||||
|
||||
# Add `uv` for extended MCP support
|
||||
COPY --from=ghcr.io/astral-sh/uv:0.9.5-python3.12-alpine /usr/local/bin/uv /usr/local/bin/uvx /bin/
|
||||
RUN uv --version
|
||||
|
||||
# Set configurable max-old-space-size with default
|
||||
ARG NODE_MAX_OLD_SPACE_SIZE=6144
|
||||
|
||||
RUN mkdir -p /app && chown node:node /app
|
||||
WORKDIR /app
|
||||
|
||||
USER node
|
||||
|
||||
COPY --chown=node:node package.json package-lock.json ./
|
||||
COPY --chown=node:node api/package.json ./api/package.json
|
||||
COPY --chown=node:node client/package.json ./client/package.json
|
||||
COPY --chown=node:node packages/data-provider/package.json ./packages/data-provider/package.json
|
||||
COPY --chown=node:node packages/data-schemas/package.json ./packages/data-schemas/package.json
|
||||
COPY --chown=node:node packages/api/package.json ./packages/api/package.json
|
||||
|
||||
RUN \
|
||||
# Allow mounting of these files, which have no default
|
||||
touch .env ; \
|
||||
# Create directories for the volumes to inherit the correct permissions
|
||||
mkdir -p /app/client/public/images /app/logs /app/uploads ; \
|
||||
npm config set fetch-retry-maxtimeout 600000 ; \
|
||||
npm config set fetch-retries 5 ; \
|
||||
npm config set fetch-retry-mintimeout 15000 ; \
|
||||
npm ci --no-audit
|
||||
|
||||
COPY --chown=node:node . .
|
||||
|
||||
RUN \
|
||||
# React client build with configurable memory
|
||||
NODE_OPTIONS="--max-old-space-size=${NODE_MAX_OLD_SPACE_SIZE}" npm run frontend; \
|
||||
npm prune --production; \
|
||||
npm cache clean --force
|
||||
|
||||
# Node API setup
|
||||
EXPOSE 3080
|
||||
ENV HOST=0.0.0.0
|
||||
CMD ["npm", "run", "backend"]
|
||||
|
||||
# Optional: for client with nginx routing
|
||||
# FROM nginx:stable-alpine AS nginx-client
|
||||
# WORKDIR /usr/share/nginx/html
|
||||
# COPY --from=node /app/client/dist /usr/share/nginx/html
|
||||
# COPY client/nginx.conf /etc/nginx/conf.d/default.conf
|
||||
# ENTRYPOINT ["nginx", "-g", "daemon off;"]
|
||||
87
Dockerfile.multi
Normal file
87
Dockerfile.multi
Normal file
@ -0,0 +1,87 @@
|
||||
# Dockerfile.multi
|
||||
# v0.8.2-rc3
|
||||
|
||||
# Set configurable max-old-space-size with default
|
||||
ARG NODE_MAX_OLD_SPACE_SIZE=6144
|
||||
|
||||
# Base for all builds
|
||||
FROM node:20-alpine AS base-min
|
||||
# Install jemalloc
|
||||
RUN apk add --no-cache jemalloc
|
||||
# Set environment variable to use jemalloc
|
||||
ENV LD_PRELOAD=/usr/lib/libjemalloc.so.2
|
||||
|
||||
WORKDIR /app
|
||||
RUN apk --no-cache add curl
|
||||
RUN npm config set fetch-retry-maxtimeout 600000 && \
|
||||
npm config set fetch-retries 5 && \
|
||||
npm config set fetch-retry-mintimeout 15000
|
||||
COPY package*.json ./
|
||||
COPY packages/data-provider/package*.json ./packages/data-provider/
|
||||
COPY packages/api/package*.json ./packages/api/
|
||||
COPY packages/data-schemas/package*.json ./packages/data-schemas/
|
||||
COPY packages/client/package*.json ./packages/client/
|
||||
COPY client/package*.json ./client/
|
||||
COPY api/package*.json ./api/
|
||||
|
||||
# Install all dependencies for every build
|
||||
FROM base-min AS base
|
||||
WORKDIR /app
|
||||
RUN npm ci
|
||||
|
||||
# Build `data-provider` package
|
||||
FROM base AS data-provider-build
|
||||
WORKDIR /app/packages/data-provider
|
||||
COPY packages/data-provider ./
|
||||
RUN npm run build
|
||||
|
||||
# Build `data-schemas` package
|
||||
FROM base AS data-schemas-build
|
||||
WORKDIR /app/packages/data-schemas
|
||||
COPY packages/data-schemas ./
|
||||
COPY --from=data-provider-build /app/packages/data-provider/dist /app/packages/data-provider/dist
|
||||
RUN npm run build
|
||||
|
||||
# Build `api` package
|
||||
FROM base AS api-package-build
|
||||
WORKDIR /app/packages/api
|
||||
COPY packages/api ./
|
||||
COPY --from=data-provider-build /app/packages/data-provider/dist /app/packages/data-provider/dist
|
||||
COPY --from=data-schemas-build /app/packages/data-schemas/dist /app/packages/data-schemas/dist
|
||||
RUN npm run build
|
||||
|
||||
# Build `client` package
|
||||
FROM base AS client-package-build
|
||||
WORKDIR /app/packages/client
|
||||
COPY packages/client ./
|
||||
RUN npm run build
|
||||
|
||||
# Client build
|
||||
FROM base AS client-build
|
||||
WORKDIR /app/client
|
||||
COPY client ./
|
||||
COPY --from=data-provider-build /app/packages/data-provider/dist /app/packages/data-provider/dist
|
||||
COPY --from=client-package-build /app/packages/client/dist /app/packages/client/dist
|
||||
COPY --from=client-package-build /app/packages/client/src /app/packages/client/src
|
||||
ARG NODE_MAX_OLD_SPACE_SIZE
|
||||
ENV NODE_OPTIONS="--max-old-space-size=${NODE_MAX_OLD_SPACE_SIZE}"
|
||||
RUN npm run build
|
||||
|
||||
# API setup (including client dist)
|
||||
FROM base-min AS api-build
|
||||
# Add `uv` for extended MCP support
|
||||
COPY --from=ghcr.io/astral-sh/uv:0.6.13 /uv /uvx /bin/
|
||||
RUN uv --version
|
||||
WORKDIR /app
|
||||
# Install only production deps
|
||||
RUN npm ci --omit=dev
|
||||
COPY api ./api
|
||||
COPY config ./config
|
||||
COPY --from=data-provider-build /app/packages/data-provider/dist ./packages/data-provider/dist
|
||||
COPY --from=data-schemas-build /app/packages/data-schemas/dist ./packages/data-schemas/dist
|
||||
COPY --from=api-package-build /app/packages/api/dist ./packages/api/dist
|
||||
COPY --from=client-build /app/client/dist ./client/dist
|
||||
WORKDIR /app/api
|
||||
EXPOSE 3080
|
||||
ENV HOST=0.0.0.0
|
||||
CMD ["node", "server/index.js"]
|
||||
151
INSTALLATION.md
Normal file
151
INSTALLATION.md
Normal file
@ -0,0 +1,151 @@
|
||||
# QognyX Installation - Environnement de Développement
|
||||
|
||||
## Installation réalisée
|
||||
|
||||
✅ Node.js 20.20.0 installé
|
||||
✅ npm 10.8.2 installé
|
||||
✅ MongoDB 7.0.28 installé et démarré
|
||||
✅ QognyX sources téléchargées et configurées
|
||||
✅ Dépendances npm installées
|
||||
|
||||
## Configuration
|
||||
|
||||
### API Personnalisée
|
||||
L'environnement est configuré pour utiliser votre API personnalisée :
|
||||
- URL: https://cpu.lsbuchet.com/v1
|
||||
- Format: Compatible OpenAI
|
||||
|
||||
### Fichiers de configuration
|
||||
- `.env` : Variables d'environnement
|
||||
- `librechat.yaml` : Configuration de l'endpoint personnalisé
|
||||
|
||||
### Base de données
|
||||
MongoDB est configuré localement sur `mongodb://127.0.0.1:27017/QognyX`
|
||||
|
||||
## Utilisation
|
||||
|
||||
### Avant de démarrer
|
||||
1. Éditer `/opt/front-dev/.env` et remplacer `your_api_key_here` par votre vraie clé API:
|
||||
```bash
|
||||
nano /opt/front-dev/.env
|
||||
# Chercher CPU_API_KEY= et remplacer la valeur
|
||||
```
|
||||
|
||||
### Démarrer l'environnement de développement
|
||||
```bash
|
||||
cd /opt/front-dev
|
||||
./start-dev.sh
|
||||
```
|
||||
ou
|
||||
```bash
|
||||
cd /opt/front-dev
|
||||
npm run dev
|
||||
```
|
||||
|
||||
L'application sera accessible sur: http://localhost:3080
|
||||
|
||||
### Build pour la production
|
||||
```bash
|
||||
cd /opt/front-dev
|
||||
./build-prod.sh
|
||||
```
|
||||
|
||||
### Démarrer en production
|
||||
```bash
|
||||
cd /opt/front-dev
|
||||
npm run start
|
||||
```
|
||||
|
||||
## Commandes utiles
|
||||
|
||||
### Vérifier le statut de MongoDB
|
||||
```bash
|
||||
systemctl status mongod
|
||||
```
|
||||
|
||||
### Arrêter MongoDB
|
||||
```bash
|
||||
systemctl stop mongod
|
||||
```
|
||||
|
||||
### Démarrer MongoDB
|
||||
```bash
|
||||
systemctl start mongod
|
||||
```
|
||||
|
||||
### Mettre à jour les sources
|
||||
```bash
|
||||
cd /opt/front-dev
|
||||
curl -L https://github.com/danny-avila/QognyX/archive/refs/heads/main.tar.gz -o update.tar.gz
|
||||
tar -xzf update.tar.gz --strip-components=1
|
||||
rm update.tar.gz
|
||||
npm ci
|
||||
```
|
||||
|
||||
### Installer de nouvelles dépendances
|
||||
```bash
|
||||
cd /opt/front-dev
|
||||
npm install <package-name>
|
||||
```
|
||||
|
||||
## Développement
|
||||
|
||||
### Structure du projet
|
||||
- `api/` : Code backend (Node.js/Express)
|
||||
- `client/` : Code frontend (React)
|
||||
- `packages/` : Packages partagés
|
||||
- `config/` : Fichiers de configuration
|
||||
|
||||
### Variables d'environnement importantes
|
||||
- `CPU_API_KEY` : Clé API pour cpu.lsbuchet.com
|
||||
- `MONGO_URI` : URI de connexion MongoDB
|
||||
- `PORT` : Port du serveur (défaut: 3080)
|
||||
- `DEBUG_LOGGING` : Active les logs de debug
|
||||
|
||||
## Logs et Debugging
|
||||
|
||||
### Voir les logs MongoDB
|
||||
```bash
|
||||
tail -f /var/log/mongodb/mongod.log
|
||||
```
|
||||
|
||||
### Activer le mode debug
|
||||
Dans `.env`, définir:
|
||||
```
|
||||
DEBUG_LOGGING=true
|
||||
DEBUG_CONSOLE=true
|
||||
```
|
||||
|
||||
## Production vs Développement
|
||||
|
||||
### Environnement de développement
|
||||
- Hot-reload activé
|
||||
- Logs verbeux
|
||||
- Pas de minification
|
||||
|
||||
### Environnement de production
|
||||
- Build optimisé
|
||||
- Minification
|
||||
- Meilleure performance
|
||||
|
||||
## Servir plusieurs instances
|
||||
|
||||
Vous pouvez configurer différentes instances sur différents ports:
|
||||
- Développement: Port 3080
|
||||
- Production: Port 3081 (par exemple)
|
||||
|
||||
Modifier `PORT=` dans `.env` pour changer le port.
|
||||
|
||||
## Support
|
||||
|
||||
Pour plus d'informations sur QognyX:
|
||||
- Documentation: https://www.librechat.ai/docs
|
||||
- GitHub: https://github.com/danny-avila/QognyX
|
||||
|
||||
## Prochaines étapes
|
||||
|
||||
1. ✏️ Éditer `.env` pour ajouter votre clé API
|
||||
2. 🚀 Démarrer l'environnement de dev: `./start-dev.sh`
|
||||
3. 🌐 Ouvrir http://localhost:3080
|
||||
4. 📝 Créer un compte utilisateur
|
||||
5. 💬 Tester l'API personnalisée
|
||||
21
LICENSE
Normal file
21
LICENSE
Normal file
@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2025 LibreChat
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
217
README.md
Normal file
217
README.md
Normal file
@ -0,0 +1,217 @@
|
||||
<p align="center">
|
||||
<a href="https://librechat.ai">
|
||||
<img src="client/public/assets/logo.svg" height="256">
|
||||
</a>
|
||||
<h1 align="center">
|
||||
<a href="https://librechat.ai">QognyX</a>
|
||||
</h1>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="https://discord.librechat.ai">
|
||||
<img
|
||||
src="https://img.shields.io/discord/1086345563026489514?label=&logo=discord&style=for-the-badge&logoWidth=20&logoColor=white&labelColor=000000&color=blueviolet">
|
||||
</a>
|
||||
<a href="https://www.youtube.com/@QognyX">
|
||||
<img
|
||||
src="https://img.shields.io/badge/YOUTUBE-red.svg?style=for-the-badge&logo=youtube&logoColor=white&labelColor=000000&logoWidth=20">
|
||||
</a>
|
||||
<a href="https://docs.librechat.ai">
|
||||
<img
|
||||
src="https://img.shields.io/badge/DOCS-blue.svg?style=for-the-badge&logo=read-the-docs&logoColor=white&labelColor=000000&logoWidth=20">
|
||||
</a>
|
||||
<a aria-label="Sponsors" href="https://github.com/sponsors/danny-avila">
|
||||
<img
|
||||
src="https://img.shields.io/badge/SPONSORS-brightgreen.svg?style=for-the-badge&logo=github-sponsors&logoColor=white&labelColor=000000&logoWidth=20">
|
||||
</a>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="https://railway.app/template/b5k2mn?referralCode=HI9hWz">
|
||||
<img src="https://railway.app/button.svg" alt="Deploy on Railway" height="30">
|
||||
</a>
|
||||
<a href="https://zeabur.com/templates/0X2ZY8">
|
||||
<img src="https://zeabur.com/button.svg" alt="Deploy on Zeabur" height="30"/>
|
||||
</a>
|
||||
<a href="https://template.cloud.sealos.io/deploy?templateName=librechat">
|
||||
<img src="https://raw.githubusercontent.com/labring-actions/templates/main/Deploy-on-Sealos.svg" alt="Deploy on Sealos" height="30">
|
||||
</a>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="https://www.librechat.ai/docs/translation">
|
||||
<img
|
||||
src="https://img.shields.io/badge/dynamic/json.svg?style=for-the-badge&color=2096F3&label=locize&query=%24.translatedPercentage&url=https://api.locize.app/badgedata/4cb2598b-ed4d-469c-9b04-2ed531a8cb45&suffix=%+translated"
|
||||
alt="Translation Progress">
|
||||
</a>
|
||||
</p>
|
||||
|
||||
|
||||
# ✨ Features
|
||||
|
||||
- 🖥️ **UI & Experience** inspired by ChatGPT with enhanced design and features
|
||||
|
||||
- 🤖 **AI Model Selection**:
|
||||
- Anthropic (Claude), AWS Bedrock, OpenAI, Azure OpenAI, Google, Vertex AI, OpenAI Responses API (incl. Azure)
|
||||
- [Custom Endpoints](https://www.librechat.ai/docs/quick_start/custom_endpoints): Use any OpenAI-compatible API with QognyX, no proxy required
|
||||
- Compatible with [Local & Remote AI Providers](https://www.librechat.ai/docs/configuration/librechat_yaml/ai_endpoints):
|
||||
- Ollama, groq, Cohere, Mistral AI, Apple MLX, koboldcpp, together.ai,
|
||||
- OpenRouter, Helicone, Perplexity, ShuttleAI, Deepseek, Qwen, and more
|
||||
|
||||
- 🔧 **[Code Interpreter API](https://www.librechat.ai/docs/features/code_interpreter)**:
|
||||
- Secure, Sandboxed Execution in Python, Node.js (JS/TS), Go, C/C++, Java, PHP, Rust, and Fortran
|
||||
- Seamless File Handling: Upload, process, and download files directly
|
||||
- No Privacy Concerns: Fully isolated and secure execution
|
||||
|
||||
- 🔦 **Agents & Tools Integration**:
|
||||
- **[QognyX Agents](https://www.librechat.ai/docs/features/agents)**:
|
||||
- No-Code Custom Assistants: Build specialized, AI-driven helpers
|
||||
- Agent Marketplace: Discover and deploy community-built agents
|
||||
- Collaborative Sharing: Share agents with specific users and groups
|
||||
- Flexible & Extensible: Use MCP Servers, tools, file search, code execution, and more
|
||||
- Compatible with Custom Endpoints, OpenAI, Azure, Anthropic, AWS Bedrock, Google, Vertex AI, Responses API, and more
|
||||
- [Model Context Protocol (MCP) Support](https://modelcontextprotocol.io/clients#librechat) for Tools
|
||||
|
||||
- 🔍 **Web Search**:
|
||||
- Search the internet and retrieve relevant information to enhance your AI context
|
||||
- Combines search providers, content scrapers, and result rerankers for optimal results
|
||||
- **Customizable Jina Reranking**: Configure custom Jina API URLs for reranking services
|
||||
- **[Learn More →](https://www.librechat.ai/docs/features/web_search)**
|
||||
|
||||
- 🪄 **Generative UI with Code Artifacts**:
|
||||
- [Code Artifacts](https://youtu.be/GfTj7O4gmd0?si=WJbdnemZpJzBrJo3) allow creation of React, HTML, and Mermaid diagrams directly in chat
|
||||
|
||||
- 🎨 **Image Generation & Editing**
|
||||
- Text-to-image and image-to-image with [GPT-Image-1](https://www.librechat.ai/docs/features/image_gen#1--openai-image-tools-recommended)
|
||||
- Text-to-image with [DALL-E (3/2)](https://www.librechat.ai/docs/features/image_gen#2--dalle-legacy), [Stable Diffusion](https://www.librechat.ai/docs/features/image_gen#3--stable-diffusion-local), [Flux](https://www.librechat.ai/docs/features/image_gen#4--flux), or any [MCP server](https://www.librechat.ai/docs/features/image_gen#5--model-context-protocol-mcp)
|
||||
- Produce stunning visuals from prompts or refine existing images with a single instruction
|
||||
|
||||
- 💾 **Presets & Context Management**:
|
||||
- Create, Save, & Share Custom Presets
|
||||
- Switch between AI Endpoints and Presets mid-chat
|
||||
- Edit, Resubmit, and Continue Messages with Conversation branching
|
||||
- Create and share prompts with specific users and groups
|
||||
- [Fork Messages & Conversations](https://www.librechat.ai/docs/features/fork) for Advanced Context control
|
||||
|
||||
- 💬 **Multimodal & File Interactions**:
|
||||
- Upload and analyze images with Claude 3, GPT-4.5, GPT-4o, o1, Llama-Vision, and Gemini 📸
|
||||
- Chat with Files using Custom Endpoints, OpenAI, Azure, Anthropic, AWS Bedrock, & Google 🗃️
|
||||
|
||||
- 🌎 **Multilingual UI**:
|
||||
- English, 中文 (简体), 中文 (繁體), العربية, Deutsch, Español, Français, Italiano
|
||||
- Polski, Português (PT), Português (BR), Русский, 日本語, Svenska, 한국어, Tiếng Việt
|
||||
- Türkçe, Nederlands, עברית, Català, Čeština, Dansk, Eesti, فارسی
|
||||
- Suomi, Magyar, Հայերեն, Bahasa Indonesia, ქართული, Latviešu, ไทย, ئۇيغۇرچە
|
||||
|
||||
- 🧠 **Reasoning UI**:
|
||||
- Dynamic Reasoning UI for Chain-of-Thought/Reasoning AI models like DeepSeek-R1
|
||||
|
||||
- 🎨 **Customizable Interface**:
|
||||
- Customizable Dropdown & Interface that adapts to both power users and newcomers
|
||||
|
||||
- 🗣️ **Speech & Audio**:
|
||||
- Chat hands-free with Speech-to-Text and Text-to-Speech
|
||||
- Automatically send and play Audio
|
||||
- Supports OpenAI, Azure OpenAI, and Elevenlabs
|
||||
|
||||
- 📥 **Import & Export Conversations**:
|
||||
- Import Conversations from QognyX, ChatGPT, Chatbot UI
|
||||
- Export conversations as screenshots, markdown, text, json
|
||||
|
||||
- 🔍 **Search & Discovery**:
|
||||
- Search all messages/conversations
|
||||
|
||||
- 👥 **Multi-User & Secure Access**:
|
||||
- Multi-User, Secure Authentication with OAuth2, LDAP, & Email Login Support
|
||||
- Built-in Moderation, and Token spend tools
|
||||
|
||||
- ⚙️ **Configuration & Deployment**:
|
||||
- Configure Proxy, Reverse Proxy, Docker, & many Deployment options
|
||||
- Use completely local or deploy on the cloud
|
||||
|
||||
- 📖 **Open-Source & Community**:
|
||||
- Completely Open-Source & Built in Public
|
||||
- Community-driven development, support, and feedback
|
||||
|
||||
[For a thorough review of our features, see our docs here](https://docs.librechat.ai/) 📚
|
||||
|
||||
## 🪶 All-In-One AI Conversations with QognyX
|
||||
|
||||
QognyX brings together the future of assistant AIs with the revolutionary technology of OpenAI's ChatGPT. Celebrating the original styling, QognyX gives you the ability to integrate multiple AI models. It also integrates and enhances original client features such as conversation and message search, prompt templates and plugins.
|
||||
|
||||
With QognyX, you no longer need to opt for ChatGPT Plus and can instead use free or pay-per-call APIs. We welcome contributions, cloning, and forking to enhance the capabilities of this advanced chatbot platform.
|
||||
|
||||
[](https://www.youtube.com/watch?v=ilfwGQtJNlI)
|
||||
|
||||
Click on the thumbnail to open the video☝️
|
||||
|
||||
---
|
||||
|
||||
## 🌐 Resources
|
||||
|
||||
**GitHub Repo:**
|
||||
- **RAG API:** [github.com/danny-avila/rag_api](https://github.com/danny-avila/rag_api)
|
||||
- **Website:** [github.com/QognyX-AI/librechat.ai](https://github.com/QognyX-AI/librechat.ai)
|
||||
|
||||
**Other:**
|
||||
- **Website:** [librechat.ai](https://librechat.ai)
|
||||
- **Documentation:** [librechat.ai/docs](https://librechat.ai/docs)
|
||||
- **Blog:** [librechat.ai/blog](https://librechat.ai/blog)
|
||||
|
||||
---
|
||||
|
||||
## 📝 Changelog
|
||||
|
||||
Keep up with the latest updates by visiting the releases page and notes:
|
||||
- [Releases](https://github.com/danny-avila/QognyX/releases)
|
||||
- [Changelog](https://www.librechat.ai/changelog)
|
||||
|
||||
**⚠️ Please consult the [changelog](https://www.librechat.ai/changelog) for breaking changes before updating.**
|
||||
|
||||
---
|
||||
|
||||
## ⭐ Star History
|
||||
|
||||
<p align="center">
|
||||
<a href="https://star-history.com/#danny-avila/QognyX&Date">
|
||||
<img alt="Star History Chart" src="https://api.star-history.com/svg?repos=danny-avila/QognyX&type=Date&theme=dark" onerror="this.src='https://api.star-history.com/svg?repos=danny-avila/QognyX&type=Date'" />
|
||||
</a>
|
||||
</p>
|
||||
<p align="center">
|
||||
<a href="https://trendshift.io/repositories/4685" target="_blank" style="padding: 10px;">
|
||||
<img src="https://trendshift.io/api/badge/repositories/4685" alt="danny-avila%2FQognyX | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
|
||||
</a>
|
||||
<a href="https://runacap.com/ross-index/q1-24/" target="_blank" rel="noopener" style="margin-left: 20px;">
|
||||
<img style="width: 260px; height: 56px" src="https://runacap.com/wp-content/uploads/2024/04/ROSS_badge_white_Q1_2024.svg" alt="ROSS Index - Fastest Growing Open-Source Startups in Q1 2024 | Runa Capital" width="260" height="56"/>
|
||||
</a>
|
||||
</p>
|
||||
|
||||
---
|
||||
|
||||
## ✨ Contributions
|
||||
|
||||
Contributions, suggestions, bug reports and fixes are welcome!
|
||||
|
||||
For new features, components, or extensions, please open an issue and discuss before sending a PR.
|
||||
|
||||
If you'd like to help translate QognyX into your language, we'd love your contribution! Improving our translations not only makes QognyX more accessible to users around the world but also enhances the overall user experience. Please check out our [Translation Guide](https://www.librechat.ai/docs/translation).
|
||||
|
||||
---
|
||||
|
||||
## 💖 This project exists in its current state thanks to all the people who contribute
|
||||
|
||||
<a href="https://github.com/danny-avila/QognyX/graphs/contributors">
|
||||
<img src="https://contrib.rocks/image?repo=danny-avila/QognyX" />
|
||||
</a>
|
||||
|
||||
---
|
||||
|
||||
## 🎉 Special Thanks
|
||||
|
||||
We thank [Locize](https://locize.com) for their translation management tools that support multiple languages in QognyX.
|
||||
|
||||
<p align="center">
|
||||
<a href="https://locize.com" target="_blank" rel="noopener noreferrer">
|
||||
<img src="https://github.com/user-attachments/assets/d6b70894-6064-475e-bb65-92a9e23e0077" alt="Locize Logo" height="50">
|
||||
</a>
|
||||
</p>
|
||||
98
RESUME.txt
Normal file
98
RESUME.txt
Normal file
@ -0,0 +1,98 @@
|
||||
================================================================================
|
||||
INSTALLATION LIBRECHAT - RÉSUMÉ
|
||||
================================================================================
|
||||
|
||||
📁 RÉPERTOIRES CRÉÉS:
|
||||
---------------------
|
||||
✅ /opt/front-Prod - Pour l'environnement de production (vide pour le moment)
|
||||
✅ /opt/front-dev - Environnement de développement CONFIGURÉ
|
||||
✅ /opt/www - Répertoire www (vide pour le moment)
|
||||
|
||||
Chaque répertoire contient un sous-dossier /copilot
|
||||
|
||||
🔧 LOGICIELS INSTALLÉS:
|
||||
-----------------------
|
||||
✅ Node.js v20.20.0
|
||||
✅ npm v10.8.2
|
||||
✅ MongoDB v7.0.28 (démarré et activé au boot)
|
||||
|
||||
📦 LIBRECHAT - ENVIRONNEMENT DE DÉVELOPPEMENT (/opt/front-dev):
|
||||
---------------------------------------------------------------
|
||||
✅ Sources complètes téléchargées depuis GitHub
|
||||
✅ Toutes les dépendances npm installées (3169 packages)
|
||||
✅ Fichier .env configuré avec secrets JWT générés
|
||||
✅ Fichier librechat.yaml configuré pour l'API cpu.lsbuchet.com
|
||||
✅ Scripts de démarrage créés
|
||||
|
||||
📝 CONFIGURATION:
|
||||
-----------------
|
||||
API personnalisée: https://cpu.lsbuchet.com/v1 (format OpenAI)
|
||||
Base de données: MongoDB local sur port 27017
|
||||
Port serveur: 3080
|
||||
Mode: Développement avec hot-reload
|
||||
|
||||
⚠️ ACTION REQUISE AVANT LE DÉMARRAGE:
|
||||
--------------------------------------
|
||||
1. Éditer le fichier /opt/front-dev/.env
|
||||
2. Remplacer "your_api_key_here" par votre vraie clé API pour cpu.lsbuchet.com:
|
||||
|
||||
Commande: nano /opt/front-dev/.env
|
||||
Chercher: CPU_API_KEY=your_api_key_here
|
||||
Remplacer par: CPU_API_KEY=votre_vraie_clé_ici
|
||||
|
||||
🚀 DÉMARRAGE:
|
||||
-------------
|
||||
Pour démarrer l'environnement de développement:
|
||||
cd /opt/front-dev
|
||||
./start-dev.sh
|
||||
|
||||
Ou directement:
|
||||
cd /opt/front-dev
|
||||
npm run dev
|
||||
|
||||
L'application sera accessible sur: http://localhost:3080
|
||||
|
||||
📚 DOCUMENTATION:
|
||||
-----------------
|
||||
Documentation complète: /opt/front-dev/INSTALLATION.md
|
||||
|
||||
🔄 COMMANDES UTILES:
|
||||
--------------------
|
||||
Démarrer MongoDB: systemctl start mongod
|
||||
Arrêter MongoDB: systemctl stop mongod
|
||||
Statut MongoDB: systemctl status mongod
|
||||
Build production: cd /opt/front-dev && ./build-prod.sh
|
||||
Démarrer en prod: cd /opt/front-dev && npm run start
|
||||
|
||||
📊 FONCTIONNALITÉS:
|
||||
-------------------
|
||||
✅ Frontend React avec hot-reload
|
||||
✅ Backend Node.js/Express
|
||||
✅ API personnalisée configurée (cpu.lsbuchet.com)
|
||||
✅ MongoDB pour le stockage
|
||||
✅ Authentification locale activée
|
||||
✅ Inscription activée
|
||||
✅ Logs de debug activés
|
||||
✅ Mise à jour des sources facile
|
||||
|
||||
🔐 SÉCURITÉ:
|
||||
------------
|
||||
✅ Secrets JWT générés automatiquement
|
||||
✅ Clés de chiffrement créées
|
||||
✅ Base de données locale sécurisée
|
||||
|
||||
📌 PROCHAINES ÉTAPES:
|
||||
---------------------
|
||||
1. Ajouter votre clé API dans /opt/front-dev/.env
|
||||
2. Démarrer le serveur de développement: ./start-dev.sh
|
||||
3. Ouvrir http://localhost:3080 dans votre navigateur
|
||||
4. Créer un compte utilisateur
|
||||
5. Commencer à développer et tester !
|
||||
|
||||
Pour la production:
|
||||
- Configurer un environnement similaire dans /opt/front-Prod
|
||||
- Builder avec: ./build-prod.sh
|
||||
- Démarrer avec: npm run start
|
||||
- Configurer un reverse proxy (nginx/apache)
|
||||
|
||||
================================================================================
|
||||
1435
api/app/clients/BaseClient.js
Normal file
1435
api/app/clients/BaseClient.js
Normal file
File diff suppressed because it is too large
Load Diff
167
api/app/clients/OllamaClient.js
Normal file
167
api/app/clients/OllamaClient.js
Normal file
@ -0,0 +1,167 @@
|
||||
const { z } = require('zod');
|
||||
const axios = require('axios');
|
||||
const { Ollama } = require('ollama');
|
||||
const { sleep } = require('@librechat/agents');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { Constants } = require('librechat-data-provider');
|
||||
const { resolveHeaders, deriveBaseURL } = require('@librechat/api');
|
||||
|
||||
const ollamaPayloadSchema = z.object({
|
||||
mirostat: z.number().optional(),
|
||||
mirostat_eta: z.number().optional(),
|
||||
mirostat_tau: z.number().optional(),
|
||||
num_ctx: z.number().optional(),
|
||||
repeat_last_n: z.number().optional(),
|
||||
repeat_penalty: z.number().optional(),
|
||||
temperature: z.number().optional(),
|
||||
seed: z.number().nullable().optional(),
|
||||
stop: z.array(z.string()).optional(),
|
||||
tfs_z: z.number().optional(),
|
||||
num_predict: z.number().optional(),
|
||||
top_k: z.number().optional(),
|
||||
top_p: z.number().optional(),
|
||||
stream: z.optional(z.boolean()),
|
||||
model: z.string(),
|
||||
});
|
||||
|
||||
/**
|
||||
* @param {string} imageUrl
|
||||
* @returns {string}
|
||||
* @throws {Error}
|
||||
*/
|
||||
const getValidBase64 = (imageUrl) => {
|
||||
const parts = imageUrl.split(';base64,');
|
||||
|
||||
if (parts.length === 2) {
|
||||
return parts[1];
|
||||
} else {
|
||||
logger.error('Invalid or no Base64 string found in URL.');
|
||||
}
|
||||
};
|
||||
|
||||
class OllamaClient {
|
||||
constructor(options = {}) {
|
||||
const host = deriveBaseURL(options.baseURL ?? 'http://localhost:11434');
|
||||
this.streamRate = options.streamRate ?? Constants.DEFAULT_STREAM_RATE;
|
||||
this.headers = options.headers ?? {};
|
||||
/** @type {Ollama} */
|
||||
this.client = new Ollama({ host });
|
||||
}
|
||||
|
||||
/**
|
||||
* Fetches Ollama models from the specified base API path.
|
||||
* @param {string} baseURL
|
||||
* @param {Object} [options] - Optional configuration
|
||||
* @param {Partial<IUser>} [options.user] - User object for header resolution
|
||||
* @param {Record<string, string>} [options.headers] - Headers to include in the request
|
||||
* @returns {Promise<string[]>} The Ollama models.
|
||||
* @throws {Error} Throws if the Ollama API request fails
|
||||
*/
|
||||
static async fetchModels(baseURL, options = {}) {
|
||||
if (!baseURL) {
|
||||
return [];
|
||||
}
|
||||
|
||||
const ollamaEndpoint = deriveBaseURL(baseURL);
|
||||
|
||||
const resolvedHeaders = resolveHeaders({
|
||||
headers: options.headers,
|
||||
user: options.user,
|
||||
});
|
||||
|
||||
/** @type {Promise<AxiosResponse<OllamaListResponse>>} */
|
||||
const response = await axios.get(`${ollamaEndpoint}/api/tags`, {
|
||||
headers: resolvedHeaders,
|
||||
timeout: 5000,
|
||||
});
|
||||
|
||||
const models = response.data.models.map((tag) => tag.name);
|
||||
return models;
|
||||
}
|
||||
|
||||
/**
|
||||
* @param {ChatCompletionMessage[]} messages
|
||||
* @returns {OllamaMessage[]}
|
||||
*/
|
||||
static formatOpenAIMessages(messages) {
|
||||
const ollamaMessages = [];
|
||||
|
||||
for (const message of messages) {
|
||||
if (typeof message.content === 'string') {
|
||||
ollamaMessages.push({
|
||||
role: message.role,
|
||||
content: message.content,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
let aggregatedText = '';
|
||||
let imageUrls = [];
|
||||
|
||||
for (const content of message.content) {
|
||||
if (content.type === 'text') {
|
||||
aggregatedText += content.text + ' ';
|
||||
} else if (content.type === 'image_url') {
|
||||
imageUrls.push(getValidBase64(content.image_url.url));
|
||||
}
|
||||
}
|
||||
|
||||
const ollamaMessage = {
|
||||
role: message.role,
|
||||
content: aggregatedText.trim(),
|
||||
};
|
||||
|
||||
if (imageUrls.length > 0) {
|
||||
ollamaMessage.images = imageUrls;
|
||||
}
|
||||
|
||||
ollamaMessages.push(ollamaMessage);
|
||||
}
|
||||
|
||||
return ollamaMessages;
|
||||
}
|
||||
|
||||
/***
|
||||
* @param {Object} params
|
||||
* @param {ChatCompletionPayload} params.payload
|
||||
* @param {onTokenProgress} params.onProgress
|
||||
* @param {AbortController} params.abortController
|
||||
*/
|
||||
async chatCompletion({ payload, onProgress, abortController = null }) {
|
||||
let intermediateReply = '';
|
||||
|
||||
const parameters = ollamaPayloadSchema.parse(payload);
|
||||
const messages = OllamaClient.formatOpenAIMessages(payload.messages);
|
||||
|
||||
if (parameters.stream) {
|
||||
const stream = await this.client.chat({
|
||||
messages,
|
||||
...parameters,
|
||||
});
|
||||
|
||||
for await (const chunk of stream) {
|
||||
const token = chunk.message.content;
|
||||
intermediateReply += token;
|
||||
onProgress(token);
|
||||
if (abortController.signal.aborted) {
|
||||
stream.controller.abort();
|
||||
break;
|
||||
}
|
||||
|
||||
await sleep(this.streamRate);
|
||||
}
|
||||
}
|
||||
// TODO: regular completion
|
||||
else {
|
||||
// const generation = await this.client.generate(payload);
|
||||
}
|
||||
|
||||
return intermediateReply;
|
||||
}
|
||||
catch(err) {
|
||||
logger.error('[OllamaClient.chatCompletion]', err);
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = { OllamaClient, ollamaPayloadSchema };
|
||||
60
api/app/clients/TextStream.js
Normal file
60
api/app/clients/TextStream.js
Normal file
@ -0,0 +1,60 @@
|
||||
const { Readable } = require('stream');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
|
||||
class TextStream extends Readable {
|
||||
constructor(text, options = {}) {
|
||||
super(options);
|
||||
this.text = text;
|
||||
this.currentIndex = 0;
|
||||
this.minChunkSize = options.minChunkSize ?? 2;
|
||||
this.maxChunkSize = options.maxChunkSize ?? 4;
|
||||
this.delay = options.delay ?? 20; // Time in milliseconds
|
||||
}
|
||||
|
||||
_read() {
|
||||
const { delay, minChunkSize, maxChunkSize } = this;
|
||||
|
||||
if (this.currentIndex < this.text.length) {
|
||||
setTimeout(() => {
|
||||
const remainingChars = this.text.length - this.currentIndex;
|
||||
const chunkSize = Math.min(this.randomInt(minChunkSize, maxChunkSize + 1), remainingChars);
|
||||
|
||||
const chunk = this.text.slice(this.currentIndex, this.currentIndex + chunkSize);
|
||||
this.push(chunk);
|
||||
this.currentIndex += chunkSize;
|
||||
}, delay);
|
||||
} else {
|
||||
this.push(null); // signal end of data
|
||||
}
|
||||
}
|
||||
|
||||
randomInt(min, max) {
|
||||
return Math.floor(Math.random() * (max - min)) + min;
|
||||
}
|
||||
|
||||
async processTextStream(onProgressCallback) {
|
||||
const streamPromise = new Promise((resolve, reject) => {
|
||||
this.on('data', (chunk) => {
|
||||
onProgressCallback(chunk.toString());
|
||||
});
|
||||
|
||||
this.on('end', () => {
|
||||
// logger.debug('[processTextStream] Stream ended');
|
||||
resolve();
|
||||
});
|
||||
|
||||
this.on('error', (err) => {
|
||||
reject(err);
|
||||
});
|
||||
});
|
||||
|
||||
try {
|
||||
await streamPromise;
|
||||
} catch (err) {
|
||||
logger.error('[processTextStream] Error in text stream:', err);
|
||||
// Handle the error appropriately, e.g., return an error message or throw an error
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = TextStream;
|
||||
7
api/app/clients/index.js
Normal file
7
api/app/clients/index.js
Normal file
@ -0,0 +1,7 @@
|
||||
const TextStream = require('./TextStream');
|
||||
const toolUtils = require('./tools/util');
|
||||
|
||||
module.exports = {
|
||||
TextStream,
|
||||
...toolUtils,
|
||||
};
|
||||
537
api/app/clients/prompts/artifacts.js
Normal file
537
api/app/clients/prompts/artifacts.js
Normal file
@ -0,0 +1,537 @@
|
||||
const dedent = require('dedent');
|
||||
const { EModelEndpoint, ArtifactModes } = require('librechat-data-provider');
|
||||
const { generateShadcnPrompt } = require('~/app/clients/prompts/shadcn-docs/generate');
|
||||
const { components } = require('~/app/clients/prompts/shadcn-docs/components');
|
||||
|
||||
/** @deprecated */
|
||||
// eslint-disable-next-line no-unused-vars
|
||||
const artifactsPromptV1 = dedent`The assistant can create and reference artifacts during conversations.
|
||||
|
||||
Artifacts are for substantial, self-contained content that users might modify or reuse, displayed in a separate UI window for clarity.
|
||||
|
||||
# Good artifacts are...
|
||||
- Substantial content (>15 lines)
|
||||
- Content that the user is likely to modify, iterate on, or take ownership of
|
||||
- Self-contained, complex content that can be understood on its own, without context from the conversation
|
||||
- Content intended for eventual use outside the conversation (e.g., reports, emails, presentations)
|
||||
- Content likely to be referenced or reused multiple times
|
||||
|
||||
# Don't use artifacts for...
|
||||
- Simple, informational, or short content, such as brief code snippets, mathematical equations, or small examples
|
||||
- Primarily explanatory, instructional, or illustrative content, such as examples provided to clarify a concept
|
||||
- Suggestions, commentary, or feedback on existing artifacts
|
||||
- Conversational or explanatory content that doesn't represent a standalone piece of work
|
||||
- Content that is dependent on the current conversational context to be useful
|
||||
- Content that is unlikely to be modified or iterated upon by the user
|
||||
- Request from users that appears to be a one-off question
|
||||
|
||||
# Usage notes
|
||||
- One artifact per message unless specifically requested
|
||||
- Prefer in-line content (don't use artifacts) when possible. Unnecessary use of artifacts can be jarring for users.
|
||||
- If a user asks the assistant to "draw an SVG" or "make a website," the assistant does not need to explain that it doesn't have these capabilities. Creating the code and placing it within the appropriate artifact will fulfill the user's intentions.
|
||||
- If asked to generate an image, the assistant can offer an SVG instead. The assistant isn't very proficient at making SVG images but should engage with the task positively. Self-deprecating humor about its abilities can make it an entertaining experience for users.
|
||||
- The assistant errs on the side of simplicity and avoids overusing artifacts for content that can be effectively presented within the conversation.
|
||||
- Always provide complete, specific, and fully functional content without any placeholders, ellipses, or 'remains the same' comments.
|
||||
|
||||
<artifact_instructions>
|
||||
When collaborating with the user on creating content that falls into compatible categories, the assistant should follow these steps:
|
||||
|
||||
1. Create the artifact using the following format:
|
||||
|
||||
:::artifact{identifier="unique-identifier" type="mime-type" title="Artifact Title"}
|
||||
\`\`\`
|
||||
Your artifact content here
|
||||
\`\`\`
|
||||
:::
|
||||
|
||||
2. Assign an identifier to the \`identifier\` attribute. For updates, reuse the prior identifier. For new artifacts, the identifier should be descriptive and relevant to the content, using kebab-case (e.g., "example-code-snippet"). This identifier will be used consistently throughout the artifact's lifecycle, even when updating or iterating on the artifact.
|
||||
3. Include a \`title\` attribute to provide a brief title or description of the content.
|
||||
4. Add a \`type\` attribute to specify the type of content the artifact represents. Assign one of the following values to the \`type\` attribute:
|
||||
- HTML: "text/html"
|
||||
- The user interface can render single file HTML pages placed within the artifact tags. HTML, JS, and CSS should be in a single file when using the \`text/html\` type.
|
||||
- Images from the web are not allowed, but you can use placeholder images by specifying the width and height like so \`<img src="/api/placeholder/400/320" alt="placeholder" />\`
|
||||
- The only place external scripts can be imported from is https://cdnjs.cloudflare.com
|
||||
- Mermaid Diagrams: "application/vnd.mermaid"
|
||||
- The user interface will render Mermaid diagrams placed within the artifact tags.
|
||||
- React Components: "application/vnd.react"
|
||||
- Use this for displaying either: React elements, e.g. \`<strong>Hello World!</strong>\`, React pure functional components, e.g. \`() => <strong>Hello World!</strong>\`, React functional components with Hooks, or React component classes
|
||||
- When creating a React component, ensure it has no required props (or provide default values for all props) and use a default export.
|
||||
- Use Tailwind classes for styling. DO NOT USE ARBITRARY VALUES (e.g. \`h-[600px]\`).
|
||||
- Base React is available to be imported. To use hooks, first import it at the top of the artifact, e.g. \`import { useState } from "react"\`
|
||||
- The lucide-react@0.263.1 library is available to be imported. e.g. \`import { Camera } from "lucide-react"\` & \`<Camera color="red" size={48} />\`
|
||||
- The recharts charting library is available to be imported, e.g. \`import { LineChart, XAxis, ... } from "recharts"\` & \`<LineChart ...><XAxis dataKey="name"> ...\`
|
||||
- The assistant can use prebuilt components from the \`shadcn/ui\` library after it is imported: \`import { Alert, AlertDescription, AlertTitle, AlertDialog, AlertDialogAction } from '/components/ui/alert';\`. If using components from the shadcn/ui library, the assistant mentions this to the user and offers to help them install the components if necessary.
|
||||
- Components MUST be imported from \`/components/ui/name\` and NOT from \`/components/name\` or \`@/components/ui/name\`.
|
||||
- NO OTHER LIBRARIES (e.g. zod, hookform) ARE INSTALLED OR ABLE TO BE IMPORTED.
|
||||
- Images from the web are not allowed, but you can use placeholder images by specifying the width and height like so \`<img src="/api/placeholder/400/320" alt="placeholder" />\`
|
||||
- If you are unable to follow the above requirements for any reason, don't use artifacts and use regular code blocks instead, which will not attempt to render the component.
|
||||
5. Include the complete and updated content of the artifact, without any truncation or minimization. Don't use "// rest of the code remains the same...".
|
||||
6. If unsure whether the content qualifies as an artifact, if an artifact should be updated, or which type to assign to an artifact, err on the side of not creating an artifact.
|
||||
7. Always use triple backticks (\`\`\`) to enclose the content within the artifact, regardless of the content type.
|
||||
</artifact_instructions>
|
||||
|
||||
Here are some examples of correct usage of artifacts:
|
||||
|
||||
<examples>
|
||||
<example_docstring>
|
||||
This example demonstrates how to create a Mermaid artifact for a simple flow chart.
|
||||
</example_docstring>
|
||||
|
||||
<example>
|
||||
<user_query>Can you create a simple flow chart showing the process of making tea using Mermaid?</user_query>
|
||||
|
||||
<assistant_response>
|
||||
Sure! Here's a simple flow chart depicting the process of making tea using Mermaid syntax:
|
||||
|
||||
:::artifact{identifier="tea-making-flowchart" type="application/vnd.mermaid" title="Flow chart: Making Tea"}
|
||||
\`\`\`mermaid
|
||||
graph TD
|
||||
A[Start] --> B{Water boiled?}
|
||||
B -->|Yes| C[Add tea leaves to cup]
|
||||
B -->|No| D[Boil water]
|
||||
D --> B
|
||||
C --> E[Pour boiling water into cup]
|
||||
E --> F[Steep tea for desired time]
|
||||
F --> G[Remove tea leaves]
|
||||
G --> H[Add milk or sugar, if desired]
|
||||
H --> I[Enjoy your tea!]
|
||||
I --> J[End]
|
||||
\`\`\`
|
||||
:::
|
||||
|
||||
This flow chart uses Mermaid syntax to visualize the steps involved in making a cup of tea. Here's a brief explanation of the process:
|
||||
|
||||
1. Start
|
||||
2. Check if water is boiled
|
||||
3. If not boiled, boil the water
|
||||
4. Once water is boiled, add tea leaves to the cup
|
||||
5. Pour boiling water into the cup
|
||||
6. Steep the tea for the desired time
|
||||
7. Remove the tea leaves
|
||||
8. Optionally add milk or sugar
|
||||
9. Enjoy your tea!
|
||||
10. End
|
||||
|
||||
This chart provides a clear visual representation of the tea-making process. You can easily modify or expand this chart if you want to add more details or steps to the process. Let me know if you'd like any changes or have any questions!
|
||||
</assistant_response>
|
||||
</example>
|
||||
</examples>`;
|
||||
|
||||
const artifactsPrompt = dedent`The assistant can create and reference artifacts during conversations.
|
||||
|
||||
Artifacts are for substantial, self-contained content that users might modify or reuse, displayed in a separate UI window for clarity.
|
||||
|
||||
# Good artifacts are...
|
||||
- Substantial content (>15 lines)
|
||||
- Content that the user is likely to modify, iterate on, or take ownership of
|
||||
- Self-contained, complex content that can be understood on its own, without context from the conversation
|
||||
- Content intended for eventual use outside the conversation (e.g., reports, emails, presentations)
|
||||
- Content likely to be referenced or reused multiple times
|
||||
|
||||
# Don't use artifacts for...
|
||||
- Simple, informational, or short content, such as brief code snippets, mathematical equations, or small examples
|
||||
- Primarily explanatory, instructional, or illustrative content, such as examples provided to clarify a concept
|
||||
- Suggestions, commentary, or feedback on existing artifacts
|
||||
- Conversational or explanatory content that doesn't represent a standalone piece of work
|
||||
- Content that is dependent on the current conversational context to be useful
|
||||
- Content that is unlikely to be modified or iterated upon by the user
|
||||
- Request from users that appears to be a one-off question
|
||||
|
||||
# Usage notes
|
||||
- One artifact per message unless specifically requested
|
||||
- Prefer in-line content (don't use artifacts) when possible. Unnecessary use of artifacts can be jarring for users.
|
||||
- If a user asks the assistant to "draw an SVG" or "make a website," the assistant does not need to explain that it doesn't have these capabilities. Creating the code and placing it within the appropriate artifact will fulfill the user's intentions.
|
||||
- If asked to generate an image, the assistant can offer an SVG instead. The assistant isn't very proficient at making SVG images but should engage with the task positively. Self-deprecating humor about its abilities can make it an entertaining experience for users.
|
||||
- The assistant errs on the side of simplicity and avoids overusing artifacts for content that can be effectively presented within the conversation.
|
||||
- Always provide complete, specific, and fully functional content for artifacts without any snippets, placeholders, ellipses, or 'remains the same' comments.
|
||||
- If an artifact is not necessary or requested, the assistant should not mention artifacts at all, and respond to the user accordingly.
|
||||
|
||||
<artifact_instructions>
|
||||
When collaborating with the user on creating content that falls into compatible categories, the assistant should follow these steps:
|
||||
|
||||
1. Create the artifact using the following format:
|
||||
|
||||
:::artifact{identifier="unique-identifier" type="mime-type" title="Artifact Title"}
|
||||
\`\`\`
|
||||
Your artifact content here
|
||||
\`\`\`
|
||||
:::
|
||||
|
||||
2. Assign an identifier to the \`identifier\` attribute. For updates, reuse the prior identifier. For new artifacts, the identifier should be descriptive and relevant to the content, using kebab-case (e.g., "example-code-snippet"). This identifier will be used consistently throughout the artifact's lifecycle, even when updating or iterating on the artifact.
|
||||
3. Include a \`title\` attribute to provide a brief title or description of the content.
|
||||
4. Add a \`type\` attribute to specify the type of content the artifact represents. Assign one of the following values to the \`type\` attribute:
|
||||
- HTML: "text/html"
|
||||
- The user interface can render single file HTML pages placed within the artifact tags. HTML, JS, and CSS should be in a single file when using the \`text/html\` type.
|
||||
- Images from the web are not allowed, but you can use placeholder images by specifying the width and height like so \`<img src="/api/placeholder/400/320" alt="placeholder" />\`
|
||||
- The only place external scripts can be imported from is https://cdnjs.cloudflare.com
|
||||
- SVG: "image/svg+xml"
|
||||
- The user interface will render the Scalable Vector Graphics (SVG) image within the artifact tags.
|
||||
- The assistant should specify the viewbox of the SVG rather than defining a width/height
|
||||
- Markdown: "text/markdown" or "text/md"
|
||||
- The user interface will render Markdown content placed within the artifact tags.
|
||||
- Supports standard Markdown syntax including headers, lists, links, images, code blocks, tables, and more.
|
||||
- Both "text/markdown" and "text/md" are accepted as valid MIME types for Markdown content.
|
||||
- Mermaid Diagrams: "application/vnd.mermaid"
|
||||
- The user interface will render Mermaid diagrams placed within the artifact tags.
|
||||
- React Components: "application/vnd.react"
|
||||
- Use this for displaying either: React elements, e.g. \`<strong>Hello World!</strong>\`, React pure functional components, e.g. \`() => <strong>Hello World!</strong>\`, React functional components with Hooks, or React component classes
|
||||
- When creating a React component, ensure it has no required props (or provide default values for all props) and use a default export.
|
||||
- Use Tailwind classes for styling. DO NOT USE ARBITRARY VALUES (e.g. \`h-[600px]\`).
|
||||
- Base React is available to be imported. To use hooks, first import it at the top of the artifact, e.g. \`import { useState } from "react"\`
|
||||
- The lucide-react@0.394.0 library is available to be imported. e.g. \`import { Camera } from "lucide-react"\` & \`<Camera color="red" size={48} />\`
|
||||
- The recharts charting library is available to be imported, e.g. \`import { LineChart, XAxis, ... } from "recharts"\` & \`<LineChart ...><XAxis dataKey="name"> ...\`
|
||||
- The three.js library is available to be imported, e.g. \`import * as THREE from "three";\`
|
||||
- The date-fns library is available to be imported, e.g. \`import { compareAsc, format } from "date-fns";\`
|
||||
- The react-day-picker library is available to be imported, e.g. \`import { DayPicker } from "react-day-picker";\`
|
||||
- The assistant can use prebuilt components from the \`shadcn/ui\` library after it is imported: \`import { Alert, AlertDescription, AlertTitle, AlertDialog, AlertDialogAction } from '/components/ui/alert';\`. If using components from the shadcn/ui library, the assistant mentions this to the user and offers to help them install the components if necessary.
|
||||
- Components MUST be imported from \`/components/ui/name\` and NOT from \`/components/name\` or \`@/components/ui/name\`.
|
||||
- NO OTHER LIBRARIES (e.g. zod, hookform) ARE INSTALLED OR ABLE TO BE IMPORTED.
|
||||
- Images from the web are not allowed, but you can use placeholder images by specifying the width and height like so \`<img src="/api/placeholder/400/320" alt="placeholder" />\`
|
||||
- When iterating on code, ensure that the code is complete and functional without any snippets, placeholders, or ellipses.
|
||||
- If you are unable to follow the above requirements for any reason, don't use artifacts and use regular code blocks instead, which will not attempt to render the component.
|
||||
5. Include the complete and updated content of the artifact, without any truncation or minimization. Don't use "// rest of the code remains the same...".
|
||||
6. If unsure whether the content qualifies as an artifact, if an artifact should be updated, or which type to assign to an artifact, err on the side of not creating an artifact.
|
||||
7. Always use triple backticks (\`\`\`) to enclose the content within the artifact, regardless of the content type.
|
||||
</artifact_instructions>
|
||||
|
||||
Here are some examples of correct usage of artifacts:
|
||||
|
||||
<examples>
|
||||
<example_docstring>
|
||||
This example demonstrates how to create a Mermaid artifact for a simple flow chart.
|
||||
</example_docstring>
|
||||
|
||||
<example>
|
||||
<user_query>Can you create a simple flow chart showing the process of making tea using Mermaid?</user_query>
|
||||
|
||||
<assistant_response>
|
||||
Sure! Here's a simple flow chart depicting the process of making tea using Mermaid syntax:
|
||||
|
||||
:::artifact{identifier="tea-making-flowchart" type="application/vnd.mermaid" title="Flow chart: Making Tea"}
|
||||
\`\`\`mermaid
|
||||
graph TD
|
||||
A[Start] --> B{Water boiled?}
|
||||
B -->|Yes| C[Add tea leaves to cup]
|
||||
B -->|No| D[Boil water]
|
||||
D --> B
|
||||
C --> E[Pour boiling water into cup]
|
||||
E --> F[Steep tea for desired time]
|
||||
F --> G[Remove tea leaves]
|
||||
G --> H[Add milk or sugar, if desired]
|
||||
H --> I[Enjoy your tea!]
|
||||
I --> J[End]
|
||||
\`\`\`
|
||||
:::
|
||||
|
||||
This flow chart uses Mermaid syntax to visualize the steps involved in making a cup of tea. Here's a brief explanation of the process:
|
||||
|
||||
1. Start
|
||||
2. Check if water is boiled
|
||||
3. If not boiled, boil the water
|
||||
4. Once water is boiled, add tea leaves to the cup
|
||||
5. Pour boiling water into the cup
|
||||
6. Steep the tea for the desired time
|
||||
7. Remove the tea leaves
|
||||
8. Optionally add milk or sugar
|
||||
9. Enjoy your tea!
|
||||
10. End
|
||||
|
||||
This chart provides a clear visual representation of the tea-making process. You can easily modify or expand this chart if you want to add more details or steps to the process. Let me know if you'd like any changes or have any questions!
|
||||
</assistant_response>
|
||||
</example>
|
||||
|
||||
<example>
|
||||
<user_query>Create a simple React counter component</user_query>
|
||||
<assistant_response>
|
||||
Here's a simple React counter component:
|
||||
|
||||
:::artifact{identifier="react-counter" type="application/vnd.react" title="React Counter"}
|
||||
\`\`\`
|
||||
import { useState } from 'react';
|
||||
|
||||
export default function Counter() {
|
||||
const [count, setCount] = useState(0);
|
||||
return (
|
||||
<div className="p-4">
|
||||
<p className="mb-2">Count: {count}</p>
|
||||
<button className="bg-blue-500 text-white px-4 py-2 rounded" onClick={() => setCount(count + 1)}>
|
||||
Increment
|
||||
</button>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
\`\`\`
|
||||
:::
|
||||
|
||||
This component creates a simple counter with an increment button.
|
||||
</assistant_response>
|
||||
</example>
|
||||
|
||||
<example>
|
||||
<user_query>Create a basic HTML structure for a blog post</user_query>
|
||||
<assistant_response>
|
||||
Here's a basic HTML structure for a blog post:
|
||||
|
||||
:::artifact{identifier="blog-post-html" type="text/html" title="Blog Post HTML"}
|
||||
\`\`\`
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>My Blog Post</title>
|
||||
<style>
|
||||
body { font-family: Arial, sans-serif; line-height: 1.6; max-width: 800px; margin: 0 auto; padding: 20px; }
|
||||
h1 { color: #333; }
|
||||
p { margin-bottom: 15px; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<header>
|
||||
<h1>My First Blog Post</h1>
|
||||
</header>
|
||||
<main>
|
||||
<article>
|
||||
<p>This is the content of my blog post. It's short and sweet!</p>
|
||||
</article>
|
||||
</main>
|
||||
<footer>
|
||||
<p>© 2023 My Blog</p>
|
||||
</footer>
|
||||
</body>
|
||||
</html>
|
||||
\`\`\`
|
||||
:::
|
||||
|
||||
This HTML structure provides a simple layout for a blog post.
|
||||
</assistant_response>
|
||||
</example>
|
||||
</examples>`;
|
||||
|
||||
const artifactsOpenAIPrompt = dedent`The assistant can create and reference artifacts during conversations.
|
||||
|
||||
Artifacts are for substantial, self-contained content that users might modify or reuse, displayed in a separate UI window for clarity.
|
||||
|
||||
# Good artifacts are...
|
||||
- Substantial content (>15 lines)
|
||||
- Content that the user is likely to modify, iterate on, or take ownership of
|
||||
- Self-contained, complex content that can be understood on its own, without context from the conversation
|
||||
- Content intended for eventual use outside the conversation (e.g., reports, emails, presentations)
|
||||
- Content likely to be referenced or reused multiple times
|
||||
|
||||
# Don't use artifacts for...
|
||||
- Simple, informational, or short content, such as brief code snippets, mathematical equations, or small examples
|
||||
- Primarily explanatory, instructional, or illustrative content, such as examples provided to clarify a concept
|
||||
- Suggestions, commentary, or feedback on existing artifacts
|
||||
- Conversational or explanatory content that doesn't represent a standalone piece of work
|
||||
- Content that is dependent on the current conversational context to be useful
|
||||
- Content that is unlikely to be modified or iterated upon by the user
|
||||
- Request from users that appears to be a one-off question
|
||||
|
||||
# Usage notes
|
||||
- One artifact per message unless specifically requested
|
||||
- Prefer in-line content (don't use artifacts) when possible. Unnecessary use of artifacts can be jarring for users.
|
||||
- If a user asks the assistant to "draw an SVG" or "make a website," the assistant does not need to explain that it doesn't have these capabilities. Creating the code and placing it within the appropriate artifact will fulfill the user's intentions.
|
||||
- If asked to generate an image, the assistant can offer an SVG instead. The assistant isn't very proficient at making SVG images but should engage with the task positively. Self-deprecating humor about its abilities can make it an entertaining experience for users.
|
||||
- The assistant errs on the side of simplicity and avoids overusing artifacts for content that can be effectively presented within the conversation.
|
||||
- Always provide complete, specific, and fully functional content for artifacts without any snippets, placeholders, ellipses, or 'remains the same' comments.
|
||||
- If an artifact is not necessary or requested, the assistant should not mention artifacts at all, and respond to the user accordingly.
|
||||
|
||||
## Artifact Instructions
|
||||
When collaborating with the user on creating content that falls into compatible categories, the assistant should follow these steps:
|
||||
|
||||
1. Create the artifact using the following remark-directive markdown format:
|
||||
|
||||
:::artifact{identifier="unique-identifier" type="mime-type" title="Artifact Title"}
|
||||
\`\`\`
|
||||
Your artifact content here
|
||||
\`\`\`
|
||||
:::
|
||||
|
||||
a. Example of correct format:
|
||||
|
||||
:::artifact{identifier="example-artifact" type="text/plain" title="Example Artifact"}
|
||||
\`\`\`
|
||||
This is the content of the artifact.
|
||||
It can span multiple lines.
|
||||
\`\`\`
|
||||
:::
|
||||
|
||||
b. Common mistakes to avoid:
|
||||
- Don't split the opening ::: line
|
||||
- Don't add extra backticks outside the artifact structure
|
||||
- Don't omit the closing :::
|
||||
|
||||
2. Assign an identifier to the \`identifier\` attribute. For updates, reuse the prior identifier. For new artifacts, the identifier should be descriptive and relevant to the content, using kebab-case (e.g., "example-code-snippet"). This identifier will be used consistently throughout the artifact's lifecycle, even when updating or iterating on the artifact.
|
||||
3. Include a \`title\` attribute to provide a brief title or description of the content.
|
||||
4. Add a \`type\` attribute to specify the type of content the artifact represents. Assign one of the following values to the \`type\` attribute:
|
||||
- HTML: "text/html"
|
||||
- The user interface can render single file HTML pages placed within the artifact tags. HTML, JS, and CSS should be in a single file when using the \`text/html\` type.
|
||||
- Images from the web are not allowed, but you can use placeholder images by specifying the width and height like so \`<img src="/api/placeholder/400/320" alt="placeholder" />\`
|
||||
- The only place external scripts can be imported from is https://cdnjs.cloudflare.com
|
||||
- SVG: "image/svg+xml"
|
||||
- The user interface will render the Scalable Vector Graphics (SVG) image within the artifact tags.
|
||||
- The assistant should specify the viewbox of the SVG rather than defining a width/height
|
||||
- Markdown: "text/markdown" or "text/md"
|
||||
- The user interface will render Markdown content placed within the artifact tags.
|
||||
- Supports standard Markdown syntax including headers, lists, links, images, code blocks, tables, and more.
|
||||
- Both "text/markdown" and "text/md" are accepted as valid MIME types for Markdown content.
|
||||
- Mermaid Diagrams: "application/vnd.mermaid"
|
||||
- The user interface will render Mermaid diagrams placed within the artifact tags.
|
||||
- React Components: "application/vnd.react"
|
||||
- Use this for displaying either: React elements, e.g. \`<strong>Hello World!</strong>\`, React pure functional components, e.g. \`() => <strong>Hello World!</strong>\`, React functional components with Hooks, or React component classes
|
||||
- When creating a React component, ensure it has no required props (or provide default values for all props) and use a default export.
|
||||
- Use Tailwind classes for styling. DO NOT USE ARBITRARY VALUES (e.g. \`h-[600px]\`).
|
||||
- Base React is available to be imported. To use hooks, first import it at the top of the artifact, e.g. \`import { useState } from "react"\`
|
||||
- The lucide-react@0.394.0 library is available to be imported. e.g. \`import { Camera } from "lucide-react"\` & \`<Camera color="red" size={48} />\`
|
||||
- The recharts charting library is available to be imported, e.g. \`import { LineChart, XAxis, ... } from "recharts"\` & \`<LineChart ...><XAxis dataKey="name"> ...\`
|
||||
- The three.js library is available to be imported, e.g. \`import * as THREE from "three";\`
|
||||
- The date-fns library is available to be imported, e.g. \`import { compareAsc, format } from "date-fns";\`
|
||||
- The react-day-picker library is available to be imported, e.g. \`import { DayPicker } from "react-day-picker";\`
|
||||
- The assistant can use prebuilt components from the \`shadcn/ui\` library after it is imported: \`import { Alert, AlertDescription, AlertTitle, AlertDialog, AlertDialogAction } from '/components/ui/alert';\`. If using components from the shadcn/ui library, the assistant mentions this to the user and offers to help them install the components if necessary.
|
||||
- Components MUST be imported from \`/components/ui/name\` and NOT from \`/components/name\` or \`@/components/ui/name\`.
|
||||
- NO OTHER LIBRARIES (e.g. zod, hookform) ARE INSTALLED OR ABLE TO BE IMPORTED.
|
||||
- Images from the web are not allowed, but you can use placeholder images by specifying the width and height like so \`<img src="/api/placeholder/400/320" alt="placeholder" />\`
|
||||
- When iterating on code, ensure that the code is complete and functional without any snippets, placeholders, or ellipses.
|
||||
- If you are unable to follow the above requirements for any reason, don't use artifacts and use regular code blocks instead, which will not attempt to render the component.
|
||||
5. Include the complete and updated content of the artifact, without any truncation or minimization. Don't use "// rest of the code remains the same...".
|
||||
6. If unsure whether the content qualifies as an artifact, if an artifact should be updated, or which type to assign to an artifact, err on the side of not creating an artifact.
|
||||
7. NEVER use triple backticks to enclose the artifact, ONLY the content within the artifact.
|
||||
|
||||
Here are some examples of correct usage of artifacts:
|
||||
|
||||
## Examples
|
||||
|
||||
### Example 1
|
||||
|
||||
This example demonstrates how to create a Mermaid artifact for a simple flow chart.
|
||||
|
||||
User: Can you create a simple flow chart showing the process of making tea using Mermaid?
|
||||
|
||||
Assistant: Sure! Here's a simple flow chart depicting the process of making tea using Mermaid syntax:
|
||||
|
||||
:::artifact{identifier="tea-making-flowchart" type="application/vnd.mermaid" title="Flow chart: Making Tea"}
|
||||
\`\`\`mermaid
|
||||
graph TD
|
||||
A[Start] --> B{Water boiled?}
|
||||
B -->|Yes| C[Add tea leaves to cup]
|
||||
B -->|No| D[Boil water]
|
||||
D --> B
|
||||
C --> E[Pour boiling water into cup]
|
||||
E --> F[Steep tea for desired time]
|
||||
F --> G[Remove tea leaves]
|
||||
G --> H[Add milk or sugar, if desired]
|
||||
H --> I[Enjoy your tea!]
|
||||
I --> J[End]
|
||||
\`\`\`
|
||||
:::
|
||||
|
||||
This flow chart uses Mermaid syntax to visualize the steps involved in making a cup of tea. Here's a brief explanation of the process:
|
||||
|
||||
1. Start
|
||||
2. Check if water is boiled
|
||||
3. If not boiled, boil the water
|
||||
4. Once water is boiled, add tea leaves to the cup
|
||||
5. Pour boiling water into the cup
|
||||
6. Steep the tea for the desired time
|
||||
7. Remove the tea leaves
|
||||
8. Optionally add milk or sugar
|
||||
9. Enjoy your tea!
|
||||
10. End
|
||||
|
||||
This chart provides a clear visual representation of the tea-making process. You can easily modify or expand this chart if you want to add more details or steps to the process. Let me know if you'd like any changes or have any questions!
|
||||
|
||||
---
|
||||
|
||||
### Example 2
|
||||
|
||||
User: Create a simple React counter component
|
||||
|
||||
Assistant: Here's a simple React counter component:
|
||||
|
||||
:::artifact{identifier="react-counter" type="application/vnd.react" title="React Counter"}
|
||||
\`\`\`
|
||||
import { useState } from 'react';
|
||||
|
||||
export default function Counter() {
|
||||
const [count, setCount] = useState(0);
|
||||
return (
|
||||
<div className="p-4">
|
||||
<p className="mb-2">Count: {count}</p>
|
||||
<button className="bg-blue-500 text-white px-4 py-2 rounded" onClick={() => setCount(count + 1)}>
|
||||
Increment
|
||||
</button>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
\`\`\`
|
||||
:::
|
||||
|
||||
This component creates a simple counter with an increment button.
|
||||
|
||||
---
|
||||
|
||||
### Example 3
|
||||
User: Create a basic HTML structure for a blog post
|
||||
Assistant: Here's a basic HTML structure for a blog post:
|
||||
|
||||
:::artifact{identifier="blog-post-html" type="text/html" title="Blog Post HTML"}
|
||||
\`\`\`
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>My Blog Post</title>
|
||||
<style>
|
||||
body { font-family: Arial, sans-serif; line-height: 1.6; max-width: 800px; margin: 0 auto; padding: 20px; }
|
||||
h1 { color: #333; }
|
||||
p { margin-bottom: 15px; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<header>
|
||||
<h1>My First Blog Post</h1>
|
||||
</header>
|
||||
<main>
|
||||
<article>
|
||||
<p>This is the content of my blog post. It's short and sweet!</p>
|
||||
</article>
|
||||
</main>
|
||||
<footer>
|
||||
<p>© 2023 My Blog</p>
|
||||
</footer>
|
||||
</body>
|
||||
</html>
|
||||
\`\`\`
|
||||
:::
|
||||
|
||||
This HTML structure provides a simple layout for a blog post.
|
||||
|
||||
---`;
|
||||
|
||||
/**
|
||||
*
|
||||
* @param {Object} params
|
||||
* @param {EModelEndpoint | string} params.endpoint - The current endpoint
|
||||
* @param {ArtifactModes} params.artifacts - The current artifact mode
|
||||
* @returns
|
||||
*/
|
||||
const generateArtifactsPrompt = ({ endpoint, artifacts }) => {
|
||||
if (artifacts === ArtifactModes.CUSTOM) {
|
||||
return null;
|
||||
}
|
||||
|
||||
let prompt = artifactsPrompt;
|
||||
if (endpoint !== EModelEndpoint.anthropic) {
|
||||
prompt = artifactsOpenAIPrompt;
|
||||
}
|
||||
|
||||
if (artifacts === ArtifactModes.SHADCNUI) {
|
||||
prompt += generateShadcnPrompt({ components, useXML: endpoint === EModelEndpoint.anthropic });
|
||||
}
|
||||
|
||||
return prompt;
|
||||
};
|
||||
|
||||
module.exports = generateArtifactsPrompt;
|
||||
160
api/app/clients/prompts/createContextHandlers.js
Normal file
160
api/app/clients/prompts/createContextHandlers.js
Normal file
@ -0,0 +1,160 @@
|
||||
const axios = require('axios');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { isEnabled, generateShortLivedToken } = require('@librechat/api');
|
||||
|
||||
const footer = `Use the context as your learned knowledge to better answer the user.
|
||||
|
||||
In your response, remember to follow these guidelines:
|
||||
- If you don't know the answer, simply say that you don't know.
|
||||
- If you are unsure how to answer, ask for clarification.
|
||||
- Avoid mentioning that you obtained the information from the context.
|
||||
`;
|
||||
|
||||
function createContextHandlers(req, userMessageContent) {
|
||||
if (!process.env.RAG_API_URL) {
|
||||
return;
|
||||
}
|
||||
|
||||
const queryPromises = [];
|
||||
const processedFiles = [];
|
||||
const processedIds = new Set();
|
||||
const jwtToken = generateShortLivedToken(req.user.id);
|
||||
const useFullContext = isEnabled(process.env.RAG_USE_FULL_CONTEXT);
|
||||
|
||||
const query = async (file) => {
|
||||
if (useFullContext) {
|
||||
return axios.get(`${process.env.RAG_API_URL}/documents/${file.file_id}/context`, {
|
||||
headers: {
|
||||
Authorization: `Bearer ${jwtToken}`,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
return axios.post(
|
||||
`${process.env.RAG_API_URL}/query`,
|
||||
{
|
||||
file_id: file.file_id,
|
||||
query: userMessageContent,
|
||||
k: 4,
|
||||
},
|
||||
{
|
||||
headers: {
|
||||
Authorization: `Bearer ${jwtToken}`,
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
},
|
||||
);
|
||||
};
|
||||
|
||||
const processFile = async (file) => {
|
||||
if (file.embedded && !processedIds.has(file.file_id)) {
|
||||
try {
|
||||
const promise = query(file);
|
||||
queryPromises.push(promise);
|
||||
processedFiles.push(file);
|
||||
processedIds.add(file.file_id);
|
||||
} catch (error) {
|
||||
logger.error(`Error processing file ${file.filename}:`, error);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const createContext = async () => {
|
||||
try {
|
||||
if (!queryPromises.length || !processedFiles.length) {
|
||||
return '';
|
||||
}
|
||||
|
||||
const oneFile = processedFiles.length === 1;
|
||||
const header = `The user has attached ${oneFile ? 'a' : processedFiles.length} file${
|
||||
!oneFile ? 's' : ''
|
||||
} to the conversation:`;
|
||||
|
||||
const files = `${
|
||||
oneFile
|
||||
? ''
|
||||
: `
|
||||
<files>`
|
||||
}${processedFiles
|
||||
.map(
|
||||
(file) => `
|
||||
<file>
|
||||
<filename>${file.filename}</filename>
|
||||
<type>${file.type}</type>
|
||||
</file>`,
|
||||
)
|
||||
.join('')}${
|
||||
oneFile
|
||||
? ''
|
||||
: `
|
||||
</files>`
|
||||
}`;
|
||||
|
||||
const resolvedQueries = await Promise.all(queryPromises);
|
||||
|
||||
const context =
|
||||
resolvedQueries.length === 0
|
||||
? '\n\tThe semantic search did not return any results.'
|
||||
: resolvedQueries
|
||||
.map((queryResult, index) => {
|
||||
const file = processedFiles[index];
|
||||
let contextItems = queryResult.data;
|
||||
|
||||
const generateContext = (currentContext) =>
|
||||
`
|
||||
<file>
|
||||
<filename>${file.filename}</filename>
|
||||
<context>${currentContext}
|
||||
</context>
|
||||
</file>`;
|
||||
|
||||
if (useFullContext) {
|
||||
return generateContext(`\n${contextItems}`);
|
||||
}
|
||||
|
||||
contextItems = queryResult.data
|
||||
.map((item) => {
|
||||
const pageContent = item[0].page_content;
|
||||
return `
|
||||
<contextItem>
|
||||
<![CDATA[${pageContent?.trim()}]]>
|
||||
</contextItem>`;
|
||||
})
|
||||
.join('');
|
||||
|
||||
return generateContext(contextItems);
|
||||
})
|
||||
.join('');
|
||||
|
||||
if (useFullContext) {
|
||||
const prompt = `${header}
|
||||
${context}
|
||||
${footer}`;
|
||||
|
||||
return prompt;
|
||||
}
|
||||
|
||||
const prompt = `${header}
|
||||
${files}
|
||||
|
||||
A semantic search was executed with the user's message as the query, retrieving the following context inside <context></context> XML tags.
|
||||
|
||||
<context>${context}
|
||||
</context>
|
||||
|
||||
${footer}`;
|
||||
|
||||
return prompt;
|
||||
} catch (error) {
|
||||
logger.error('Error creating context:', error);
|
||||
throw error;
|
||||
}
|
||||
};
|
||||
|
||||
return {
|
||||
processFile,
|
||||
createContext,
|
||||
};
|
||||
}
|
||||
|
||||
module.exports = createContextHandlers;
|
||||
34
api/app/clients/prompts/createVisionPrompt.js
Normal file
34
api/app/clients/prompts/createVisionPrompt.js
Normal file
@ -0,0 +1,34 @@
|
||||
/**
|
||||
* Generates a prompt instructing the user to describe an image in detail, tailored to different types of visual content.
|
||||
* @param {boolean} pluralized - Whether to pluralize the prompt for multiple images.
|
||||
* @returns {string} - The generated vision prompt.
|
||||
*/
|
||||
const createVisionPrompt = (pluralized = false) => {
|
||||
return `Please describe the image${
|
||||
pluralized ? 's' : ''
|
||||
} in detail, covering relevant aspects such as:
|
||||
|
||||
For photographs, illustrations, or artwork:
|
||||
- The main subject(s) and their appearance, positioning, and actions
|
||||
- The setting, background, and any notable objects or elements
|
||||
- Colors, lighting, and overall mood or atmosphere
|
||||
- Any interesting details, textures, or patterns
|
||||
- The style, technique, or medium used (if discernible)
|
||||
|
||||
For screenshots or images containing text:
|
||||
- The content and purpose of the text
|
||||
- The layout, formatting, and organization of the information
|
||||
- Any notable visual elements, such as logos, icons, or graphics
|
||||
- The overall context or message conveyed by the screenshot
|
||||
|
||||
For graphs, charts, or data visualizations:
|
||||
- The type of graph or chart (e.g., bar graph, line chart, pie chart)
|
||||
- The variables being compared or analyzed
|
||||
- Any trends, patterns, or outliers in the data
|
||||
- The axis labels, scales, and units of measurement
|
||||
- The title, legend, and any additional context provided
|
||||
|
||||
Be as specific and descriptive as possible while maintaining clarity and concision.`;
|
||||
};
|
||||
|
||||
module.exports = createVisionPrompt;
|
||||
361
api/app/clients/prompts/formatAgentMessages.spec.js
Normal file
361
api/app/clients/prompts/formatAgentMessages.spec.js
Normal file
@ -0,0 +1,361 @@
|
||||
const { ToolMessage } = require('@langchain/core/messages');
|
||||
const { ContentTypes } = require('librechat-data-provider');
|
||||
const { HumanMessage, AIMessage, SystemMessage } = require('@langchain/core/messages');
|
||||
const { formatAgentMessages } = require('./formatMessages');
|
||||
|
||||
describe('formatAgentMessages', () => {
|
||||
it('should format simple user and AI messages', () => {
|
||||
const payload = [
|
||||
{ role: 'user', content: 'Hello' },
|
||||
{ role: 'assistant', content: 'Hi there!' },
|
||||
];
|
||||
const result = formatAgentMessages(payload);
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]).toBeInstanceOf(HumanMessage);
|
||||
expect(result[1]).toBeInstanceOf(AIMessage);
|
||||
});
|
||||
|
||||
it('should handle system messages', () => {
|
||||
const payload = [{ role: 'system', content: 'You are a helpful assistant.' }];
|
||||
const result = formatAgentMessages(payload);
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result[0]).toBeInstanceOf(SystemMessage);
|
||||
});
|
||||
|
||||
it('should format messages with content arrays', () => {
|
||||
const payload = [
|
||||
{
|
||||
role: 'user',
|
||||
content: [{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Hello' }],
|
||||
},
|
||||
];
|
||||
const result = formatAgentMessages(payload);
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result[0]).toBeInstanceOf(HumanMessage);
|
||||
});
|
||||
|
||||
it('should handle tool calls and create ToolMessages', () => {
|
||||
const payload = [
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [
|
||||
{
|
||||
type: ContentTypes.TEXT,
|
||||
[ContentTypes.TEXT]: 'Let me check that for you.',
|
||||
tool_call_ids: ['123'],
|
||||
},
|
||||
{
|
||||
type: ContentTypes.TOOL_CALL,
|
||||
tool_call: {
|
||||
id: '123',
|
||||
name: 'search',
|
||||
args: '{"query":"weather"}',
|
||||
output: 'The weather is sunny.',
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
];
|
||||
const result = formatAgentMessages(payload);
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]).toBeInstanceOf(AIMessage);
|
||||
expect(result[1]).toBeInstanceOf(ToolMessage);
|
||||
expect(result[0].tool_calls).toHaveLength(1);
|
||||
expect(result[1].tool_call_id).toBe('123');
|
||||
});
|
||||
|
||||
it('should handle multiple content parts in assistant messages', () => {
|
||||
const payload = [
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [
|
||||
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Part 1' },
|
||||
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Part 2' },
|
||||
],
|
||||
},
|
||||
];
|
||||
const result = formatAgentMessages(payload);
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result[0]).toBeInstanceOf(AIMessage);
|
||||
expect(result[0].content).toHaveLength(2);
|
||||
});
|
||||
|
||||
it('should throw an error for invalid tool call structure', () => {
|
||||
const payload = [
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [
|
||||
{
|
||||
type: ContentTypes.TOOL_CALL,
|
||||
tool_call: {
|
||||
id: '123',
|
||||
name: 'search',
|
||||
args: '{"query":"weather"}',
|
||||
output: 'The weather is sunny.',
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
];
|
||||
expect(() => formatAgentMessages(payload)).toThrow('Invalid tool call structure');
|
||||
});
|
||||
|
||||
it('should handle tool calls with non-JSON args', () => {
|
||||
const payload = [
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [
|
||||
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Checking...', tool_call_ids: ['123'] },
|
||||
{
|
||||
type: ContentTypes.TOOL_CALL,
|
||||
tool_call: {
|
||||
id: '123',
|
||||
name: 'search',
|
||||
args: 'non-json-string',
|
||||
output: 'Result',
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
];
|
||||
const result = formatAgentMessages(payload);
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0].tool_calls[0].args).toStrictEqual({ input: 'non-json-string' });
|
||||
});
|
||||
|
||||
it('should handle complex tool calls with multiple steps', () => {
|
||||
const payload = [
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [
|
||||
{
|
||||
type: ContentTypes.TEXT,
|
||||
[ContentTypes.TEXT]: "I'll search for that information.",
|
||||
tool_call_ids: ['search_1'],
|
||||
},
|
||||
{
|
||||
type: ContentTypes.TOOL_CALL,
|
||||
tool_call: {
|
||||
id: 'search_1',
|
||||
name: 'search',
|
||||
args: '{"query":"weather in New York"}',
|
||||
output: 'The weather in New York is currently sunny with a temperature of 75°F.',
|
||||
},
|
||||
},
|
||||
{
|
||||
type: ContentTypes.TEXT,
|
||||
[ContentTypes.TEXT]: "Now, I'll convert the temperature.",
|
||||
tool_call_ids: ['convert_1'],
|
||||
},
|
||||
{
|
||||
type: ContentTypes.TOOL_CALL,
|
||||
tool_call: {
|
||||
id: 'convert_1',
|
||||
name: 'convert_temperature',
|
||||
args: '{"temperature": 75, "from": "F", "to": "C"}',
|
||||
output: '23.89°C',
|
||||
},
|
||||
},
|
||||
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: "Here's your answer." },
|
||||
],
|
||||
},
|
||||
];
|
||||
|
||||
const result = formatAgentMessages(payload);
|
||||
|
||||
expect(result).toHaveLength(5);
|
||||
expect(result[0]).toBeInstanceOf(AIMessage);
|
||||
expect(result[1]).toBeInstanceOf(ToolMessage);
|
||||
expect(result[2]).toBeInstanceOf(AIMessage);
|
||||
expect(result[3]).toBeInstanceOf(ToolMessage);
|
||||
expect(result[4]).toBeInstanceOf(AIMessage);
|
||||
|
||||
// Check first AIMessage
|
||||
expect(result[0].content).toBe("I'll search for that information.");
|
||||
expect(result[0].tool_calls).toHaveLength(1);
|
||||
expect(result[0].tool_calls[0]).toEqual({
|
||||
id: 'search_1',
|
||||
name: 'search',
|
||||
args: { query: 'weather in New York' },
|
||||
});
|
||||
|
||||
// Check first ToolMessage
|
||||
expect(result[1].tool_call_id).toBe('search_1');
|
||||
expect(result[1].name).toBe('search');
|
||||
expect(result[1].content).toBe(
|
||||
'The weather in New York is currently sunny with a temperature of 75°F.',
|
||||
);
|
||||
|
||||
// Check second AIMessage
|
||||
expect(result[2].content).toBe("Now, I'll convert the temperature.");
|
||||
expect(result[2].tool_calls).toHaveLength(1);
|
||||
expect(result[2].tool_calls[0]).toEqual({
|
||||
id: 'convert_1',
|
||||
name: 'convert_temperature',
|
||||
args: { temperature: 75, from: 'F', to: 'C' },
|
||||
});
|
||||
|
||||
// Check second ToolMessage
|
||||
expect(result[3].tool_call_id).toBe('convert_1');
|
||||
expect(result[3].name).toBe('convert_temperature');
|
||||
expect(result[3].content).toBe('23.89°C');
|
||||
|
||||
// Check final AIMessage
|
||||
expect(result[4].content).toStrictEqual([
|
||||
{ [ContentTypes.TEXT]: "Here's your answer.", type: ContentTypes.TEXT },
|
||||
]);
|
||||
});
|
||||
|
||||
it.skip('should not produce two consecutive assistant messages and format content correctly', () => {
|
||||
const payload = [
|
||||
{ role: 'user', content: 'Hello' },
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Hi there!' }],
|
||||
},
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'How can I help you?' }],
|
||||
},
|
||||
{ role: 'user', content: "What's the weather?" },
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [
|
||||
{
|
||||
type: ContentTypes.TEXT,
|
||||
[ContentTypes.TEXT]: 'Let me check that for you.',
|
||||
tool_call_ids: ['weather_1'],
|
||||
},
|
||||
{
|
||||
type: ContentTypes.TOOL_CALL,
|
||||
tool_call: {
|
||||
id: 'weather_1',
|
||||
name: 'check_weather',
|
||||
args: '{"location":"New York"}',
|
||||
output: 'Sunny, 75°F',
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [
|
||||
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: "Here's the weather information." },
|
||||
],
|
||||
},
|
||||
];
|
||||
|
||||
const result = formatAgentMessages(payload);
|
||||
|
||||
// Check correct message count and types
|
||||
expect(result).toHaveLength(6);
|
||||
expect(result[0]).toBeInstanceOf(HumanMessage);
|
||||
expect(result[1]).toBeInstanceOf(AIMessage);
|
||||
expect(result[2]).toBeInstanceOf(HumanMessage);
|
||||
expect(result[3]).toBeInstanceOf(AIMessage);
|
||||
expect(result[4]).toBeInstanceOf(ToolMessage);
|
||||
expect(result[5]).toBeInstanceOf(AIMessage);
|
||||
|
||||
// Check content of messages
|
||||
expect(result[0].content).toStrictEqual([
|
||||
{ [ContentTypes.TEXT]: 'Hello', type: ContentTypes.TEXT },
|
||||
]);
|
||||
expect(result[1].content).toStrictEqual([
|
||||
{ [ContentTypes.TEXT]: 'Hi there!', type: ContentTypes.TEXT },
|
||||
{ [ContentTypes.TEXT]: 'How can I help you?', type: ContentTypes.TEXT },
|
||||
]);
|
||||
expect(result[2].content).toStrictEqual([
|
||||
{ [ContentTypes.TEXT]: "What's the weather?", type: ContentTypes.TEXT },
|
||||
]);
|
||||
expect(result[3].content).toBe('Let me check that for you.');
|
||||
expect(result[4].content).toBe('Sunny, 75°F');
|
||||
expect(result[5].content).toStrictEqual([
|
||||
{ [ContentTypes.TEXT]: "Here's the weather information.", type: ContentTypes.TEXT },
|
||||
]);
|
||||
|
||||
// Check that there are no consecutive AIMessages
|
||||
const messageTypes = result.map((message) => message.constructor);
|
||||
for (let i = 0; i < messageTypes.length - 1; i++) {
|
||||
expect(messageTypes[i] === AIMessage && messageTypes[i + 1] === AIMessage).toBe(false);
|
||||
}
|
||||
|
||||
// Additional check to ensure the consecutive assistant messages were combined
|
||||
expect(result[1].content).toHaveLength(2);
|
||||
});
|
||||
|
||||
it('should skip THINK type content parts', () => {
|
||||
const payload = [
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [
|
||||
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Initial response' },
|
||||
{ type: ContentTypes.THINK, [ContentTypes.THINK]: 'Reasoning about the problem...' },
|
||||
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Final answer' },
|
||||
],
|
||||
},
|
||||
];
|
||||
|
||||
const result = formatAgentMessages(payload);
|
||||
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result[0]).toBeInstanceOf(AIMessage);
|
||||
expect(result[0].content).toEqual('Initial response\nFinal answer');
|
||||
});
|
||||
|
||||
it('should join TEXT content as string when THINK content type is present', () => {
|
||||
const payload = [
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [
|
||||
{ type: ContentTypes.THINK, [ContentTypes.THINK]: 'Analyzing the problem...' },
|
||||
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'First part of response' },
|
||||
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Second part of response' },
|
||||
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Final part of response' },
|
||||
],
|
||||
},
|
||||
];
|
||||
|
||||
const result = formatAgentMessages(payload);
|
||||
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result[0]).toBeInstanceOf(AIMessage);
|
||||
expect(typeof result[0].content).toBe('string');
|
||||
expect(result[0].content).toBe(
|
||||
'First part of response\nSecond part of response\nFinal part of response',
|
||||
);
|
||||
expect(result[0].content).not.toContain('Analyzing the problem...');
|
||||
});
|
||||
|
||||
it('should exclude ERROR type content parts', () => {
|
||||
const payload = [
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [
|
||||
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Hello there' },
|
||||
{
|
||||
type: ContentTypes.ERROR,
|
||||
[ContentTypes.ERROR]:
|
||||
'An error occurred while processing the request: Something went wrong',
|
||||
},
|
||||
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Final answer' },
|
||||
],
|
||||
},
|
||||
];
|
||||
|
||||
const result = formatAgentMessages(payload);
|
||||
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result[0]).toBeInstanceOf(AIMessage);
|
||||
expect(result[0].content).toEqual([
|
||||
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Hello there' },
|
||||
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Final answer' },
|
||||
]);
|
||||
|
||||
// Make sure no error content exists in the result
|
||||
const hasErrorContent = result[0].content.some(
|
||||
(item) =>
|
||||
item.type === ContentTypes.ERROR || JSON.stringify(item).includes('An error occurred'),
|
||||
);
|
||||
expect(hasErrorContent).toBe(false);
|
||||
});
|
||||
});
|
||||
42
api/app/clients/prompts/formatGoogleInputs.js
Normal file
42
api/app/clients/prompts/formatGoogleInputs.js
Normal file
@ -0,0 +1,42 @@
|
||||
/**
|
||||
* Formats an object to match the struct_val, list_val, string_val, float_val, and int_val format.
|
||||
*
|
||||
* @param {Object} obj - The object to be formatted.
|
||||
* @returns {Object} The formatted object.
|
||||
*
|
||||
* Handles different types:
|
||||
* - Arrays are wrapped in list_val and each element is processed.
|
||||
* - Objects are recursively processed.
|
||||
* - Strings are wrapped in string_val.
|
||||
* - Numbers are wrapped in float_val or int_val depending on whether they are floating-point or integers.
|
||||
*/
|
||||
function formatGoogleInputs(obj) {
|
||||
const formattedObj = {};
|
||||
|
||||
for (const key in obj) {
|
||||
if (Object.prototype.hasOwnProperty.call(obj, key)) {
|
||||
const value = obj[key];
|
||||
|
||||
// Handle arrays
|
||||
if (Array.isArray(value)) {
|
||||
formattedObj[key] = { list_val: value.map((item) => formatGoogleInputs(item)) };
|
||||
}
|
||||
// Handle objects
|
||||
else if (typeof value === 'object' && value !== null) {
|
||||
formattedObj[key] = formatGoogleInputs(value);
|
||||
}
|
||||
// Handle numbers
|
||||
else if (typeof value === 'number') {
|
||||
formattedObj[key] = Number.isInteger(value) ? { int_val: value } : { float_val: value };
|
||||
}
|
||||
// Handle other types (e.g., strings)
|
||||
else {
|
||||
formattedObj[key] = { string_val: [value] };
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return { struct_val: formattedObj };
|
||||
}
|
||||
|
||||
module.exports = formatGoogleInputs;
|
||||
274
api/app/clients/prompts/formatGoogleInputs.spec.js
Normal file
274
api/app/clients/prompts/formatGoogleInputs.spec.js
Normal file
@ -0,0 +1,274 @@
|
||||
const formatGoogleInputs = require('./formatGoogleInputs');
|
||||
|
||||
describe('formatGoogleInputs', () => {
|
||||
it('formats message correctly', () => {
|
||||
const input = {
|
||||
messages: [
|
||||
{
|
||||
content: 'hi',
|
||||
author: 'user',
|
||||
},
|
||||
],
|
||||
context: 'context',
|
||||
examples: [
|
||||
{
|
||||
input: {
|
||||
author: 'user',
|
||||
content: 'user input',
|
||||
},
|
||||
output: {
|
||||
author: 'bot',
|
||||
content: 'bot output',
|
||||
},
|
||||
},
|
||||
],
|
||||
parameters: {
|
||||
temperature: 0.2,
|
||||
topP: 0.8,
|
||||
topK: 40,
|
||||
maxOutputTokens: 1024,
|
||||
},
|
||||
};
|
||||
|
||||
const expectedOutput = {
|
||||
struct_val: {
|
||||
messages: {
|
||||
list_val: [
|
||||
{
|
||||
struct_val: {
|
||||
content: {
|
||||
string_val: ['hi'],
|
||||
},
|
||||
author: {
|
||||
string_val: ['user'],
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
context: {
|
||||
string_val: ['context'],
|
||||
},
|
||||
examples: {
|
||||
list_val: [
|
||||
{
|
||||
struct_val: {
|
||||
input: {
|
||||
struct_val: {
|
||||
author: {
|
||||
string_val: ['user'],
|
||||
},
|
||||
content: {
|
||||
string_val: ['user input'],
|
||||
},
|
||||
},
|
||||
},
|
||||
output: {
|
||||
struct_val: {
|
||||
author: {
|
||||
string_val: ['bot'],
|
||||
},
|
||||
content: {
|
||||
string_val: ['bot output'],
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
parameters: {
|
||||
struct_val: {
|
||||
temperature: {
|
||||
float_val: 0.2,
|
||||
},
|
||||
topP: {
|
||||
float_val: 0.8,
|
||||
},
|
||||
topK: {
|
||||
int_val: 40,
|
||||
},
|
||||
maxOutputTokens: {
|
||||
int_val: 1024,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const result = formatGoogleInputs(input);
|
||||
expect(JSON.stringify(result)).toEqual(JSON.stringify(expectedOutput));
|
||||
});
|
||||
|
||||
it('formats real payload parts', () => {
|
||||
const input = {
|
||||
instances: [
|
||||
{
|
||||
context: 'context',
|
||||
examples: [
|
||||
{
|
||||
input: {
|
||||
author: 'user',
|
||||
content: 'user input',
|
||||
},
|
||||
output: {
|
||||
author: 'bot',
|
||||
content: 'user output',
|
||||
},
|
||||
},
|
||||
],
|
||||
messages: [
|
||||
{
|
||||
author: 'user',
|
||||
content: 'hi',
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
parameters: {
|
||||
candidateCount: 1,
|
||||
maxOutputTokens: 1024,
|
||||
temperature: 0.2,
|
||||
topP: 0.8,
|
||||
topK: 40,
|
||||
},
|
||||
};
|
||||
const expectedOutput = {
|
||||
struct_val: {
|
||||
instances: {
|
||||
list_val: [
|
||||
{
|
||||
struct_val: {
|
||||
context: { string_val: ['context'] },
|
||||
examples: {
|
||||
list_val: [
|
||||
{
|
||||
struct_val: {
|
||||
input: {
|
||||
struct_val: {
|
||||
author: { string_val: ['user'] },
|
||||
content: { string_val: ['user input'] },
|
||||
},
|
||||
},
|
||||
output: {
|
||||
struct_val: {
|
||||
author: { string_val: ['bot'] },
|
||||
content: { string_val: ['user output'] },
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
messages: {
|
||||
list_val: [
|
||||
{
|
||||
struct_val: {
|
||||
author: { string_val: ['user'] },
|
||||
content: { string_val: ['hi'] },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
parameters: {
|
||||
struct_val: {
|
||||
candidateCount: { int_val: 1 },
|
||||
maxOutputTokens: { int_val: 1024 },
|
||||
temperature: { float_val: 0.2 },
|
||||
topP: { float_val: 0.8 },
|
||||
topK: { int_val: 40 },
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const result = formatGoogleInputs(input);
|
||||
expect(JSON.stringify(result)).toEqual(JSON.stringify(expectedOutput));
|
||||
});
|
||||
|
||||
it('helps create valid payload parts', () => {
|
||||
const instances = {
|
||||
context: 'context',
|
||||
examples: [
|
||||
{
|
||||
input: {
|
||||
author: 'user',
|
||||
content: 'user input',
|
||||
},
|
||||
output: {
|
||||
author: 'bot',
|
||||
content: 'user output',
|
||||
},
|
||||
},
|
||||
],
|
||||
messages: [
|
||||
{
|
||||
author: 'user',
|
||||
content: 'hi',
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const expectedInstances = {
|
||||
struct_val: {
|
||||
context: { string_val: ['context'] },
|
||||
examples: {
|
||||
list_val: [
|
||||
{
|
||||
struct_val: {
|
||||
input: {
|
||||
struct_val: {
|
||||
author: { string_val: ['user'] },
|
||||
content: { string_val: ['user input'] },
|
||||
},
|
||||
},
|
||||
output: {
|
||||
struct_val: {
|
||||
author: { string_val: ['bot'] },
|
||||
content: { string_val: ['user output'] },
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
messages: {
|
||||
list_val: [
|
||||
{
|
||||
struct_val: {
|
||||
author: { string_val: ['user'] },
|
||||
content: { string_val: ['hi'] },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const parameters = {
|
||||
candidateCount: 1,
|
||||
maxOutputTokens: 1024,
|
||||
temperature: 0.2,
|
||||
topP: 0.8,
|
||||
topK: 40,
|
||||
};
|
||||
const expectedParameters = {
|
||||
struct_val: {
|
||||
candidateCount: { int_val: 1 },
|
||||
maxOutputTokens: { int_val: 1024 },
|
||||
temperature: { float_val: 0.2 },
|
||||
topP: { float_val: 0.8 },
|
||||
topK: { int_val: 40 },
|
||||
},
|
||||
};
|
||||
|
||||
const instancesResult = formatGoogleInputs(instances);
|
||||
const parametersResult = formatGoogleInputs(parameters);
|
||||
expect(JSON.stringify(instancesResult)).toEqual(JSON.stringify(expectedInstances));
|
||||
expect(JSON.stringify(parametersResult)).toEqual(JSON.stringify(expectedParameters));
|
||||
});
|
||||
});
|
||||
245
api/app/clients/prompts/formatMessages.js
Normal file
245
api/app/clients/prompts/formatMessages.js
Normal file
@ -0,0 +1,245 @@
|
||||
const { ToolMessage } = require('@langchain/core/messages');
|
||||
const { EModelEndpoint, ContentTypes } = require('librechat-data-provider');
|
||||
const { HumanMessage, AIMessage, SystemMessage } = require('@langchain/core/messages');
|
||||
|
||||
/**
|
||||
* Formats a message to OpenAI Vision API payload format.
|
||||
*
|
||||
* @param {Object} params - The parameters for formatting.
|
||||
* @param {Object} params.message - The message object to format.
|
||||
* @param {string} [params.message.role] - The role of the message sender (must be 'user').
|
||||
* @param {string} [params.message.content] - The text content of the message.
|
||||
* @param {EModelEndpoint} [params.endpoint] - Identifier for specific endpoint handling
|
||||
* @param {Array<string>} [params.image_urls] - The image_urls to attach to the message.
|
||||
* @returns {(Object)} - The formatted message.
|
||||
*/
|
||||
const formatVisionMessage = ({ message, image_urls, endpoint }) => {
|
||||
if (endpoint === EModelEndpoint.anthropic) {
|
||||
message.content = [...image_urls, { type: ContentTypes.TEXT, text: message.content }];
|
||||
return message;
|
||||
}
|
||||
|
||||
message.content = [{ type: ContentTypes.TEXT, text: message.content }, ...image_urls];
|
||||
|
||||
return message;
|
||||
};
|
||||
|
||||
/**
|
||||
* Formats a message to OpenAI payload format based on the provided options.
|
||||
*
|
||||
* @param {Object} params - The parameters for formatting.
|
||||
* @param {Object} params.message - The message object to format.
|
||||
* @param {string} [params.message.role] - The role of the message sender (e.g., 'user', 'assistant').
|
||||
* @param {string} [params.message._name] - The name associated with the message.
|
||||
* @param {string} [params.message.sender] - The sender of the message.
|
||||
* @param {string} [params.message.text] - The text content of the message.
|
||||
* @param {string} [params.message.content] - The content of the message.
|
||||
* @param {Array<string>} [params.message.image_urls] - The image_urls attached to the message for Vision API.
|
||||
* @param {string} [params.userName] - The name of the user.
|
||||
* @param {string} [params.assistantName] - The name of the assistant.
|
||||
* @param {string} [params.endpoint] - Identifier for specific endpoint handling
|
||||
* @param {boolean} [params.langChain=false] - Whether to return a LangChain message object.
|
||||
* @returns {(Object|HumanMessage|AIMessage|SystemMessage)} - The formatted message.
|
||||
*/
|
||||
const formatMessage = ({ message, userName, assistantName, endpoint, langChain = false }) => {
|
||||
let { role: _role, _name, sender, text, content: _content, lc_id } = message;
|
||||
if (lc_id && lc_id[2] && !langChain) {
|
||||
const roleMapping = {
|
||||
SystemMessage: 'system',
|
||||
HumanMessage: 'user',
|
||||
AIMessage: 'assistant',
|
||||
};
|
||||
_role = roleMapping[lc_id[2]];
|
||||
}
|
||||
const role = _role ?? (sender && sender?.toLowerCase() === 'user' ? 'user' : 'assistant');
|
||||
const content = _content ?? text ?? '';
|
||||
const formattedMessage = {
|
||||
role,
|
||||
content,
|
||||
};
|
||||
|
||||
const { image_urls } = message;
|
||||
if (Array.isArray(image_urls) && image_urls.length > 0 && role === 'user') {
|
||||
return formatVisionMessage({
|
||||
message: formattedMessage,
|
||||
image_urls: message.image_urls,
|
||||
endpoint,
|
||||
});
|
||||
}
|
||||
|
||||
if (_name) {
|
||||
formattedMessage.name = _name;
|
||||
}
|
||||
|
||||
if (userName && formattedMessage.role === 'user') {
|
||||
formattedMessage.name = userName;
|
||||
}
|
||||
|
||||
if (assistantName && formattedMessage.role === 'assistant') {
|
||||
formattedMessage.name = assistantName;
|
||||
}
|
||||
|
||||
if (formattedMessage.name) {
|
||||
// Conform to API regex: ^[a-zA-Z0-9_-]{1,64}$
|
||||
// https://community.openai.com/t/the-format-of-the-name-field-in-the-documentation-is-incorrect/175684/2
|
||||
formattedMessage.name = formattedMessage.name.replace(/[^a-zA-Z0-9_-]/g, '_');
|
||||
|
||||
if (formattedMessage.name.length > 64) {
|
||||
formattedMessage.name = formattedMessage.name.substring(0, 64);
|
||||
}
|
||||
}
|
||||
|
||||
if (!langChain) {
|
||||
return formattedMessage;
|
||||
}
|
||||
|
||||
if (role === 'user') {
|
||||
return new HumanMessage(formattedMessage);
|
||||
} else if (role === 'assistant') {
|
||||
return new AIMessage(formattedMessage);
|
||||
} else {
|
||||
return new SystemMessage(formattedMessage);
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Formats an array of messages for LangChain.
|
||||
*
|
||||
* @param {Array<Object>} messages - The array of messages to format.
|
||||
* @param {Object} formatOptions - The options for formatting each message.
|
||||
* @param {string} [formatOptions.userName] - The name of the user.
|
||||
* @param {string} [formatOptions.assistantName] - The name of the assistant.
|
||||
* @returns {Array<(HumanMessage|AIMessage|SystemMessage)>} - The array of formatted LangChain messages.
|
||||
*/
|
||||
const formatLangChainMessages = (messages, formatOptions) =>
|
||||
messages.map((msg) => formatMessage({ ...formatOptions, message: msg, langChain: true }));
|
||||
|
||||
/**
|
||||
* Formats a LangChain message object by merging properties from `lc_kwargs` or `kwargs` and `additional_kwargs`.
|
||||
*
|
||||
* @param {Object} message - The message object to format.
|
||||
* @param {Object} [message.lc_kwargs] - Contains properties to be merged. Either this or `message.kwargs` should be provided.
|
||||
* @param {Object} [message.kwargs] - Contains properties to be merged. Either this or `message.lc_kwargs` should be provided.
|
||||
* @param {Object} [message.kwargs.additional_kwargs] - Additional properties to be merged.
|
||||
*
|
||||
* @returns {Object} The formatted LangChain message.
|
||||
*/
|
||||
const formatFromLangChain = (message) => {
|
||||
const { additional_kwargs, ...message_kwargs } = message.lc_kwargs ?? message.kwargs;
|
||||
return {
|
||||
...message_kwargs,
|
||||
...additional_kwargs,
|
||||
};
|
||||
};
|
||||
|
||||
/**
|
||||
* Formats an array of messages for LangChain, handling tool calls and creating ToolMessage instances.
|
||||
*
|
||||
* @param {Array<Partial<TMessage>>} payload - The array of messages to format.
|
||||
* @returns {Array<(HumanMessage|AIMessage|SystemMessage|ToolMessage)>} - The array of formatted LangChain messages, including ToolMessages for tool calls.
|
||||
*/
|
||||
const formatAgentMessages = (payload) => {
|
||||
const messages = [];
|
||||
|
||||
for (const message of payload) {
|
||||
if (typeof message.content === 'string') {
|
||||
message.content = [{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: message.content }];
|
||||
}
|
||||
if (message.role !== 'assistant') {
|
||||
messages.push(formatMessage({ message, langChain: true }));
|
||||
continue;
|
||||
}
|
||||
|
||||
let currentContent = [];
|
||||
let lastAIMessage = null;
|
||||
|
||||
let hasReasoning = false;
|
||||
for (const part of message.content) {
|
||||
if (part.type === ContentTypes.TEXT && part.tool_call_ids) {
|
||||
/*
|
||||
If there's pending content, it needs to be aggregated as a single string to prepare for tool calls.
|
||||
For Anthropic models, the "tool_calls" field on a message is only respected if content is a string.
|
||||
*/
|
||||
if (currentContent.length > 0) {
|
||||
let content = currentContent.reduce((acc, curr) => {
|
||||
if (curr.type === ContentTypes.TEXT) {
|
||||
return `${acc}${curr[ContentTypes.TEXT]}\n`;
|
||||
}
|
||||
return acc;
|
||||
}, '');
|
||||
content = `${content}\n${part[ContentTypes.TEXT] ?? ''}`.trim();
|
||||
lastAIMessage = new AIMessage({ content });
|
||||
messages.push(lastAIMessage);
|
||||
currentContent = [];
|
||||
continue;
|
||||
}
|
||||
|
||||
// Create a new AIMessage with this text and prepare for tool calls
|
||||
lastAIMessage = new AIMessage({
|
||||
content: part.text || '',
|
||||
});
|
||||
|
||||
messages.push(lastAIMessage);
|
||||
} else if (part.type === ContentTypes.TOOL_CALL) {
|
||||
if (!lastAIMessage) {
|
||||
throw new Error('Invalid tool call structure: No preceding AIMessage with tool_call_ids');
|
||||
}
|
||||
|
||||
// Note: `tool_calls` list is defined when constructed by `AIMessage` class, and outputs should be excluded from it
|
||||
const { output, args: _args, ...tool_call } = part.tool_call;
|
||||
// TODO: investigate; args as dictionary may need to be provider-or-tool-specific
|
||||
let args = _args;
|
||||
try {
|
||||
args = JSON.parse(_args);
|
||||
} catch (e) {
|
||||
if (typeof _args === 'string') {
|
||||
args = { input: _args };
|
||||
}
|
||||
}
|
||||
|
||||
tool_call.args = args;
|
||||
lastAIMessage.tool_calls.push(tool_call);
|
||||
|
||||
// Add the corresponding ToolMessage
|
||||
messages.push(
|
||||
new ToolMessage({
|
||||
tool_call_id: tool_call.id,
|
||||
name: tool_call.name,
|
||||
content: output || '',
|
||||
}),
|
||||
);
|
||||
} else if (part.type === ContentTypes.THINK) {
|
||||
hasReasoning = true;
|
||||
continue;
|
||||
} else if (part.type === ContentTypes.ERROR || part.type === ContentTypes.AGENT_UPDATE) {
|
||||
continue;
|
||||
} else {
|
||||
currentContent.push(part);
|
||||
}
|
||||
}
|
||||
|
||||
if (hasReasoning) {
|
||||
currentContent = currentContent
|
||||
.reduce((acc, curr) => {
|
||||
if (curr.type === ContentTypes.TEXT) {
|
||||
return `${acc}${curr[ContentTypes.TEXT]}\n`;
|
||||
}
|
||||
return acc;
|
||||
}, '')
|
||||
.trim();
|
||||
}
|
||||
|
||||
if (currentContent.length > 0) {
|
||||
messages.push(new AIMessage({ content: currentContent }));
|
||||
}
|
||||
}
|
||||
|
||||
return messages;
|
||||
};
|
||||
|
||||
module.exports = {
|
||||
formatMessage,
|
||||
formatFromLangChain,
|
||||
formatAgentMessages,
|
||||
formatLangChainMessages,
|
||||
};
|
||||
276
api/app/clients/prompts/formatMessages.spec.js
Normal file
276
api/app/clients/prompts/formatMessages.spec.js
Normal file
@ -0,0 +1,276 @@
|
||||
const { Constants } = require('librechat-data-provider');
|
||||
const { HumanMessage, AIMessage, SystemMessage } = require('@langchain/core/messages');
|
||||
const { formatMessage, formatLangChainMessages, formatFromLangChain } = require('./formatMessages');
|
||||
|
||||
describe('formatMessage', () => {
|
||||
it('formats user message', () => {
|
||||
const input = {
|
||||
message: {
|
||||
sender: 'user',
|
||||
text: 'Hello',
|
||||
},
|
||||
userName: 'John',
|
||||
};
|
||||
const result = formatMessage(input);
|
||||
expect(result).toEqual({
|
||||
role: 'user',
|
||||
content: 'Hello',
|
||||
name: 'John',
|
||||
});
|
||||
});
|
||||
|
||||
it('sanitizes the name by replacing invalid characters (per OpenAI)', () => {
|
||||
const input = {
|
||||
message: {
|
||||
sender: 'user',
|
||||
text: 'Hello',
|
||||
},
|
||||
userName: ' John$Doe@Example! ',
|
||||
};
|
||||
const result = formatMessage(input);
|
||||
expect(result).toEqual({
|
||||
role: 'user',
|
||||
content: 'Hello',
|
||||
name: '_John_Doe_Example__',
|
||||
});
|
||||
});
|
||||
|
||||
it('trims the name to a maximum length of 64 characters', () => {
|
||||
const longName = 'a'.repeat(100);
|
||||
const input = {
|
||||
message: {
|
||||
sender: 'user',
|
||||
text: 'Hello',
|
||||
},
|
||||
userName: longName,
|
||||
};
|
||||
const result = formatMessage(input);
|
||||
expect(result.name.length).toBe(64);
|
||||
expect(result.name).toBe('a'.repeat(64));
|
||||
});
|
||||
|
||||
it('formats a realistic user message', () => {
|
||||
const input = {
|
||||
message: {
|
||||
_id: '6512cdfb92cbf69fea615331',
|
||||
messageId: 'b620bf73-c5c3-4a38-b724-76886aac24c4',
|
||||
__v: 0,
|
||||
conversationId: '5c23d24f-941f-4aab-85df-127b596c8aa5',
|
||||
createdAt: Date.now(),
|
||||
error: false,
|
||||
finish_reason: null,
|
||||
isCreatedByUser: true,
|
||||
model: null,
|
||||
parentMessageId: Constants.NO_PARENT,
|
||||
sender: 'User',
|
||||
text: 'hi',
|
||||
tokenCount: 5,
|
||||
unfinished: false,
|
||||
updatedAt: Date.now(),
|
||||
user: '6512cdf475f05c86d44c31d2',
|
||||
},
|
||||
userName: 'John',
|
||||
};
|
||||
const result = formatMessage(input);
|
||||
expect(result).toEqual({
|
||||
role: 'user',
|
||||
content: 'hi',
|
||||
name: 'John',
|
||||
});
|
||||
});
|
||||
|
||||
it('formats assistant message', () => {
|
||||
const input = {
|
||||
message: {
|
||||
sender: 'assistant',
|
||||
text: 'Hi there',
|
||||
},
|
||||
assistantName: 'Assistant',
|
||||
};
|
||||
const result = formatMessage(input);
|
||||
expect(result).toEqual({
|
||||
role: 'assistant',
|
||||
content: 'Hi there',
|
||||
name: 'Assistant',
|
||||
});
|
||||
});
|
||||
|
||||
it('formats system message', () => {
|
||||
const input = {
|
||||
message: {
|
||||
role: 'system',
|
||||
text: 'Hi there',
|
||||
},
|
||||
};
|
||||
const result = formatMessage(input);
|
||||
expect(result).toEqual({
|
||||
role: 'system',
|
||||
content: 'Hi there',
|
||||
});
|
||||
});
|
||||
|
||||
it('formats user message with langChain', () => {
|
||||
const input = {
|
||||
message: {
|
||||
sender: 'user',
|
||||
text: 'Hello',
|
||||
},
|
||||
userName: 'John',
|
||||
langChain: true,
|
||||
};
|
||||
const result = formatMessage(input);
|
||||
expect(result).toBeInstanceOf(HumanMessage);
|
||||
expect(result.lc_kwargs.content).toEqual(input.message.text);
|
||||
expect(result.lc_kwargs.name).toEqual(input.userName);
|
||||
});
|
||||
|
||||
it('formats assistant message with langChain', () => {
|
||||
const input = {
|
||||
message: {
|
||||
sender: 'assistant',
|
||||
text: 'Hi there',
|
||||
},
|
||||
assistantName: 'Assistant',
|
||||
langChain: true,
|
||||
};
|
||||
const result = formatMessage(input);
|
||||
expect(result).toBeInstanceOf(AIMessage);
|
||||
expect(result.lc_kwargs.content).toEqual(input.message.text);
|
||||
expect(result.lc_kwargs.name).toEqual(input.assistantName);
|
||||
});
|
||||
|
||||
it('formats system message with langChain', () => {
|
||||
const input = {
|
||||
message: {
|
||||
role: 'system',
|
||||
text: 'This is a system message.',
|
||||
},
|
||||
langChain: true,
|
||||
};
|
||||
const result = formatMessage(input);
|
||||
expect(result).toBeInstanceOf(SystemMessage);
|
||||
expect(result.lc_kwargs.content).toEqual(input.message.text);
|
||||
});
|
||||
|
||||
it('formats langChain messages into OpenAI payload format', () => {
|
||||
const human = {
|
||||
message: new HumanMessage({
|
||||
content: 'Hello',
|
||||
}),
|
||||
};
|
||||
const system = {
|
||||
message: new SystemMessage({
|
||||
content: 'Hello',
|
||||
}),
|
||||
};
|
||||
const ai = {
|
||||
message: new AIMessage({
|
||||
content: 'Hello',
|
||||
}),
|
||||
};
|
||||
const humanResult = formatMessage(human);
|
||||
const systemResult = formatMessage(system);
|
||||
const aiResult = formatMessage(ai);
|
||||
expect(humanResult).toEqual({
|
||||
role: 'user',
|
||||
content: 'Hello',
|
||||
});
|
||||
expect(systemResult).toEqual({
|
||||
role: 'system',
|
||||
content: 'Hello',
|
||||
});
|
||||
expect(aiResult).toEqual({
|
||||
role: 'assistant',
|
||||
content: 'Hello',
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('formatLangChainMessages', () => {
|
||||
it('formats an array of messages for LangChain', () => {
|
||||
const messages = [
|
||||
{
|
||||
role: 'system',
|
||||
content: 'This is a system message',
|
||||
},
|
||||
{
|
||||
sender: 'user',
|
||||
text: 'Hello',
|
||||
},
|
||||
{
|
||||
sender: 'assistant',
|
||||
text: 'Hi there',
|
||||
},
|
||||
];
|
||||
const formatOptions = {
|
||||
userName: 'John',
|
||||
assistantName: 'Assistant',
|
||||
};
|
||||
const result = formatLangChainMessages(messages, formatOptions);
|
||||
expect(result).toHaveLength(3);
|
||||
expect(result[0]).toBeInstanceOf(SystemMessage);
|
||||
expect(result[1]).toBeInstanceOf(HumanMessage);
|
||||
expect(result[2]).toBeInstanceOf(AIMessage);
|
||||
|
||||
expect(result[0].lc_kwargs.content).toEqual(messages[0].content);
|
||||
expect(result[1].lc_kwargs.content).toEqual(messages[1].text);
|
||||
expect(result[2].lc_kwargs.content).toEqual(messages[2].text);
|
||||
|
||||
expect(result[1].lc_kwargs.name).toEqual(formatOptions.userName);
|
||||
expect(result[2].lc_kwargs.name).toEqual(formatOptions.assistantName);
|
||||
});
|
||||
|
||||
describe('formatFromLangChain', () => {
|
||||
it('should merge kwargs and additional_kwargs', () => {
|
||||
const message = {
|
||||
kwargs: {
|
||||
content: 'some content',
|
||||
name: 'dan',
|
||||
additional_kwargs: {
|
||||
function_call: {
|
||||
name: 'dall-e',
|
||||
arguments: '{\n "input": "Subject: hedgehog, Style: cute"\n}',
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const expected = {
|
||||
content: 'some content',
|
||||
name: 'dan',
|
||||
function_call: {
|
||||
name: 'dall-e',
|
||||
arguments: '{\n "input": "Subject: hedgehog, Style: cute"\n}',
|
||||
},
|
||||
};
|
||||
|
||||
expect(formatFromLangChain(message)).toEqual(expected);
|
||||
});
|
||||
|
||||
it('should handle messages without additional_kwargs', () => {
|
||||
const message = {
|
||||
kwargs: {
|
||||
content: 'some content',
|
||||
name: 'dan',
|
||||
},
|
||||
};
|
||||
|
||||
const expected = {
|
||||
content: 'some content',
|
||||
name: 'dan',
|
||||
};
|
||||
|
||||
expect(formatFromLangChain(message)).toEqual(expected);
|
||||
});
|
||||
|
||||
it('should handle empty messages', () => {
|
||||
const message = {
|
||||
kwargs: {},
|
||||
};
|
||||
|
||||
const expected = {};
|
||||
|
||||
expect(formatFromLangChain(message)).toEqual(expected);
|
||||
});
|
||||
});
|
||||
});
|
||||
13
api/app/clients/prompts/index.js
Normal file
13
api/app/clients/prompts/index.js
Normal file
@ -0,0 +1,13 @@
|
||||
const formatMessages = require('./formatMessages');
|
||||
const summaryPrompts = require('./summaryPrompts');
|
||||
const truncate = require('./truncate');
|
||||
const createVisionPrompt = require('./createVisionPrompt');
|
||||
const createContextHandlers = require('./createContextHandlers');
|
||||
|
||||
module.exports = {
|
||||
...formatMessages,
|
||||
...summaryPrompts,
|
||||
...truncate,
|
||||
createVisionPrompt,
|
||||
createContextHandlers,
|
||||
};
|
||||
495
api/app/clients/prompts/shadcn-docs/components.js
Normal file
495
api/app/clients/prompts/shadcn-docs/components.js
Normal file
@ -0,0 +1,495 @@
|
||||
// Essential Components
|
||||
const essentialComponents = {
|
||||
avatar: {
|
||||
componentName: 'Avatar',
|
||||
importDocs: 'import { Avatar, AvatarFallback, AvatarImage } from "/components/ui/avatar"',
|
||||
usageDocs: `
|
||||
<Avatar>
|
||||
<AvatarImage src="https://github.com/shadcn.png" />
|
||||
<AvatarFallback>CN</AvatarFallback>
|
||||
</Avatar>`,
|
||||
},
|
||||
button: {
|
||||
componentName: 'Button',
|
||||
importDocs: 'import { Button } from "/components/ui/button"',
|
||||
usageDocs: `
|
||||
<Button variant="outline">Button</Button>`,
|
||||
},
|
||||
card: {
|
||||
componentName: 'Card',
|
||||
importDocs: `
|
||||
import {
|
||||
Card,
|
||||
CardContent,
|
||||
CardDescription,
|
||||
CardFooter,
|
||||
CardHeader,
|
||||
CardTitle,
|
||||
} from "/components/ui/card"`,
|
||||
usageDocs: `
|
||||
<Card>
|
||||
<CardHeader>
|
||||
<CardTitle>Card Title</CardTitle>
|
||||
<CardDescription>Card Description</CardDescription>
|
||||
</CardHeader>
|
||||
<CardContent>
|
||||
<p>Card Content</p>
|
||||
</CardContent>
|
||||
<CardFooter>
|
||||
<p>Card Footer</p>
|
||||
</CardFooter>
|
||||
</Card>`,
|
||||
},
|
||||
checkbox: {
|
||||
componentName: 'Checkbox',
|
||||
importDocs: 'import { Checkbox } from "/components/ui/checkbox"',
|
||||
usageDocs: '<Checkbox />',
|
||||
},
|
||||
input: {
|
||||
componentName: 'Input',
|
||||
importDocs: 'import { Input } from "/components/ui/input"',
|
||||
usageDocs: '<Input />',
|
||||
},
|
||||
label: {
|
||||
componentName: 'Label',
|
||||
importDocs: 'import { Label } from "/components/ui/label"',
|
||||
usageDocs: '<Label htmlFor="email">Your email address</Label>',
|
||||
},
|
||||
radioGroup: {
|
||||
componentName: 'RadioGroup',
|
||||
importDocs: `
|
||||
import { Label } from "/components/ui/label"
|
||||
import { RadioGroup, RadioGroupItem } from "/components/ui/radio-group"`,
|
||||
usageDocs: `
|
||||
<RadioGroup defaultValue="option-one">
|
||||
<div className="flex items-center space-x-2">
|
||||
<RadioGroupItem value="option-one" id="option-one" />
|
||||
<Label htmlFor="option-one">Option One</Label>
|
||||
</div>
|
||||
<div className="flex items-center space-x-2">
|
||||
<RadioGroupItem value="option-two" id="option-two" />
|
||||
<Label htmlFor="option-two">Option Two</Label>
|
||||
</div>
|
||||
</RadioGroup>`,
|
||||
},
|
||||
select: {
|
||||
componentName: 'Select',
|
||||
importDocs: `
|
||||
import {
|
||||
Select,
|
||||
SelectContent,
|
||||
SelectItem,
|
||||
SelectTrigger,
|
||||
SelectValue,
|
||||
} from "/components/ui/select"`,
|
||||
usageDocs: `
|
||||
<Select>
|
||||
<SelectTrigger className="w-[180px]">
|
||||
<SelectValue placeholder="Theme" />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="light">Light</SelectItem>
|
||||
<SelectItem value="dark">Dark</SelectItem>
|
||||
<SelectItem value="system">System</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>`,
|
||||
},
|
||||
textarea: {
|
||||
componentName: 'Textarea',
|
||||
importDocs: 'import { Textarea } from "/components/ui/textarea"',
|
||||
usageDocs: '<Textarea />',
|
||||
},
|
||||
};
|
||||
|
||||
// Extra Components
|
||||
const extraComponents = {
|
||||
accordion: {
|
||||
componentName: 'Accordion',
|
||||
importDocs: `
|
||||
import {
|
||||
Accordion,
|
||||
AccordionContent,
|
||||
AccordionItem,
|
||||
AccordionTrigger,
|
||||
} from "/components/ui/accordion"`,
|
||||
usageDocs: `
|
||||
<Accordion type="single" collapsible>
|
||||
<AccordionItem value="item-1">
|
||||
<AccordionTrigger>Is it accessible?</AccordionTrigger>
|
||||
<AccordionContent>
|
||||
Yes. It adheres to the WAI-ARIA design pattern.
|
||||
</AccordionContent>
|
||||
</AccordionItem>
|
||||
</Accordion>`,
|
||||
},
|
||||
alertDialog: {
|
||||
componentName: 'AlertDialog',
|
||||
importDocs: `
|
||||
import {
|
||||
AlertDialog,
|
||||
AlertDialogAction,
|
||||
AlertDialogCancel,
|
||||
AlertDialogContent,
|
||||
AlertDialogDescription,
|
||||
AlertDialogFooter,
|
||||
AlertDialogHeader,
|
||||
AlertDialogTitle,
|
||||
AlertDialogTrigger,
|
||||
} from "/components/ui/alert-dialog"`,
|
||||
usageDocs: `
|
||||
<AlertDialog>
|
||||
<AlertDialogTrigger>Open</AlertDialogTrigger>
|
||||
<AlertDialogContent>
|
||||
<AlertDialogHeader>
|
||||
<AlertDialogTitle>Are you absolutely sure?</AlertDialogTitle>
|
||||
<AlertDialogDescription>
|
||||
This action cannot be undone.
|
||||
</AlertDialogDescription>
|
||||
</AlertDialogHeader>
|
||||
<AlertDialogFooter>
|
||||
<AlertDialogCancel>Cancel</AlertDialogCancel>
|
||||
<AlertDialogAction>Continue</AlertDialogAction>
|
||||
</AlertDialogFooter>
|
||||
</AlertDialogContent>
|
||||
</AlertDialog>`,
|
||||
},
|
||||
alert: {
|
||||
componentName: 'Alert',
|
||||
importDocs: `
|
||||
import {
|
||||
Alert,
|
||||
AlertDescription,
|
||||
AlertTitle,
|
||||
} from "/components/ui/alert"`,
|
||||
usageDocs: `
|
||||
<Alert>
|
||||
<AlertTitle>Heads up!</AlertTitle>
|
||||
<AlertDescription>
|
||||
You can add components to your app using the cli.
|
||||
</AlertDescription>
|
||||
</Alert>`,
|
||||
},
|
||||
aspectRatio: {
|
||||
componentName: 'AspectRatio',
|
||||
importDocs: 'import { AspectRatio } from "/components/ui/aspect-ratio"',
|
||||
usageDocs: `
|
||||
<AspectRatio ratio={16 / 9}>
|
||||
<Image src="..." alt="Image" className="rounded-md object-cover" />
|
||||
</AspectRatio>`,
|
||||
},
|
||||
badge: {
|
||||
componentName: 'Badge',
|
||||
importDocs: 'import { Badge } from "/components/ui/badge"',
|
||||
usageDocs: '<Badge>Badge</Badge>',
|
||||
},
|
||||
calendar: {
|
||||
componentName: 'Calendar',
|
||||
importDocs: 'import { Calendar } from "/components/ui/calendar"',
|
||||
usageDocs: '<Calendar />',
|
||||
},
|
||||
carousel: {
|
||||
componentName: 'Carousel',
|
||||
importDocs: `
|
||||
import {
|
||||
Carousel,
|
||||
CarouselContent,
|
||||
CarouselItem,
|
||||
CarouselNext,
|
||||
CarouselPrevious,
|
||||
} from "/components/ui/carousel"`,
|
||||
usageDocs: `
|
||||
<Carousel>
|
||||
<CarouselContent>
|
||||
<CarouselItem>...</CarouselItem>
|
||||
<CarouselItem>...</CarouselItem>
|
||||
<CarouselItem>...</CarouselItem>
|
||||
</CarouselContent>
|
||||
<CarouselPrevious />
|
||||
<CarouselNext />
|
||||
</Carousel>`,
|
||||
},
|
||||
collapsible: {
|
||||
componentName: 'Collapsible',
|
||||
importDocs: `
|
||||
import {
|
||||
Collapsible,
|
||||
CollapsibleContent,
|
||||
CollapsibleTrigger,
|
||||
} from "/components/ui/collapsible"`,
|
||||
usageDocs: `
|
||||
<Collapsible>
|
||||
<CollapsibleTrigger>Can I use this in my project?</CollapsibleTrigger>
|
||||
<CollapsibleContent>
|
||||
Yes. Free to use for personal and commercial projects. No attribution required.
|
||||
</CollapsibleContent>
|
||||
</Collapsible>`,
|
||||
},
|
||||
dialog: {
|
||||
componentName: 'Dialog',
|
||||
importDocs: `
|
||||
import {
|
||||
Dialog,
|
||||
DialogContent,
|
||||
DialogDescription,
|
||||
DialogHeader,
|
||||
DialogTitle,
|
||||
DialogTrigger,
|
||||
} from "/components/ui/dialog"`,
|
||||
usageDocs: `
|
||||
<Dialog>
|
||||
<DialogTrigger>Open</DialogTrigger>
|
||||
<DialogContent>
|
||||
<DialogHeader>
|
||||
<DialogTitle>Are you sure absolutely sure?</DialogTitle>
|
||||
<DialogDescription>
|
||||
This action cannot be undone.
|
||||
</DialogDescription>
|
||||
</DialogHeader>
|
||||
</DialogContent>
|
||||
</Dialog>`,
|
||||
},
|
||||
dropdownMenu: {
|
||||
componentName: 'DropdownMenu',
|
||||
importDocs: `
|
||||
import {
|
||||
DropdownMenu,
|
||||
DropdownMenuContent,
|
||||
DropdownMenuItem,
|
||||
DropdownMenuLabel,
|
||||
DropdownMenuSeparator,
|
||||
DropdownMenuTrigger,
|
||||
} from "/components/ui/dropdown-menu"`,
|
||||
usageDocs: `
|
||||
<DropdownMenu>
|
||||
<DropdownMenuTrigger>Open</DropdownMenuTrigger>
|
||||
<DropdownMenuContent>
|
||||
<DropdownMenuLabel>My Account</DropdownMenuLabel>
|
||||
<DropdownMenuSeparator />
|
||||
<DropdownMenuItem>Profile</DropdownMenuItem>
|
||||
<DropdownMenuItem>Billing</DropdownMenuItem>
|
||||
<DropdownMenuItem>Team</DropdownMenuItem>
|
||||
<DropdownMenuItem>Subscription</DropdownMenuItem>
|
||||
</DropdownMenuContent>
|
||||
</DropdownMenu>`,
|
||||
},
|
||||
menubar: {
|
||||
componentName: 'Menubar',
|
||||
importDocs: `
|
||||
import {
|
||||
Menubar,
|
||||
MenubarContent,
|
||||
MenubarItem,
|
||||
MenubarMenu,
|
||||
MenubarSeparator,
|
||||
MenubarShortcut,
|
||||
MenubarTrigger,
|
||||
} from "/components/ui/menubar"`,
|
||||
usageDocs: `
|
||||
<Menubar>
|
||||
<MenubarMenu>
|
||||
<MenubarTrigger>File</MenubarTrigger>
|
||||
<MenubarContent>
|
||||
<MenubarItem>
|
||||
New Tab <MenubarShortcut>⌘T</MenubarShortcut>
|
||||
</MenubarItem>
|
||||
<MenubarItem>New Window</MenubarItem>
|
||||
<MenubarSeparator />
|
||||
<MenubarItem>Share</MenubarItem>
|
||||
<MenubarSeparator />
|
||||
<MenubarItem>Print</MenubarItem>
|
||||
</MenubarContent>
|
||||
</MenubarMenu>
|
||||
</Menubar>`,
|
||||
},
|
||||
navigationMenu: {
|
||||
componentName: 'NavigationMenu',
|
||||
importDocs: `
|
||||
import {
|
||||
NavigationMenu,
|
||||
NavigationMenuContent,
|
||||
NavigationMenuItem,
|
||||
NavigationMenuLink,
|
||||
NavigationMenuList,
|
||||
NavigationMenuTrigger,
|
||||
navigationMenuTriggerStyle,
|
||||
} from "/components/ui/navigation-menu"`,
|
||||
usageDocs: `
|
||||
<NavigationMenu>
|
||||
<NavigationMenuList>
|
||||
<NavigationMenuItem>
|
||||
<NavigationMenuTrigger>Item One</NavigationMenuTrigger>
|
||||
<NavigationMenuContent>
|
||||
<NavigationMenuLink>Link</NavigationMenuLink>
|
||||
</NavigationMenuContent>
|
||||
</NavigationMenuItem>
|
||||
</NavigationMenuList>
|
||||
</NavigationMenu>`,
|
||||
},
|
||||
popover: {
|
||||
componentName: 'Popover',
|
||||
importDocs: `
|
||||
import {
|
||||
Popover,
|
||||
PopoverContent,
|
||||
PopoverTrigger,
|
||||
} from "/components/ui/popover"`,
|
||||
usageDocs: `
|
||||
<Popover>
|
||||
<PopoverTrigger>Open</PopoverTrigger>
|
||||
<PopoverContent>Place content for the popover here.</PopoverContent>
|
||||
</Popover>`,
|
||||
},
|
||||
progress: {
|
||||
componentName: 'Progress',
|
||||
importDocs: 'import { Progress } from "/components/ui/progress"',
|
||||
usageDocs: '<Progress value={33} />',
|
||||
},
|
||||
separator: {
|
||||
componentName: 'Separator',
|
||||
importDocs: 'import { Separator } from "/components/ui/separator"',
|
||||
usageDocs: '<Separator />',
|
||||
},
|
||||
sheet: {
|
||||
componentName: 'Sheet',
|
||||
importDocs: `
|
||||
import {
|
||||
Sheet,
|
||||
SheetContent,
|
||||
SheetDescription,
|
||||
SheetHeader,
|
||||
SheetTitle,
|
||||
SheetTrigger,
|
||||
} from "/components/ui/sheet"`,
|
||||
usageDocs: `
|
||||
<Sheet>
|
||||
<SheetTrigger>Open</SheetTrigger>
|
||||
<SheetContent>
|
||||
<SheetHeader>
|
||||
<SheetTitle>Are you sure absolutely sure?</SheetTitle>
|
||||
<SheetDescription>
|
||||
This action cannot be undone.
|
||||
</SheetDescription>
|
||||
</SheetHeader>
|
||||
</SheetContent>
|
||||
</Sheet>`,
|
||||
},
|
||||
skeleton: {
|
||||
componentName: 'Skeleton',
|
||||
importDocs: 'import { Skeleton } from "/components/ui/skeleton"',
|
||||
usageDocs: '<Skeleton className="w-[100px] h-[20px] rounded-full" />',
|
||||
},
|
||||
slider: {
|
||||
componentName: 'Slider',
|
||||
importDocs: 'import { Slider } from "/components/ui/slider"',
|
||||
usageDocs: '<Slider defaultValue={[33]} max={100} step={1} />',
|
||||
},
|
||||
switch: {
|
||||
componentName: 'Switch',
|
||||
importDocs: 'import { Switch } from "/components/ui/switch"',
|
||||
usageDocs: '<Switch />',
|
||||
},
|
||||
table: {
|
||||
componentName: 'Table',
|
||||
importDocs: `
|
||||
import {
|
||||
Table,
|
||||
TableBody,
|
||||
TableCaption,
|
||||
TableCell,
|
||||
TableHead,
|
||||
TableHeader,
|
||||
TableRow,
|
||||
} from "/components/ui/table"`,
|
||||
usageDocs: `
|
||||
<Table>
|
||||
<TableCaption>A list of your recent invoices.</TableCaption>
|
||||
<TableHeader>
|
||||
<TableRow>
|
||||
<TableHead className="w-[100px]">Invoice</TableHead>
|
||||
<TableHead>Status</TableHead>
|
||||
<TableHead>Method</TableHead>
|
||||
<TableHead className="text-right">Amount</TableHead>
|
||||
</TableRow>
|
||||
</TableHeader>
|
||||
<TableBody>
|
||||
<TableRow>
|
||||
<TableCell className="font-medium">INV001</TableCell>
|
||||
<TableCell>Paid</TableCell>
|
||||
<TableCell>Credit Card</TableCell>
|
||||
<TableCell className="text-right">$250.00</TableCell>
|
||||
</TableRow>
|
||||
</TableBody>
|
||||
</Table>`,
|
||||
},
|
||||
tabs: {
|
||||
componentName: 'Tabs',
|
||||
importDocs: `
|
||||
import {
|
||||
Tabs,
|
||||
TabsContent,
|
||||
TabsList,
|
||||
TabsTrigger,
|
||||
} from "/components/ui/tabs"`,
|
||||
usageDocs: `
|
||||
<Tabs defaultValue="account" className="w-[400px]">
|
||||
<TabsList>
|
||||
<TabsTrigger value="account">Account</TabsTrigger>
|
||||
<TabsTrigger value="password">Password</TabsTrigger>
|
||||
</TabsList>
|
||||
<TabsContent value="account">Make changes to your account here.</TabsContent>
|
||||
<TabsContent value="password">Change your password here.</TabsContent>
|
||||
</Tabs>`,
|
||||
},
|
||||
toast: {
|
||||
componentName: 'Toast',
|
||||
importDocs: `
|
||||
import { useToast } from "/components/ui/use-toast"
|
||||
import { Button } from "/components/ui/button"`,
|
||||
usageDocs: `
|
||||
export function ToastDemo() {
|
||||
const { toast } = useToast()
|
||||
return (
|
||||
<Button
|
||||
onClick={() => {
|
||||
toast({
|
||||
title: "Scheduled: Catch up",
|
||||
description: "Friday, February 10, 2023 at 5:57 PM",
|
||||
})
|
||||
}}
|
||||
>
|
||||
Show Toast
|
||||
</Button>
|
||||
)
|
||||
}`,
|
||||
},
|
||||
toggle: {
|
||||
componentName: 'Toggle',
|
||||
importDocs: 'import { Toggle } from "/components/ui/toggle"',
|
||||
usageDocs: '<Toggle>Toggle</Toggle>',
|
||||
},
|
||||
tooltip: {
|
||||
componentName: 'Tooltip',
|
||||
importDocs: `
|
||||
import {
|
||||
Tooltip,
|
||||
TooltipContent,
|
||||
TooltipProvider,
|
||||
TooltipTrigger,
|
||||
} from "/components/ui/tooltip"`,
|
||||
usageDocs: `
|
||||
<TooltipProvider>
|
||||
<Tooltip>
|
||||
<TooltipTrigger>Hover</TooltipTrigger>
|
||||
<TooltipContent>
|
||||
<p>Add to library</p>
|
||||
</TooltipContent>
|
||||
</Tooltip>
|
||||
</TooltipProvider>`,
|
||||
},
|
||||
};
|
||||
|
||||
const components = Object.assign({}, essentialComponents, extraComponents);
|
||||
|
||||
module.exports = {
|
||||
components,
|
||||
};
|
||||
50
api/app/clients/prompts/shadcn-docs/generate.js
Normal file
50
api/app/clients/prompts/shadcn-docs/generate.js
Normal file
@ -0,0 +1,50 @@
|
||||
const dedent = require('dedent');
|
||||
|
||||
/**
|
||||
* Generate system prompt for AI-assisted React component creation
|
||||
* @param {Object} options - Configuration options
|
||||
* @param {Object} options.components - Documentation for shadcn components
|
||||
* @param {boolean} [options.useXML=false] - Whether to use XML-style formatting for component instructions
|
||||
* @returns {string} The generated system prompt
|
||||
*/
|
||||
function generateShadcnPrompt(options) {
|
||||
const { components, useXML = false } = options;
|
||||
|
||||
let systemPrompt = dedent`
|
||||
## Additional Artifact Instructions for React Components: "application/vnd.react"
|
||||
|
||||
There are some prestyled components (primitives) available for use. Please use your best judgement to use any of these components if the app calls for one.
|
||||
|
||||
Here are the components that are available, along with how to import them, and how to use them:
|
||||
|
||||
${Object.values(components)
|
||||
.map((component) => {
|
||||
if (useXML) {
|
||||
return dedent`
|
||||
<component>
|
||||
<name>${component.componentName}</name>
|
||||
<import-instructions>${component.importDocs}</import-instructions>
|
||||
<usage-instructions>${component.usageDocs}</usage-instructions>
|
||||
</component>
|
||||
`;
|
||||
} else {
|
||||
return dedent`
|
||||
# ${component.componentName}
|
||||
|
||||
## Import Instructions
|
||||
${component.importDocs}
|
||||
|
||||
## Usage Instructions
|
||||
${component.usageDocs}
|
||||
`;
|
||||
}
|
||||
})
|
||||
.join('\n\n')}
|
||||
`;
|
||||
|
||||
return systemPrompt;
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
generateShadcnPrompt,
|
||||
};
|
||||
53
api/app/clients/prompts/summaryPrompts.js
Normal file
53
api/app/clients/prompts/summaryPrompts.js
Normal file
@ -0,0 +1,53 @@
|
||||
const { PromptTemplate } = require('@langchain/core/prompts');
|
||||
/*
|
||||
* Without `{summary}` and `{new_lines}`, token count is 98
|
||||
* We are counting this towards the max context tokens for summaries, +3 for the assistant label (101)
|
||||
* If this prompt changes, use https://tiktokenizer.vercel.app/ to count the tokens
|
||||
*/
|
||||
const _DEFAULT_SUMMARIZER_TEMPLATE = `Summarize the conversation by integrating new lines into the current summary.
|
||||
|
||||
EXAMPLE:
|
||||
Current summary:
|
||||
The human inquires about the AI's view on artificial intelligence. The AI believes it's beneficial.
|
||||
|
||||
New lines:
|
||||
Human: Why is it beneficial?
|
||||
AI: It helps humans achieve their potential.
|
||||
|
||||
New summary:
|
||||
The human inquires about the AI's view on artificial intelligence. The AI believes it's beneficial because it helps humans achieve their potential.
|
||||
|
||||
Current summary:
|
||||
{summary}
|
||||
|
||||
New lines:
|
||||
{new_lines}
|
||||
|
||||
New summary:`;
|
||||
|
||||
const SUMMARY_PROMPT = new PromptTemplate({
|
||||
inputVariables: ['summary', 'new_lines'],
|
||||
template: _DEFAULT_SUMMARIZER_TEMPLATE,
|
||||
});
|
||||
|
||||
/*
|
||||
* Without `{new_lines}`, token count is 27
|
||||
* We are counting this towards the max context tokens for summaries, rounded up to 30
|
||||
* If this prompt changes, use https://tiktokenizer.vercel.app/ to count the tokens
|
||||
*/
|
||||
const _CUT_OFF_SUMMARIZER = `The following text is cut-off:
|
||||
{new_lines}
|
||||
|
||||
Summarize the content as best as you can, noting that it was cut-off.
|
||||
|
||||
Summary:`;
|
||||
|
||||
const CUT_OFF_PROMPT = new PromptTemplate({
|
||||
inputVariables: ['new_lines'],
|
||||
template: _CUT_OFF_SUMMARIZER,
|
||||
});
|
||||
|
||||
module.exports = {
|
||||
SUMMARY_PROMPT,
|
||||
CUT_OFF_PROMPT,
|
||||
};
|
||||
115
api/app/clients/prompts/truncate.js
Normal file
115
api/app/clients/prompts/truncate.js
Normal file
@ -0,0 +1,115 @@
|
||||
const MAX_CHAR = 255;
|
||||
|
||||
/**
|
||||
* Truncates a given text to a specified maximum length, appending ellipsis and a notification
|
||||
* if the original text exceeds the maximum length.
|
||||
*
|
||||
* @param {string} text - The text to be truncated.
|
||||
* @param {number} [maxLength=MAX_CHAR] - The maximum length of the text after truncation. Defaults to MAX_CHAR.
|
||||
* @returns {string} The truncated text if the original text length exceeds maxLength, otherwise returns the original text.
|
||||
*/
|
||||
function truncateText(text, maxLength = MAX_CHAR) {
|
||||
if (text.length > maxLength) {
|
||||
return `${text.slice(0, maxLength)}... [text truncated for brevity]`;
|
||||
}
|
||||
return text;
|
||||
}
|
||||
|
||||
/**
|
||||
* Truncates a given text to a specified maximum length by showing the first half and the last half of the text,
|
||||
* separated by ellipsis. This method ensures the output does not exceed the maximum length, including the addition
|
||||
* of ellipsis and notification if the original text exceeds the maximum length.
|
||||
*
|
||||
* @param {string} text - The text to be truncated.
|
||||
* @param {number} [maxLength=MAX_CHAR] - The maximum length of the output text after truncation. Defaults to MAX_CHAR.
|
||||
* @returns {string} The truncated text showing the first half and the last half, or the original text if it does not exceed maxLength.
|
||||
*/
|
||||
function smartTruncateText(text, maxLength = MAX_CHAR) {
|
||||
const ellipsis = '...';
|
||||
const notification = ' [text truncated for brevity]';
|
||||
const halfMaxLength = Math.floor((maxLength - ellipsis.length - notification.length) / 2);
|
||||
|
||||
if (text.length > maxLength) {
|
||||
const startLastHalf = text.length - halfMaxLength;
|
||||
return `${text.slice(0, halfMaxLength)}${ellipsis}${text.slice(startLastHalf)}${notification}`;
|
||||
}
|
||||
|
||||
return text;
|
||||
}
|
||||
|
||||
/**
|
||||
* @param {TMessage[]} _messages
|
||||
* @param {number} maxContextTokens
|
||||
* @param {function({role: string, content: TMessageContent[]}): number} getTokenCountForMessage
|
||||
*
|
||||
* @returns {{
|
||||
* dbMessages: TMessage[],
|
||||
* editedIndices: number[]
|
||||
* }}
|
||||
*/
|
||||
function truncateToolCallOutputs(_messages, maxContextTokens, getTokenCountForMessage) {
|
||||
const THRESHOLD_PERCENTAGE = 0.5;
|
||||
const targetTokenLimit = maxContextTokens * THRESHOLD_PERCENTAGE;
|
||||
|
||||
let currentTokenCount = 3;
|
||||
const messages = [..._messages];
|
||||
const processedMessages = [];
|
||||
let currentIndex = messages.length;
|
||||
const editedIndices = new Set();
|
||||
while (messages.length > 0) {
|
||||
currentIndex--;
|
||||
const message = messages.pop();
|
||||
currentTokenCount += message.tokenCount;
|
||||
if (currentTokenCount < targetTokenLimit) {
|
||||
processedMessages.push(message);
|
||||
continue;
|
||||
}
|
||||
|
||||
if (!message.content || !Array.isArray(message.content)) {
|
||||
processedMessages.push(message);
|
||||
continue;
|
||||
}
|
||||
|
||||
const toolCallIndices = message.content
|
||||
.map((item, index) => (item.type === 'tool_call' ? index : -1))
|
||||
.filter((index) => index !== -1)
|
||||
.reverse();
|
||||
|
||||
if (toolCallIndices.length === 0) {
|
||||
processedMessages.push(message);
|
||||
continue;
|
||||
}
|
||||
|
||||
const newContent = [...message.content];
|
||||
|
||||
// Truncate all tool outputs since we're over threshold
|
||||
for (const index of toolCallIndices) {
|
||||
const toolCall = newContent[index].tool_call;
|
||||
if (!toolCall || !toolCall.output) {
|
||||
continue;
|
||||
}
|
||||
|
||||
editedIndices.add(currentIndex);
|
||||
|
||||
newContent[index] = {
|
||||
...newContent[index],
|
||||
tool_call: {
|
||||
...toolCall,
|
||||
output: '[OUTPUT_OMITTED_FOR_BREVITY]',
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
const truncatedMessage = {
|
||||
...message,
|
||||
content: newContent,
|
||||
tokenCount: getTokenCountForMessage({ role: 'assistant', content: newContent }),
|
||||
};
|
||||
|
||||
processedMessages.push(truncatedMessage);
|
||||
}
|
||||
|
||||
return { dbMessages: processedMessages.reverse(), editedIndices: Array.from(editedIndices) };
|
||||
}
|
||||
|
||||
module.exports = { truncateText, smartTruncateText, truncateToolCallOutputs };
|
||||
931
api/app/clients/specs/BaseClient.test.js
Normal file
931
api/app/clients/specs/BaseClient.test.js
Normal file
@ -0,0 +1,931 @@
|
||||
const { Constants } = require('librechat-data-provider');
|
||||
const { initializeFakeClient } = require('./FakeClient');
|
||||
|
||||
jest.mock('~/db/connect');
|
||||
jest.mock('~/server/services/Config', () => ({
|
||||
getAppConfig: jest.fn().mockResolvedValue({
|
||||
// Default app config for tests
|
||||
paths: { uploads: '/tmp' },
|
||||
fileStrategy: 'local',
|
||||
memory: { disabled: false },
|
||||
}),
|
||||
}));
|
||||
jest.mock('~/models', () => ({
|
||||
User: jest.fn(),
|
||||
Key: jest.fn(),
|
||||
Session: jest.fn(),
|
||||
Balance: jest.fn(),
|
||||
Transaction: jest.fn(),
|
||||
getMessages: jest.fn().mockResolvedValue([]),
|
||||
saveMessage: jest.fn(),
|
||||
updateMessage: jest.fn(),
|
||||
deleteMessagesSince: jest.fn(),
|
||||
deleteMessages: jest.fn(),
|
||||
getConvoTitle: jest.fn(),
|
||||
getConvo: jest.fn(),
|
||||
saveConvo: jest.fn(),
|
||||
deleteConvos: jest.fn(),
|
||||
getPreset: jest.fn(),
|
||||
getPresets: jest.fn(),
|
||||
savePreset: jest.fn(),
|
||||
deletePresets: jest.fn(),
|
||||
findFileById: jest.fn(),
|
||||
createFile: jest.fn(),
|
||||
updateFile: jest.fn(),
|
||||
deleteFile: jest.fn(),
|
||||
deleteFiles: jest.fn(),
|
||||
getFiles: jest.fn(),
|
||||
updateFileUsage: jest.fn(),
|
||||
}));
|
||||
|
||||
const { getConvo, saveConvo } = require('~/models');
|
||||
|
||||
jest.mock('@librechat/agents', () => {
|
||||
const { Providers } = jest.requireActual('@librechat/agents');
|
||||
return {
|
||||
Providers,
|
||||
ChatOpenAI: jest.fn().mockImplementation(() => {
|
||||
return {};
|
||||
}),
|
||||
};
|
||||
});
|
||||
|
||||
let parentMessageId;
|
||||
let conversationId;
|
||||
const fakeMessages = [];
|
||||
const userMessage = 'Hello, ChatGPT!';
|
||||
const apiKey = 'fake-api-key';
|
||||
|
||||
const messageHistory = [
|
||||
{ role: 'user', isCreatedByUser: true, text: 'Hello', messageId: '1' },
|
||||
{ role: 'assistant', isCreatedByUser: false, text: 'Hi', messageId: '2', parentMessageId: '1' },
|
||||
{
|
||||
role: 'user',
|
||||
isCreatedByUser: true,
|
||||
text: "What's up",
|
||||
messageId: '3',
|
||||
parentMessageId: '2',
|
||||
},
|
||||
];
|
||||
|
||||
describe('BaseClient', () => {
|
||||
let TestClient;
|
||||
const options = {
|
||||
// debug: true,
|
||||
modelOptions: {
|
||||
model: 'gpt-4o-mini',
|
||||
temperature: 0,
|
||||
},
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
TestClient = initializeFakeClient(apiKey, options, fakeMessages);
|
||||
TestClient.summarizeMessages = jest.fn().mockResolvedValue({
|
||||
summaryMessage: {
|
||||
role: 'system',
|
||||
content: 'Refined answer',
|
||||
},
|
||||
summaryTokenCount: 5,
|
||||
});
|
||||
});
|
||||
|
||||
test('returns the input messages without instructions when addInstructions() is called with empty instructions', () => {
|
||||
const messages = [{ content: 'Hello' }, { content: 'How are you?' }, { content: 'Goodbye' }];
|
||||
const instructions = '';
|
||||
const result = TestClient.addInstructions(messages, instructions);
|
||||
expect(result).toEqual(messages);
|
||||
});
|
||||
|
||||
test('returns the input messages with instructions properly added when addInstructions() is called with non-empty instructions', () => {
|
||||
const messages = [{ content: 'Hello' }, { content: 'How are you?' }, { content: 'Goodbye' }];
|
||||
const instructions = { content: 'Please respond to the question.' };
|
||||
const result = TestClient.addInstructions(messages, instructions);
|
||||
const expected = [
|
||||
{ content: 'Please respond to the question.' },
|
||||
{ content: 'Hello' },
|
||||
{ content: 'How are you?' },
|
||||
{ content: 'Goodbye' },
|
||||
];
|
||||
expect(result).toEqual(expected);
|
||||
});
|
||||
|
||||
test('returns the input messages with instructions properly added when addInstructions() with legacy flag', () => {
|
||||
const messages = [{ content: 'Hello' }, { content: 'How are you?' }, { content: 'Goodbye' }];
|
||||
const instructions = { content: 'Please respond to the question.' };
|
||||
const result = TestClient.addInstructions(messages, instructions, true);
|
||||
const expected = [
|
||||
{ content: 'Hello' },
|
||||
{ content: 'How are you?' },
|
||||
{ content: 'Please respond to the question.' },
|
||||
{ content: 'Goodbye' },
|
||||
];
|
||||
expect(result).toEqual(expected);
|
||||
});
|
||||
|
||||
test('concats messages correctly in concatenateMessages()', () => {
|
||||
const messages = [
|
||||
{ name: 'User', content: 'Hello' },
|
||||
{ name: 'Assistant', content: 'How can I help you?' },
|
||||
{ name: 'User', content: 'I have a question.' },
|
||||
];
|
||||
const result = TestClient.concatenateMessages(messages);
|
||||
const expected =
|
||||
'User:\nHello\n\nAssistant:\nHow can I help you?\n\nUser:\nI have a question.\n\n';
|
||||
expect(result).toBe(expected);
|
||||
});
|
||||
|
||||
test('refines messages correctly in summarizeMessages()', async () => {
|
||||
const messagesToRefine = [
|
||||
{ role: 'user', content: 'Hello', tokenCount: 10 },
|
||||
{ role: 'assistant', content: 'How can I help you?', tokenCount: 20 },
|
||||
];
|
||||
const remainingContextTokens = 100;
|
||||
const expectedRefinedMessage = {
|
||||
role: 'system',
|
||||
content: 'Refined answer',
|
||||
};
|
||||
|
||||
const result = await TestClient.summarizeMessages({ messagesToRefine, remainingContextTokens });
|
||||
expect(result.summaryMessage).toEqual(expectedRefinedMessage);
|
||||
});
|
||||
|
||||
test('gets messages within token limit (under limit) correctly in getMessagesWithinTokenLimit()', async () => {
|
||||
TestClient.maxContextTokens = 100;
|
||||
TestClient.shouldSummarize = true;
|
||||
|
||||
const messages = [
|
||||
{ role: 'user', content: 'Hello', tokenCount: 5 },
|
||||
{ role: 'assistant', content: 'How can I help you?', tokenCount: 19 },
|
||||
{ role: 'user', content: 'I have a question.', tokenCount: 18 },
|
||||
];
|
||||
const expectedContext = [
|
||||
{ role: 'user', content: 'Hello', tokenCount: 5 }, // 'Hello'.length
|
||||
{ role: 'assistant', content: 'How can I help you?', tokenCount: 19 },
|
||||
{ role: 'user', content: 'I have a question.', tokenCount: 18 },
|
||||
];
|
||||
// Subtract 3 tokens for Assistant Label priming after all messages have been counted.
|
||||
const expectedRemainingContextTokens = 58 - 3; // (100 - 5 - 19 - 18) - 3
|
||||
const expectedMessagesToRefine = [];
|
||||
|
||||
const lastExpectedMessage =
|
||||
expectedMessagesToRefine?.[expectedMessagesToRefine.length - 1] ?? {};
|
||||
const expectedIndex = messages.findIndex((msg) => msg.content === lastExpectedMessage?.content);
|
||||
|
||||
const result = await TestClient.getMessagesWithinTokenLimit({ messages });
|
||||
|
||||
expect(result.context).toEqual(expectedContext);
|
||||
expect(result.messagesToRefine.length - 1).toEqual(expectedIndex);
|
||||
expect(result.remainingContextTokens).toBe(expectedRemainingContextTokens);
|
||||
expect(result.messagesToRefine).toEqual(expectedMessagesToRefine);
|
||||
});
|
||||
|
||||
test('gets result over token limit correctly in getMessagesWithinTokenLimit()', async () => {
|
||||
TestClient.maxContextTokens = 50; // Set a lower limit
|
||||
TestClient.shouldSummarize = true;
|
||||
|
||||
const messages = [
|
||||
{ role: 'user', content: 'Hello', tokenCount: 30 },
|
||||
{ role: 'assistant', content: 'How can I help you?', tokenCount: 30 },
|
||||
{ role: 'user', content: 'I have a question.', tokenCount: 5 },
|
||||
{ role: 'user', content: 'I need a coffee, stat!', tokenCount: 19 },
|
||||
{ role: 'assistant', content: 'Sure, I can help with that.', tokenCount: 18 },
|
||||
];
|
||||
|
||||
// Subtract 3 tokens for Assistant Label priming after all messages have been counted.
|
||||
const expectedRemainingContextTokens = 5; // (50 - 18 - 19 - 5) - 3
|
||||
const expectedMessagesToRefine = [
|
||||
{ role: 'user', content: 'Hello', tokenCount: 30 },
|
||||
{ role: 'assistant', content: 'How can I help you?', tokenCount: 30 },
|
||||
];
|
||||
const expectedContext = [
|
||||
{ role: 'user', content: 'I have a question.', tokenCount: 5 },
|
||||
{ role: 'user', content: 'I need a coffee, stat!', tokenCount: 19 },
|
||||
{ role: 'assistant', content: 'Sure, I can help with that.', tokenCount: 18 },
|
||||
];
|
||||
|
||||
const lastExpectedMessage =
|
||||
expectedMessagesToRefine?.[expectedMessagesToRefine.length - 1] ?? {};
|
||||
const expectedIndex = messages.findIndex((msg) => msg.content === lastExpectedMessage?.content);
|
||||
|
||||
const result = await TestClient.getMessagesWithinTokenLimit({ messages });
|
||||
|
||||
expect(result.context).toEqual(expectedContext);
|
||||
expect(result.messagesToRefine.length - 1).toEqual(expectedIndex);
|
||||
expect(result.remainingContextTokens).toBe(expectedRemainingContextTokens);
|
||||
expect(result.messagesToRefine).toEqual(expectedMessagesToRefine);
|
||||
});
|
||||
|
||||
describe('getMessagesForConversation', () => {
|
||||
it('should return an empty array if the parentMessageId does not exist', () => {
|
||||
const result = TestClient.constructor.getMessagesForConversation({
|
||||
messages: unorderedMessages,
|
||||
parentMessageId: '999',
|
||||
});
|
||||
expect(result).toEqual([]);
|
||||
});
|
||||
|
||||
it('should handle messages with messageId property', () => {
|
||||
const messagesWithMessageId = [
|
||||
{ messageId: '1', parentMessageId: null, text: 'Message 1' },
|
||||
{ messageId: '2', parentMessageId: '1', text: 'Message 2' },
|
||||
];
|
||||
const result = TestClient.constructor.getMessagesForConversation({
|
||||
messages: messagesWithMessageId,
|
||||
parentMessageId: '2',
|
||||
});
|
||||
expect(result).toEqual([
|
||||
{ messageId: '1', parentMessageId: null, text: 'Message 1' },
|
||||
{ messageId: '2', parentMessageId: '1', text: 'Message 2' },
|
||||
]);
|
||||
});
|
||||
|
||||
const messagesWithNullParent = [
|
||||
{ id: '1', parentMessageId: null, text: 'Message 1' },
|
||||
{ id: '2', parentMessageId: null, text: 'Message 2' },
|
||||
];
|
||||
|
||||
it('should handle messages with null parentMessageId that are not root', () => {
|
||||
const result = TestClient.constructor.getMessagesForConversation({
|
||||
messages: messagesWithNullParent,
|
||||
parentMessageId: '2',
|
||||
});
|
||||
expect(result).toEqual([{ id: '2', parentMessageId: null, text: 'Message 2' }]);
|
||||
});
|
||||
|
||||
const cyclicMessages = [
|
||||
{ id: '3', parentMessageId: '2', text: 'Message 3' },
|
||||
{ id: '1', parentMessageId: '3', text: 'Message 1' },
|
||||
{ id: '2', parentMessageId: '1', text: 'Message 2' },
|
||||
];
|
||||
|
||||
it('should handle cyclic references without going into an infinite loop', () => {
|
||||
const result = TestClient.constructor.getMessagesForConversation({
|
||||
messages: cyclicMessages,
|
||||
parentMessageId: '3',
|
||||
});
|
||||
expect(result).toEqual([
|
||||
{ id: '1', parentMessageId: '3', text: 'Message 1' },
|
||||
{ id: '2', parentMessageId: '1', text: 'Message 2' },
|
||||
{ id: '3', parentMessageId: '2', text: 'Message 3' },
|
||||
]);
|
||||
});
|
||||
|
||||
const unorderedMessages = [
|
||||
{ id: '3', parentMessageId: '2', text: 'Message 3' },
|
||||
{ id: '2', parentMessageId: '1', text: 'Message 2' },
|
||||
{ id: '1', parentMessageId: Constants.NO_PARENT, text: 'Message 1' },
|
||||
];
|
||||
|
||||
it('should return ordered messages based on parentMessageId', () => {
|
||||
const result = TestClient.constructor.getMessagesForConversation({
|
||||
messages: unorderedMessages,
|
||||
parentMessageId: '3',
|
||||
});
|
||||
expect(result).toEqual([
|
||||
{ id: '1', parentMessageId: Constants.NO_PARENT, text: 'Message 1' },
|
||||
{ id: '2', parentMessageId: '1', text: 'Message 2' },
|
||||
{ id: '3', parentMessageId: '2', text: 'Message 3' },
|
||||
]);
|
||||
});
|
||||
|
||||
const unorderedBranchedMessages = [
|
||||
{ id: '4', parentMessageId: '2', text: 'Message 4', summary: 'Summary for Message 4' },
|
||||
{ id: '10', parentMessageId: '7', text: 'Message 10' },
|
||||
{ id: '1', parentMessageId: null, text: 'Message 1' },
|
||||
{ id: '6', parentMessageId: '5', text: 'Message 7' },
|
||||
{ id: '7', parentMessageId: '5', text: 'Message 7' },
|
||||
{ id: '2', parentMessageId: '1', text: 'Message 2' },
|
||||
{ id: '8', parentMessageId: '6', text: 'Message 8' },
|
||||
{ id: '5', parentMessageId: '3', text: 'Message 5' },
|
||||
{ id: '3', parentMessageId: '1', text: 'Message 3' },
|
||||
{ id: '6', parentMessageId: '4', text: 'Message 6' },
|
||||
{ id: '8', parentMessageId: '7', text: 'Message 9' },
|
||||
{ id: '9', parentMessageId: '7', text: 'Message 9' },
|
||||
{ id: '11', parentMessageId: '2', text: 'Message 11', summary: 'Summary for Message 11' },
|
||||
];
|
||||
|
||||
it('should return ordered messages from a branched array based on parentMessageId', () => {
|
||||
const result = TestClient.constructor.getMessagesForConversation({
|
||||
messages: unorderedBranchedMessages,
|
||||
parentMessageId: '10',
|
||||
summary: true,
|
||||
});
|
||||
expect(result).toEqual([
|
||||
{ id: '1', parentMessageId: null, text: 'Message 1' },
|
||||
{ id: '3', parentMessageId: '1', text: 'Message 3' },
|
||||
{ id: '5', parentMessageId: '3', text: 'Message 5' },
|
||||
{ id: '7', parentMessageId: '5', text: 'Message 7' },
|
||||
{ id: '10', parentMessageId: '7', text: 'Message 10' },
|
||||
]);
|
||||
});
|
||||
|
||||
it('should return an empty array if no messages are provided', () => {
|
||||
const result = TestClient.constructor.getMessagesForConversation({
|
||||
messages: [],
|
||||
parentMessageId: '3',
|
||||
});
|
||||
expect(result).toEqual([]);
|
||||
});
|
||||
|
||||
it('should map over the ordered messages if mapMethod is provided', () => {
|
||||
const mapMethod = (msg) => msg.text;
|
||||
const result = TestClient.constructor.getMessagesForConversation({
|
||||
messages: unorderedMessages,
|
||||
parentMessageId: '3',
|
||||
mapMethod,
|
||||
});
|
||||
expect(result).toEqual(['Message 1', 'Message 2', 'Message 3']);
|
||||
});
|
||||
|
||||
let unorderedMessagesWithSummary = [
|
||||
{ id: '4', parentMessageId: '3', text: 'Message 4' },
|
||||
{ id: '2', parentMessageId: '1', text: 'Message 2', summary: 'Summary for Message 2' },
|
||||
{ id: '3', parentMessageId: '2', text: 'Message 3', summary: 'Summary for Message 3' },
|
||||
{ id: '1', parentMessageId: null, text: 'Message 1' },
|
||||
];
|
||||
|
||||
it('should start with the message that has a summary property and continue until the specified parentMessageId', () => {
|
||||
const result = TestClient.constructor.getMessagesForConversation({
|
||||
messages: unorderedMessagesWithSummary,
|
||||
parentMessageId: '4',
|
||||
summary: true,
|
||||
});
|
||||
expect(result).toEqual([
|
||||
{
|
||||
id: '3',
|
||||
parentMessageId: '2',
|
||||
role: 'system',
|
||||
text: 'Summary for Message 3',
|
||||
summary: 'Summary for Message 3',
|
||||
},
|
||||
{ id: '4', parentMessageId: '3', text: 'Message 4' },
|
||||
]);
|
||||
});
|
||||
|
||||
it('should handle multiple summaries and return the branch from the latest to the parentMessageId', () => {
|
||||
unorderedMessagesWithSummary = [
|
||||
{ id: '5', parentMessageId: '4', text: 'Message 5' },
|
||||
{ id: '2', parentMessageId: '1', text: 'Message 2', summary: 'Summary for Message 2' },
|
||||
{ id: '3', parentMessageId: '2', text: 'Message 3', summary: 'Summary for Message 3' },
|
||||
{ id: '4', parentMessageId: '3', text: 'Message 4', summary: 'Summary for Message 4' },
|
||||
{ id: '1', parentMessageId: null, text: 'Message 1' },
|
||||
];
|
||||
const result = TestClient.constructor.getMessagesForConversation({
|
||||
messages: unorderedMessagesWithSummary,
|
||||
parentMessageId: '5',
|
||||
summary: true,
|
||||
});
|
||||
expect(result).toEqual([
|
||||
{
|
||||
id: '4',
|
||||
parentMessageId: '3',
|
||||
role: 'system',
|
||||
text: 'Summary for Message 4',
|
||||
summary: 'Summary for Message 4',
|
||||
},
|
||||
{ id: '5', parentMessageId: '4', text: 'Message 5' },
|
||||
]);
|
||||
});
|
||||
|
||||
it('should handle summary at root edge case and continue until the parentMessageId', () => {
|
||||
unorderedMessagesWithSummary = [
|
||||
{ id: '5', parentMessageId: '4', text: 'Message 5' },
|
||||
{ id: '1', parentMessageId: null, text: 'Message 1', summary: 'Summary for Message 1' },
|
||||
{ id: '4', parentMessageId: '3', text: 'Message 4', summary: 'Summary for Message 4' },
|
||||
{ id: '2', parentMessageId: '1', text: 'Message 2', summary: 'Summary for Message 2' },
|
||||
{ id: '3', parentMessageId: '2', text: 'Message 3', summary: 'Summary for Message 3' },
|
||||
];
|
||||
const result = TestClient.constructor.getMessagesForConversation({
|
||||
messages: unorderedMessagesWithSummary,
|
||||
parentMessageId: '5',
|
||||
summary: true,
|
||||
});
|
||||
expect(result).toEqual([
|
||||
{
|
||||
id: '4',
|
||||
parentMessageId: '3',
|
||||
role: 'system',
|
||||
text: 'Summary for Message 4',
|
||||
summary: 'Summary for Message 4',
|
||||
},
|
||||
{ id: '5', parentMessageId: '4', text: 'Message 5' },
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('sendMessage', () => {
|
||||
test('sendMessage should return a response message', async () => {
|
||||
const expectedResult = expect.objectContaining({
|
||||
sender: TestClient.sender,
|
||||
text: expect.any(String),
|
||||
isCreatedByUser: false,
|
||||
messageId: expect.any(String),
|
||||
parentMessageId: expect.any(String),
|
||||
conversationId: expect.any(String),
|
||||
});
|
||||
|
||||
const response = await TestClient.sendMessage(userMessage);
|
||||
parentMessageId = response.messageId;
|
||||
conversationId = response.conversationId;
|
||||
expect(response).toEqual(expectedResult);
|
||||
});
|
||||
|
||||
test('should replace responseMessageId with new UUID when isRegenerate is true and messageId ends with underscore', async () => {
|
||||
const mockCrypto = require('crypto');
|
||||
const newUUID = 'new-uuid-1234';
|
||||
jest.spyOn(mockCrypto, 'randomUUID').mockReturnValue(newUUID);
|
||||
|
||||
const opts = {
|
||||
isRegenerate: true,
|
||||
responseMessageId: 'existing-message-id_',
|
||||
};
|
||||
|
||||
await TestClient.setMessageOptions(opts);
|
||||
|
||||
expect(TestClient.responseMessageId).toBe(newUUID);
|
||||
expect(TestClient.responseMessageId).not.toBe('existing-message-id_');
|
||||
|
||||
mockCrypto.randomUUID.mockRestore();
|
||||
});
|
||||
|
||||
test('should not replace responseMessageId when isRegenerate is false', async () => {
|
||||
const opts = {
|
||||
isRegenerate: false,
|
||||
responseMessageId: 'existing-message-id_',
|
||||
};
|
||||
|
||||
await TestClient.setMessageOptions(opts);
|
||||
|
||||
expect(TestClient.responseMessageId).toBe('existing-message-id_');
|
||||
});
|
||||
|
||||
test('should not replace responseMessageId when it does not end with underscore', async () => {
|
||||
const opts = {
|
||||
isRegenerate: true,
|
||||
responseMessageId: 'existing-message-id',
|
||||
};
|
||||
|
||||
await TestClient.setMessageOptions(opts);
|
||||
|
||||
expect(TestClient.responseMessageId).toBe('existing-message-id');
|
||||
});
|
||||
|
||||
test('sendMessage should work with provided conversationId and parentMessageId', async () => {
|
||||
const userMessage = 'Second message in the conversation';
|
||||
const opts = {
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
getReqData: jest.fn(),
|
||||
onStart: jest.fn(),
|
||||
};
|
||||
|
||||
const expectedResult = expect.objectContaining({
|
||||
sender: TestClient.sender,
|
||||
text: expect.any(String),
|
||||
isCreatedByUser: false,
|
||||
messageId: expect.any(String),
|
||||
parentMessageId: expect.any(String),
|
||||
conversationId: opts.conversationId,
|
||||
});
|
||||
|
||||
const response = await TestClient.sendMessage(userMessage, opts);
|
||||
parentMessageId = response.messageId;
|
||||
expect(response.conversationId).toEqual(conversationId);
|
||||
expect(response).toEqual(expectedResult);
|
||||
expect(opts.getReqData).toHaveBeenCalled();
|
||||
expect(opts.onStart).toHaveBeenCalled();
|
||||
expect(TestClient.getBuildMessagesOptions).toHaveBeenCalled();
|
||||
expect(TestClient.getSaveOptions).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test('should return chat history', async () => {
|
||||
TestClient = initializeFakeClient(apiKey, options, messageHistory);
|
||||
const chatMessages = await TestClient.loadHistory(conversationId, '2');
|
||||
expect(TestClient.currentMessages).toHaveLength(2);
|
||||
expect(chatMessages[0].text).toEqual('Hello');
|
||||
|
||||
const chatMessages2 = await TestClient.loadHistory(conversationId, '3');
|
||||
expect(TestClient.currentMessages).toHaveLength(3);
|
||||
expect(chatMessages2[chatMessages2.length - 1].text).toEqual("What's up");
|
||||
});
|
||||
|
||||
/* Most of the new sendMessage logic revolving around edited/continued AI messages
|
||||
* can be summarized by the following test. The condition will load the entire history up to
|
||||
* the message that is being edited, which will trigger the AI API to 'continue' the response.
|
||||
* The 'userMessage' is only passed by convention and is not necessary for the generation.
|
||||
*/
|
||||
it('should not push userMessage to currentMessages when isEdited is true and vice versa', async () => {
|
||||
const overrideParentMessageId = 'user-message-id';
|
||||
const responseMessageId = 'response-message-id';
|
||||
const newHistory = messageHistory.slice();
|
||||
newHistory.push({
|
||||
role: 'assistant',
|
||||
isCreatedByUser: false,
|
||||
text: 'test message',
|
||||
messageId: responseMessageId,
|
||||
parentMessageId: '3',
|
||||
});
|
||||
|
||||
TestClient = initializeFakeClient(apiKey, options, newHistory);
|
||||
const sendMessageOptions = {
|
||||
isEdited: true,
|
||||
overrideParentMessageId,
|
||||
parentMessageId: '3',
|
||||
responseMessageId,
|
||||
};
|
||||
|
||||
await TestClient.sendMessage('test message', sendMessageOptions);
|
||||
const currentMessages = TestClient.currentMessages;
|
||||
expect(currentMessages[currentMessages.length - 1].messageId).not.toEqual(
|
||||
overrideParentMessageId,
|
||||
);
|
||||
|
||||
// Test the opposite case
|
||||
sendMessageOptions.isEdited = false;
|
||||
await TestClient.sendMessage('test message', sendMessageOptions);
|
||||
const currentMessages2 = TestClient.currentMessages;
|
||||
expect(currentMessages2[currentMessages2.length - 1].messageId).toEqual(
|
||||
overrideParentMessageId,
|
||||
);
|
||||
});
|
||||
|
||||
test('setOptions is called with the correct arguments only when replaceOptions is set to true', async () => {
|
||||
TestClient.setOptions = jest.fn();
|
||||
const opts = { conversationId: '123', parentMessageId: '456', replaceOptions: true };
|
||||
await TestClient.sendMessage('Hello, world!', opts);
|
||||
expect(TestClient.setOptions).toHaveBeenCalledWith(opts);
|
||||
TestClient.setOptions.mockClear();
|
||||
});
|
||||
|
||||
test('loadHistory is called with the correct arguments', async () => {
|
||||
const opts = { conversationId: '123', parentMessageId: '456' };
|
||||
await TestClient.sendMessage('Hello, world!', opts);
|
||||
expect(TestClient.loadHistory).toHaveBeenCalledWith(
|
||||
opts.conversationId,
|
||||
opts.parentMessageId,
|
||||
);
|
||||
});
|
||||
|
||||
test('getReqData is called with the correct arguments', async () => {
|
||||
const getReqData = jest.fn();
|
||||
const opts = { getReqData };
|
||||
const response = await TestClient.sendMessage('Hello, world!', opts);
|
||||
expect(getReqData).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
userMessage: expect.objectContaining({ text: 'Hello, world!' }),
|
||||
conversationId: response.conversationId,
|
||||
responseMessageId: response.messageId,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
test('onStart is called with the correct arguments', async () => {
|
||||
const onStart = jest.fn();
|
||||
const opts = { onStart };
|
||||
await TestClient.sendMessage('Hello, world!', opts);
|
||||
|
||||
expect(onStart).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ text: 'Hello, world!' }),
|
||||
expect.any(String),
|
||||
/** `isNewConvo` */
|
||||
true,
|
||||
);
|
||||
});
|
||||
|
||||
test('saveMessageToDatabase is called with the correct arguments', async () => {
|
||||
const saveOptions = TestClient.getSaveOptions();
|
||||
const user = {};
|
||||
const opts = { user };
|
||||
const saveSpy = jest.spyOn(TestClient, 'saveMessageToDatabase');
|
||||
await TestClient.sendMessage('Hello, world!', opts);
|
||||
expect(saveSpy).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
sender: expect.any(String),
|
||||
text: expect.any(String),
|
||||
isCreatedByUser: expect.any(Boolean),
|
||||
messageId: expect.any(String),
|
||||
parentMessageId: expect.any(String),
|
||||
conversationId: expect.any(String),
|
||||
}),
|
||||
saveOptions,
|
||||
user,
|
||||
);
|
||||
});
|
||||
|
||||
test('should handle existing conversation when getConvo retrieves one', async () => {
|
||||
const existingConvo = {
|
||||
conversationId: 'existing-convo-id',
|
||||
endpoint: 'openai',
|
||||
endpointType: 'openai',
|
||||
model: 'gpt-3.5-turbo',
|
||||
messages: [
|
||||
{ role: 'user', content: 'Existing message 1' },
|
||||
{ role: 'assistant', content: 'Existing response 1' },
|
||||
],
|
||||
temperature: 1,
|
||||
};
|
||||
|
||||
const { temperature: _temp, ...newConvo } = existingConvo;
|
||||
|
||||
const user = {
|
||||
id: 'user-id',
|
||||
};
|
||||
|
||||
getConvo.mockResolvedValue(existingConvo);
|
||||
saveConvo.mockResolvedValue(newConvo);
|
||||
|
||||
TestClient = initializeFakeClient(
|
||||
apiKey,
|
||||
{
|
||||
...options,
|
||||
req: {
|
||||
user,
|
||||
},
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
||||
const saveSpy = jest.spyOn(TestClient, 'saveMessageToDatabase');
|
||||
|
||||
const newMessage = 'New message in existing conversation';
|
||||
const response = await TestClient.sendMessage(newMessage, {
|
||||
user,
|
||||
conversationId: existingConvo.conversationId,
|
||||
});
|
||||
|
||||
expect(getConvo).toHaveBeenCalledWith(user.id, existingConvo.conversationId);
|
||||
expect(TestClient.conversationId).toBe(existingConvo.conversationId);
|
||||
expect(response.conversationId).toBe(existingConvo.conversationId);
|
||||
expect(TestClient.fetchedConvo).toBe(true);
|
||||
|
||||
expect(saveSpy).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
conversationId: existingConvo.conversationId,
|
||||
text: newMessage,
|
||||
}),
|
||||
expect.any(Object),
|
||||
expect.any(Object),
|
||||
);
|
||||
|
||||
expect(saveConvo).toHaveBeenCalledTimes(2);
|
||||
expect(saveConvo).toHaveBeenCalledWith(
|
||||
expect.any(Object),
|
||||
expect.objectContaining({
|
||||
conversationId: existingConvo.conversationId,
|
||||
}),
|
||||
expect.objectContaining({
|
||||
context: 'api/app/clients/BaseClient.js - saveMessageToDatabase #saveConvo',
|
||||
unsetFields: {
|
||||
temperature: 1,
|
||||
},
|
||||
}),
|
||||
);
|
||||
|
||||
await TestClient.sendMessage('Another message', {
|
||||
conversationId: existingConvo.conversationId,
|
||||
});
|
||||
expect(getConvo).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
test('should correctly handle existing conversation and unset fields appropriately', async () => {
|
||||
const existingConvo = {
|
||||
conversationId: 'existing-convo-id',
|
||||
endpoint: 'openai',
|
||||
endpointType: 'openai',
|
||||
model: 'gpt-3.5-turbo',
|
||||
messages: [
|
||||
{ role: 'user', content: 'Existing message 1' },
|
||||
{ role: 'assistant', content: 'Existing response 1' },
|
||||
],
|
||||
title: 'Existing Conversation',
|
||||
someExistingField: 'existingValue',
|
||||
anotherExistingField: 'anotherValue',
|
||||
temperature: 0.7,
|
||||
modelLabel: 'GPT-3.5',
|
||||
};
|
||||
|
||||
getConvo.mockResolvedValue(existingConvo);
|
||||
saveConvo.mockResolvedValue(existingConvo);
|
||||
|
||||
TestClient = initializeFakeClient(
|
||||
apiKey,
|
||||
{
|
||||
...options,
|
||||
modelOptions: {
|
||||
model: 'gpt-4',
|
||||
temperature: 0.5,
|
||||
},
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
||||
const newMessage = 'New message in existing conversation';
|
||||
await TestClient.sendMessage(newMessage, {
|
||||
conversationId: existingConvo.conversationId,
|
||||
});
|
||||
|
||||
expect(saveConvo).toHaveBeenCalledTimes(2);
|
||||
|
||||
const saveConvoCall = saveConvo.mock.calls[0];
|
||||
const [, savedFields, saveOptions] = saveConvoCall;
|
||||
|
||||
// Instead of checking all excludedKeys, we'll just check specific fields
|
||||
// that we know should be excluded
|
||||
expect(savedFields).not.toHaveProperty('messages');
|
||||
expect(savedFields).not.toHaveProperty('title');
|
||||
|
||||
// Only check that someExistingField is in unsetFields
|
||||
expect(saveOptions.unsetFields).toHaveProperty('someExistingField', 1);
|
||||
|
||||
// Mock saveConvo to return the expected fields
|
||||
saveConvo.mockImplementation((req, fields) => {
|
||||
return Promise.resolve({
|
||||
...fields,
|
||||
endpoint: 'openai',
|
||||
endpointType: 'openai',
|
||||
model: 'gpt-4',
|
||||
temperature: 0.5,
|
||||
});
|
||||
});
|
||||
|
||||
// Only check the conversationId since that's the only field we can be sure about
|
||||
expect(savedFields).toHaveProperty('conversationId', 'existing-convo-id');
|
||||
|
||||
expect(TestClient.fetchedConvo).toBe(true);
|
||||
|
||||
await TestClient.sendMessage('Another message', {
|
||||
conversationId: existingConvo.conversationId,
|
||||
});
|
||||
|
||||
expect(getConvo).toHaveBeenCalledTimes(1);
|
||||
|
||||
const secondSaveConvoCall = saveConvo.mock.calls[1];
|
||||
expect(secondSaveConvoCall[2]).toHaveProperty('unsetFields', {});
|
||||
});
|
||||
|
||||
test('sendCompletion is called with the correct arguments', async () => {
|
||||
const payload = {}; // Mock payload
|
||||
TestClient.buildMessages.mockReturnValue({ prompt: payload, tokenCountMap: null });
|
||||
const opts = {};
|
||||
await TestClient.sendMessage('Hello, world!', opts);
|
||||
expect(TestClient.sendCompletion).toHaveBeenCalledWith(payload, opts);
|
||||
});
|
||||
|
||||
test('getTokenCount for response is called with the correct arguments', async () => {
|
||||
const tokenCountMap = {}; // Mock tokenCountMap
|
||||
TestClient.buildMessages.mockReturnValue({ prompt: [], tokenCountMap });
|
||||
TestClient.getTokenCountForResponse = jest.fn();
|
||||
const response = await TestClient.sendMessage('Hello, world!', {});
|
||||
expect(TestClient.getTokenCountForResponse).toHaveBeenCalledWith(response);
|
||||
});
|
||||
|
||||
test('returns an object with the correct shape', async () => {
|
||||
const response = await TestClient.sendMessage('Hello, world!', {});
|
||||
expect(response).toEqual(
|
||||
expect.objectContaining({
|
||||
sender: expect.any(String),
|
||||
text: expect.any(String),
|
||||
isCreatedByUser: expect.any(Boolean),
|
||||
messageId: expect.any(String),
|
||||
parentMessageId: expect.any(String),
|
||||
conversationId: expect.any(String),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
test('userMessagePromise is awaited before saving response message', async () => {
|
||||
// Mock the saveMessageToDatabase method
|
||||
TestClient.saveMessageToDatabase = jest.fn().mockImplementation(() => {
|
||||
return new Promise((resolve) => setTimeout(resolve, 100)); // Simulate a delay
|
||||
});
|
||||
|
||||
// Send a message
|
||||
const messagePromise = TestClient.sendMessage('Hello, world!');
|
||||
|
||||
// Wait a short time to ensure the user message save has started
|
||||
await new Promise((resolve) => setTimeout(resolve, 50));
|
||||
|
||||
// Check that saveMessageToDatabase has been called once (for the user message)
|
||||
expect(TestClient.saveMessageToDatabase).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Wait for the message to be fully processed
|
||||
await messagePromise;
|
||||
|
||||
// Check that saveMessageToDatabase has been called twice (once for user message, once for response)
|
||||
expect(TestClient.saveMessageToDatabase).toHaveBeenCalledTimes(2);
|
||||
|
||||
// Check the order of calls
|
||||
const calls = TestClient.saveMessageToDatabase.mock.calls;
|
||||
expect(calls[0][0].isCreatedByUser).toBe(true); // First call should be for user message
|
||||
expect(calls[1][0].isCreatedByUser).toBe(false); // Second call should be for response message
|
||||
});
|
||||
});
|
||||
|
||||
describe('getMessagesWithinTokenLimit with instructions', () => {
|
||||
test('should always include instructions when present', async () => {
|
||||
TestClient.maxContextTokens = 50;
|
||||
const instructions = {
|
||||
role: 'system',
|
||||
content: 'System instructions',
|
||||
tokenCount: 20,
|
||||
};
|
||||
|
||||
const messages = [
|
||||
instructions,
|
||||
{ role: 'user', content: 'Hello', tokenCount: 10 },
|
||||
{ role: 'assistant', content: 'Hi there', tokenCount: 15 },
|
||||
];
|
||||
|
||||
const result = await TestClient.getMessagesWithinTokenLimit({
|
||||
messages,
|
||||
instructions,
|
||||
});
|
||||
|
||||
expect(result.context[0]).toBe(instructions);
|
||||
expect(result.remainingContextTokens).toBe(2);
|
||||
});
|
||||
|
||||
test('should handle case when messages exceed limit but instructions must be preserved', async () => {
|
||||
TestClient.maxContextTokens = 30;
|
||||
const instructions = {
|
||||
role: 'system',
|
||||
content: 'System instructions',
|
||||
tokenCount: 20,
|
||||
};
|
||||
|
||||
const messages = [
|
||||
instructions,
|
||||
{ role: 'user', content: 'Hello', tokenCount: 10 },
|
||||
{ role: 'assistant', content: 'Hi there', tokenCount: 15 },
|
||||
];
|
||||
|
||||
const result = await TestClient.getMessagesWithinTokenLimit({
|
||||
messages,
|
||||
instructions,
|
||||
});
|
||||
|
||||
// Should only include instructions and the last message that fits
|
||||
expect(result.context).toHaveLength(1);
|
||||
expect(result.context[0].content).toBe(instructions.content);
|
||||
expect(result.messagesToRefine).toHaveLength(2);
|
||||
expect(result.remainingContextTokens).toBe(7); // 30 - 20 - 3 (assistant label)
|
||||
});
|
||||
|
||||
test('should work correctly without instructions (1/2)', async () => {
|
||||
TestClient.maxContextTokens = 50;
|
||||
const messages = [
|
||||
{ role: 'user', content: 'Hello', tokenCount: 10 },
|
||||
{ role: 'assistant', content: 'Hi there', tokenCount: 15 },
|
||||
];
|
||||
|
||||
const result = await TestClient.getMessagesWithinTokenLimit({
|
||||
messages,
|
||||
});
|
||||
|
||||
expect(result.context).toHaveLength(2);
|
||||
expect(result.remainingContextTokens).toBe(22); // 50 - 10 - 15 - 3(assistant label)
|
||||
expect(result.messagesToRefine).toHaveLength(0);
|
||||
});
|
||||
|
||||
test('should work correctly without instructions (2/2)', async () => {
|
||||
TestClient.maxContextTokens = 30;
|
||||
const messages = [
|
||||
{ role: 'user', content: 'Hello', tokenCount: 10 },
|
||||
{ role: 'assistant', content: 'Hi there', tokenCount: 20 },
|
||||
];
|
||||
|
||||
const result = await TestClient.getMessagesWithinTokenLimit({
|
||||
messages,
|
||||
});
|
||||
|
||||
expect(result.context).toHaveLength(1);
|
||||
expect(result.remainingContextTokens).toBe(7);
|
||||
expect(result.messagesToRefine).toHaveLength(1);
|
||||
});
|
||||
|
||||
test('should handle case when only instructions fit within limit', async () => {
|
||||
TestClient.maxContextTokens = 25;
|
||||
const instructions = {
|
||||
role: 'system',
|
||||
content: 'System instructions',
|
||||
tokenCount: 20,
|
||||
};
|
||||
|
||||
const messages = [
|
||||
instructions,
|
||||
{ role: 'user', content: 'Hello', tokenCount: 10 },
|
||||
{ role: 'assistant', content: 'Hi there', tokenCount: 15 },
|
||||
];
|
||||
|
||||
const result = await TestClient.getMessagesWithinTokenLimit({
|
||||
messages,
|
||||
instructions,
|
||||
});
|
||||
|
||||
expect(result.context).toHaveLength(1);
|
||||
expect(result.context[0]).toBe(instructions);
|
||||
expect(result.messagesToRefine).toHaveLength(2);
|
||||
expect(result.remainingContextTokens).toBe(2); // 25 - 20 - 3(assistant label)
|
||||
});
|
||||
});
|
||||
});
|
||||
126
api/app/clients/specs/FakeClient.js
Normal file
126
api/app/clients/specs/FakeClient.js
Normal file
@ -0,0 +1,126 @@
|
||||
const { getModelMaxTokens } = require('@librechat/api');
|
||||
const BaseClient = require('../BaseClient');
|
||||
|
||||
class FakeClient extends BaseClient {
|
||||
constructor(apiKey, options = {}) {
|
||||
super(apiKey, options);
|
||||
this.sender = 'AI Assistant';
|
||||
this.setOptions(options);
|
||||
}
|
||||
setOptions(options) {
|
||||
if (this.options && !this.options.replaceOptions) {
|
||||
this.options.modelOptions = {
|
||||
...this.options.modelOptions,
|
||||
...options.modelOptions,
|
||||
};
|
||||
delete options.modelOptions;
|
||||
this.options = {
|
||||
...this.options,
|
||||
...options,
|
||||
};
|
||||
} else {
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
if (this.options.openaiApiKey) {
|
||||
this.apiKey = this.options.openaiApiKey;
|
||||
}
|
||||
|
||||
const modelOptions = this.options.modelOptions || {};
|
||||
if (!this.modelOptions) {
|
||||
this.modelOptions = {
|
||||
...modelOptions,
|
||||
model: modelOptions.model || 'gpt-3.5-turbo',
|
||||
temperature:
|
||||
typeof modelOptions.temperature === 'undefined' ? 0.8 : modelOptions.temperature,
|
||||
top_p: typeof modelOptions.top_p === 'undefined' ? 1 : modelOptions.top_p,
|
||||
presence_penalty:
|
||||
typeof modelOptions.presence_penalty === 'undefined' ? 1 : modelOptions.presence_penalty,
|
||||
stop: modelOptions.stop,
|
||||
};
|
||||
}
|
||||
|
||||
this.maxContextTokens =
|
||||
this.options.maxContextTokens ?? getModelMaxTokens(this.modelOptions.model) ?? 4097;
|
||||
}
|
||||
buildMessages() {}
|
||||
getTokenCount(str) {
|
||||
return str.length;
|
||||
}
|
||||
getTokenCountForMessage(message) {
|
||||
return message?.content?.length || message.length;
|
||||
}
|
||||
}
|
||||
|
||||
const initializeFakeClient = (apiKey, options, fakeMessages) => {
|
||||
let TestClient = new FakeClient(apiKey);
|
||||
TestClient.options = options;
|
||||
TestClient.abortController = { abort: jest.fn() };
|
||||
TestClient.loadHistory = jest
|
||||
.fn()
|
||||
.mockImplementation((conversationId, parentMessageId = null) => {
|
||||
if (!conversationId) {
|
||||
TestClient.currentMessages = [];
|
||||
return Promise.resolve([]);
|
||||
}
|
||||
|
||||
const orderedMessages = TestClient.constructor.getMessagesForConversation({
|
||||
messages: fakeMessages,
|
||||
parentMessageId,
|
||||
});
|
||||
|
||||
TestClient.currentMessages = orderedMessages;
|
||||
return Promise.resolve(orderedMessages);
|
||||
});
|
||||
|
||||
TestClient.getSaveOptions = jest.fn().mockImplementation(() => {
|
||||
return {};
|
||||
});
|
||||
|
||||
TestClient.getBuildMessagesOptions = jest.fn().mockImplementation(() => {
|
||||
return {};
|
||||
});
|
||||
|
||||
TestClient.sendCompletion = jest.fn(async () => {
|
||||
return {
|
||||
completion: 'Mock response text',
|
||||
metadata: undefined,
|
||||
};
|
||||
});
|
||||
|
||||
TestClient.getCompletion = jest.fn().mockImplementation(async (..._args) => {
|
||||
return {
|
||||
choices: [
|
||||
{
|
||||
message: {
|
||||
content: 'Mock response text',
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
});
|
||||
|
||||
TestClient.buildMessages = jest.fn(async (messages, parentMessageId) => {
|
||||
const orderedMessages = TestClient.constructor.getMessagesForConversation({
|
||||
messages,
|
||||
parentMessageId,
|
||||
});
|
||||
const formattedMessages = orderedMessages.map((message) => {
|
||||
let { role: _role, sender, text } = message;
|
||||
const role = _role ?? sender;
|
||||
const content = text ?? '';
|
||||
return {
|
||||
role: role?.toLowerCase() === 'user' ? 'user' : 'assistant',
|
||||
content,
|
||||
};
|
||||
});
|
||||
return {
|
||||
prompt: formattedMessages,
|
||||
tokenCountMap: null, // Simplified for the mock
|
||||
};
|
||||
});
|
||||
|
||||
return TestClient;
|
||||
};
|
||||
|
||||
module.exports = { FakeClient, initializeFakeClient };
|
||||
30
api/app/clients/tools/index.js
Normal file
30
api/app/clients/tools/index.js
Normal file
@ -0,0 +1,30 @@
|
||||
const manifest = require('./manifest');
|
||||
|
||||
// Structured Tools
|
||||
const DALLE3 = require('./structured/DALLE3');
|
||||
const FluxAPI = require('./structured/FluxAPI');
|
||||
const OpenWeather = require('./structured/OpenWeather');
|
||||
const StructuredWolfram = require('./structured/Wolfram');
|
||||
const StructuredACS = require('./structured/AzureAISearch');
|
||||
const StructuredSD = require('./structured/StableDiffusion');
|
||||
const GoogleSearchAPI = require('./structured/GoogleSearch');
|
||||
const TraversaalSearch = require('./structured/TraversaalSearch');
|
||||
const createOpenAIImageTools = require('./structured/OpenAIImageTools');
|
||||
const TavilySearchResults = require('./structured/TavilySearchResults');
|
||||
const createGeminiImageTool = require('./structured/GeminiImageGen');
|
||||
|
||||
module.exports = {
|
||||
...manifest,
|
||||
// Structured Tools
|
||||
DALLE3,
|
||||
FluxAPI,
|
||||
OpenWeather,
|
||||
StructuredSD,
|
||||
StructuredACS,
|
||||
GoogleSearchAPI,
|
||||
TraversaalSearch,
|
||||
StructuredWolfram,
|
||||
TavilySearchResults,
|
||||
createOpenAIImageTools,
|
||||
createGeminiImageTool,
|
||||
};
|
||||
20
api/app/clients/tools/manifest.js
Normal file
20
api/app/clients/tools/manifest.js
Normal file
@ -0,0 +1,20 @@
|
||||
const availableTools = require('./manifest.json');
|
||||
|
||||
/** @type {Record<string, TPlugin | undefined>} */
|
||||
const manifestToolMap = {};
|
||||
|
||||
/** @type {Array<TPlugin>} */
|
||||
const toolkits = [];
|
||||
|
||||
availableTools.forEach((tool) => {
|
||||
manifestToolMap[tool.pluginKey] = tool;
|
||||
if (tool.toolkit === true) {
|
||||
toolkits.push(tool);
|
||||
}
|
||||
});
|
||||
|
||||
module.exports = {
|
||||
toolkits,
|
||||
availableTools,
|
||||
manifestToolMap,
|
||||
};
|
||||
183
api/app/clients/tools/manifest.json
Normal file
183
api/app/clients/tools/manifest.json
Normal file
@ -0,0 +1,183 @@
|
||||
[
|
||||
{
|
||||
"name": "Traversaal",
|
||||
"pluginKey": "traversaal_search",
|
||||
"description": "Traversaal is a robust search API tailored for LLM Agents. Get an API key here: https://api.traversaal.ai",
|
||||
"icon": "https://traversaal.ai/favicon.ico",
|
||||
"authConfig": [
|
||||
{
|
||||
"authField": "TRAVERSAAL_API_KEY",
|
||||
"label": "Traversaal API Key",
|
||||
"description": "Get your API key here: <a href=\"https://api.traversaal.ai\" target=\"_blank\">https://api.traversaal.ai</a>"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "Google",
|
||||
"pluginKey": "google",
|
||||
"description": "Use Google Search to find information about the weather, news, sports, and more.",
|
||||
"icon": "https://i.imgur.com/SMmVkNB.png",
|
||||
"authConfig": [
|
||||
{
|
||||
"authField": "GOOGLE_CSE_ID",
|
||||
"label": "Google CSE ID",
|
||||
"description": "This is your Google Custom Search Engine ID. For instructions on how to obtain this, see <a href='https://github.com/danny-avila/QognyX/blob/main/docs/features/plugins/google_search.md'>Our Docs</a>."
|
||||
},
|
||||
{
|
||||
"authField": "GOOGLE_SEARCH_API_KEY",
|
||||
"label": "Google API Key",
|
||||
"description": "This is your Google Custom Search API Key. For instructions on how to obtain this, see <a href='https://github.com/danny-avila/QognyX/blob/main/docs/features/plugins/google_search.md'>Our Docs</a>."
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "OpenAI Image Tools",
|
||||
"pluginKey": "image_gen_oai",
|
||||
"toolkit": true,
|
||||
"description": "Image Generation and Editing using OpenAI's latest state-of-the-art models",
|
||||
"icon": "assets/image_gen_oai.png",
|
||||
"authConfig": [
|
||||
{
|
||||
"authField": "IMAGE_GEN_OAI_API_KEY",
|
||||
"label": "OpenAI Image Tools API Key",
|
||||
"description": "Your OpenAI API Key for Image Generation and Editing"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "Wolfram",
|
||||
"pluginKey": "wolfram",
|
||||
"description": "Access computation, math, curated knowledge & real-time data through Wolfram|Alpha and Wolfram Language.",
|
||||
"icon": "https://www.wolframcdn.com/images/icons/Wolfram.png",
|
||||
"authConfig": [
|
||||
{
|
||||
"authField": "WOLFRAM_APP_ID",
|
||||
"label": "Wolfram App ID",
|
||||
"description": "An AppID must be supplied in all calls to the Wolfram|Alpha API. You can get one by registering at <a href='http://products.wolframalpha.com/api/'>Wolfram|Alpha</a> and going to the <a href='https://developer.wolframalpha.com/portal/myapps/'>Developer Portal</a>."
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "Browser",
|
||||
"pluginKey": "web-browser",
|
||||
"description": "Scrape and summarize webpage data",
|
||||
"icon": "assets/web-browser.svg",
|
||||
"authConfig": [
|
||||
{
|
||||
"authField": "OPENAI_API_KEY",
|
||||
"label": "OpenAI API Key",
|
||||
"description": "Browser makes use of OpenAI embeddings"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "DALL-E-3",
|
||||
"pluginKey": "dalle",
|
||||
"description": "[DALL-E-3] Create realistic images and art from a description in natural language",
|
||||
"icon": "https://i.imgur.com/u2TzXzH.png",
|
||||
"authConfig": [
|
||||
{
|
||||
"authField": "DALLE3_API_KEY||DALLE_API_KEY",
|
||||
"label": "OpenAI API Key",
|
||||
"description": "You can use DALL-E with your API Key from OpenAI."
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "Tavily Search",
|
||||
"pluginKey": "tavily_search_results_json",
|
||||
"description": "Tavily Search is a robust search API tailored for LLM Agents. It seamlessly integrates with diverse data sources to ensure a superior, relevant search experience.",
|
||||
"icon": "https://tavily.com/favicon.ico",
|
||||
"authConfig": [
|
||||
{
|
||||
"authField": "TAVILY_API_KEY",
|
||||
"label": "Tavily API Key",
|
||||
"description": "Get your API key here: https://app.tavily.com/"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "Calculator",
|
||||
"pluginKey": "calculator",
|
||||
"description": "Perform simple and complex mathematical calculations.",
|
||||
"icon": "https://i.imgur.com/RHsSG5h.png",
|
||||
"authConfig": []
|
||||
},
|
||||
{
|
||||
"name": "Stable Diffusion",
|
||||
"pluginKey": "stable-diffusion",
|
||||
"description": "Generate photo-realistic images given any text input.",
|
||||
"icon": "https://i.imgur.com/Yr466dp.png",
|
||||
"authConfig": [
|
||||
{
|
||||
"authField": "SD_WEBUI_URL",
|
||||
"label": "Your Stable Diffusion WebUI API URL",
|
||||
"description": "You need to provide the URL of your Stable Diffusion WebUI API. For instructions on how to obtain this, see <a href='url'>Our Docs</a>."
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "Azure AI Search",
|
||||
"pluginKey": "azure-ai-search",
|
||||
"description": "Use Azure AI Search to find information",
|
||||
"icon": "https://i.imgur.com/E7crPze.png",
|
||||
"authConfig": [
|
||||
{
|
||||
"authField": "AZURE_AI_SEARCH_SERVICE_ENDPOINT",
|
||||
"label": "Azure AI Search Endpoint",
|
||||
"description": "You need to provide your Endpoint for Azure AI Search."
|
||||
},
|
||||
{
|
||||
"authField": "AZURE_AI_SEARCH_INDEX_NAME",
|
||||
"label": "Azure AI Search Index Name",
|
||||
"description": "You need to provide your Index Name for Azure AI Search."
|
||||
},
|
||||
{
|
||||
"authField": "AZURE_AI_SEARCH_API_KEY",
|
||||
"label": "Azure AI Search API Key",
|
||||
"description": "You need to provide your API Key for Azure AI Search."
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "OpenWeather",
|
||||
"pluginKey": "open_weather",
|
||||
"description": "Get weather forecasts and historical data from the OpenWeather API",
|
||||
"icon": "assets/openweather.png",
|
||||
"authConfig": [
|
||||
{
|
||||
"authField": "OPENWEATHER_API_KEY",
|
||||
"label": "OpenWeather API Key",
|
||||
"description": "Sign up at <a href=\"https://home.openweathermap.org/users/sign_up\" target=\"_blank\">OpenWeather</a>, then get your key at <a href=\"https://home.openweathermap.org/api_keys\" target=\"_blank\">API keys</a>."
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "Flux",
|
||||
"pluginKey": "flux",
|
||||
"description": "Generate images using text with the Flux API.",
|
||||
"icon": "https://blackforestlabs.ai/wp-content/uploads/2024/07/bfl_logo_retraced_blk.png",
|
||||
"isAuthRequired": "true",
|
||||
"authConfig": [
|
||||
{
|
||||
"authField": "FLUX_API_KEY",
|
||||
"label": "Your Flux API Key",
|
||||
"description": "Provide your Flux API key from your user profile."
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "Gemini Image Tools",
|
||||
"pluginKey": "gemini_image_gen",
|
||||
"toolkit": true,
|
||||
"description": "Generate high-quality images using Google's Gemini Image Models. Supports Gemini API or Vertex AI.",
|
||||
"icon": "assets/gemini_image_gen.svg",
|
||||
"authConfig": [
|
||||
{
|
||||
"authField": "GEMINI_API_KEY||GOOGLE_KEY||GEMINI_VERTEX_ENABLED",
|
||||
"label": "Gemini API Key (Optional if Vertex AI is configured)",
|
||||
"description": "Your Google Gemini API Key from <a href='https://aistudio.google.com/app/apikey' target='_blank'>Google AI Studio</a>. Leave blank if using Vertex AI with service account."
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
104
api/app/clients/tools/structured/AzureAISearch.js
Normal file
104
api/app/clients/tools/structured/AzureAISearch.js
Normal file
@ -0,0 +1,104 @@
|
||||
const { z } = require('zod');
|
||||
const { Tool } = require('@langchain/core/tools');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { SearchClient, AzureKeyCredential } = require('@azure/search-documents');
|
||||
|
||||
class AzureAISearch extends Tool {
|
||||
// Constants for default values
|
||||
static DEFAULT_API_VERSION = '2023-11-01';
|
||||
static DEFAULT_QUERY_TYPE = 'simple';
|
||||
static DEFAULT_TOP = 5;
|
||||
|
||||
// Helper function for initializing properties
|
||||
_initializeField(field, envVar, defaultValue) {
|
||||
return field || process.env[envVar] || defaultValue;
|
||||
}
|
||||
|
||||
constructor(fields = {}) {
|
||||
super();
|
||||
this.name = 'azure-ai-search';
|
||||
this.description =
|
||||
"Use the 'azure-ai-search' tool to retrieve search results relevant to your input";
|
||||
/* Used to initialize the Tool without necessary variables. */
|
||||
this.override = fields.override ?? false;
|
||||
|
||||
// Define schema
|
||||
this.schema = z.object({
|
||||
query: z.string().describe('Search word or phrase to Azure AI Search'),
|
||||
});
|
||||
|
||||
// Initialize properties using helper function
|
||||
this.serviceEndpoint = this._initializeField(
|
||||
fields.AZURE_AI_SEARCH_SERVICE_ENDPOINT,
|
||||
'AZURE_AI_SEARCH_SERVICE_ENDPOINT',
|
||||
);
|
||||
this.indexName = this._initializeField(
|
||||
fields.AZURE_AI_SEARCH_INDEX_NAME,
|
||||
'AZURE_AI_SEARCH_INDEX_NAME',
|
||||
);
|
||||
this.apiKey = this._initializeField(fields.AZURE_AI_SEARCH_API_KEY, 'AZURE_AI_SEARCH_API_KEY');
|
||||
this.apiVersion = this._initializeField(
|
||||
fields.AZURE_AI_SEARCH_API_VERSION,
|
||||
'AZURE_AI_SEARCH_API_VERSION',
|
||||
AzureAISearch.DEFAULT_API_VERSION,
|
||||
);
|
||||
this.queryType = this._initializeField(
|
||||
fields.AZURE_AI_SEARCH_SEARCH_OPTION_QUERY_TYPE,
|
||||
'AZURE_AI_SEARCH_SEARCH_OPTION_QUERY_TYPE',
|
||||
AzureAISearch.DEFAULT_QUERY_TYPE,
|
||||
);
|
||||
this.top = this._initializeField(
|
||||
fields.AZURE_AI_SEARCH_SEARCH_OPTION_TOP,
|
||||
'AZURE_AI_SEARCH_SEARCH_OPTION_TOP',
|
||||
AzureAISearch.DEFAULT_TOP,
|
||||
);
|
||||
this.select = this._initializeField(
|
||||
fields.AZURE_AI_SEARCH_SEARCH_OPTION_SELECT,
|
||||
'AZURE_AI_SEARCH_SEARCH_OPTION_SELECT',
|
||||
);
|
||||
|
||||
// Check for required fields
|
||||
if (!this.override && (!this.serviceEndpoint || !this.indexName || !this.apiKey)) {
|
||||
throw new Error(
|
||||
'Missing AZURE_AI_SEARCH_SERVICE_ENDPOINT, AZURE_AI_SEARCH_INDEX_NAME, or AZURE_AI_SEARCH_API_KEY environment variable.',
|
||||
);
|
||||
}
|
||||
|
||||
if (this.override) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Create SearchClient
|
||||
this.client = new SearchClient(
|
||||
this.serviceEndpoint,
|
||||
this.indexName,
|
||||
new AzureKeyCredential(this.apiKey),
|
||||
{ apiVersion: this.apiVersion },
|
||||
);
|
||||
}
|
||||
|
||||
// Improved error handling and logging
|
||||
async _call(data) {
|
||||
const { query } = data;
|
||||
try {
|
||||
const searchOption = {
|
||||
queryType: this.queryType,
|
||||
top: typeof this.top === 'string' ? Number(this.top) : this.top,
|
||||
};
|
||||
if (this.select) {
|
||||
searchOption.select = this.select.split(',');
|
||||
}
|
||||
const searchResults = await this.client.search(query, searchOption);
|
||||
const resultDocuments = [];
|
||||
for await (const result of searchResults.results) {
|
||||
resultDocuments.push(result.document);
|
||||
}
|
||||
return JSON.stringify(resultDocuments);
|
||||
} catch (error) {
|
||||
logger.error('Azure AI Search request failed', error);
|
||||
return 'There was an error with Azure AI Search.';
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = AzureAISearch;
|
||||
230
api/app/clients/tools/structured/DALLE3.js
Normal file
230
api/app/clients/tools/structured/DALLE3.js
Normal file
@ -0,0 +1,230 @@
|
||||
const { z } = require('zod');
|
||||
const path = require('path');
|
||||
const OpenAI = require('openai');
|
||||
const { v4: uuidv4 } = require('uuid');
|
||||
const { ProxyAgent, fetch } = require('undici');
|
||||
const { Tool } = require('@langchain/core/tools');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { getImageBasename, extractBaseURL } = require('@librechat/api');
|
||||
const { FileContext, ContentTypes } = require('librechat-data-provider');
|
||||
|
||||
const displayMessage =
|
||||
"DALL-E displayed an image. All generated images are already plainly visible, so don't repeat the descriptions in detail. Do not list download links as they are available in the UI already. The user may download the images by clicking on them, but do not mention anything about downloading to the user.";
|
||||
class DALLE3 extends Tool {
|
||||
constructor(fields = {}) {
|
||||
super();
|
||||
/** @type {boolean} Used to initialize the Tool without necessary variables. */
|
||||
this.override = fields.override ?? false;
|
||||
/** @type {boolean} Necessary for output to contain all image metadata. */
|
||||
this.returnMetadata = fields.returnMetadata ?? false;
|
||||
|
||||
this.userId = fields.userId;
|
||||
this.fileStrategy = fields.fileStrategy;
|
||||
/** @type {boolean} */
|
||||
this.isAgent = fields.isAgent;
|
||||
if (fields.processFileURL) {
|
||||
/** @type {processFileURL} Necessary for output to contain all image metadata. */
|
||||
this.processFileURL = fields.processFileURL.bind(this);
|
||||
}
|
||||
|
||||
let apiKey = fields.DALLE3_API_KEY ?? fields.DALLE_API_KEY ?? this.getApiKey();
|
||||
const config = { apiKey };
|
||||
if (process.env.DALLE_REVERSE_PROXY) {
|
||||
config.baseURL = extractBaseURL(process.env.DALLE_REVERSE_PROXY);
|
||||
}
|
||||
|
||||
if (process.env.DALLE3_AZURE_API_VERSION && process.env.DALLE3_BASEURL) {
|
||||
config.baseURL = process.env.DALLE3_BASEURL;
|
||||
config.defaultQuery = { 'api-version': process.env.DALLE3_AZURE_API_VERSION };
|
||||
config.defaultHeaders = {
|
||||
'api-key': process.env.DALLE3_API_KEY,
|
||||
'Content-Type': 'application/json',
|
||||
};
|
||||
config.apiKey = process.env.DALLE3_API_KEY;
|
||||
}
|
||||
|
||||
if (process.env.PROXY) {
|
||||
const proxyAgent = new ProxyAgent(process.env.PROXY);
|
||||
config.fetchOptions = {
|
||||
dispatcher: proxyAgent,
|
||||
};
|
||||
}
|
||||
|
||||
/** @type {OpenAI} */
|
||||
this.openai = new OpenAI(config);
|
||||
this.name = 'dalle';
|
||||
this.description = `Use DALLE to create images from text descriptions.
|
||||
- It requires prompts to be in English, detailed, and to specify image type and human features for diversity.
|
||||
- Create only one image, without repeating or listing descriptions outside the "prompts" field.
|
||||
- Maintains the original intent of the description, with parameters for image style, quality, and size to tailor the output.`;
|
||||
this.description_for_model =
|
||||
process.env.DALLE3_SYSTEM_PROMPT ??
|
||||
`// Whenever a description of an image is given, generate prompts (following these rules), and use dalle to create the image. If the user does not ask for a specific number of images, default to creating 2 prompts to send to dalle that are written to be as diverse as possible. All prompts sent to dalle must abide by the following policies:
|
||||
// 1. Prompts must be in English. Translate to English if needed.
|
||||
// 2. One image per function call. Create only 1 image per request unless explicitly told to generate more than 1 image.
|
||||
// 3. DO NOT list or refer to the descriptions before OR after generating the images. They should ONLY ever be written out ONCE, in the \`"prompts"\` field of the request. You do not need to ask for permission to generate, just do it!
|
||||
// 4. Always mention the image type (photo, oil painting, watercolor painting, illustration, cartoon, drawing, vector, render, etc.) at the beginning of the caption. Unless the captions suggests otherwise, make one of the images a photo.
|
||||
// 5. Diversify depictions of ALL images with people to always include always DESCENT and GENDER for EACH person using direct terms. Adjust only human descriptions.
|
||||
// - EXPLICITLY specify these attributes, not abstractly reference them. The attributes should be specified in a minimal way and should directly describe their physical form.
|
||||
// - Your choices should be grounded in reality. For example, all of a given OCCUPATION should not be the same gender or race. Additionally, focus on creating diverse, inclusive, and exploratory scenes via the properties you choose during rewrites. Make choices that may be insightful or unique sometimes.
|
||||
// - Use "various" or "diverse" ONLY IF the description refers to groups of more than 3 people. Do not change the number of people requested in the original description.
|
||||
// - Don't alter memes, fictional character origins, or unseen people. Maintain the original prompt's intent and prioritize quality.
|
||||
// The prompt must intricately describe every part of the image in concrete, objective detail. THINK about what the end goal of the description is, and extrapolate that to what would make satisfying images.
|
||||
// All descriptions sent to dalle should be a paragraph of text that is extremely descriptive and detailed. Each should be more than 3 sentences long.
|
||||
// - The "vivid" style is HIGHLY preferred, but "natural" is also supported.`;
|
||||
this.schema = z.object({
|
||||
prompt: z
|
||||
.string()
|
||||
.max(4000)
|
||||
.describe(
|
||||
'A text description of the desired image, following the rules, up to 4000 characters.',
|
||||
),
|
||||
style: z
|
||||
.enum(['vivid', 'natural'])
|
||||
.describe(
|
||||
'Must be one of `vivid` or `natural`. `vivid` generates hyper-real and dramatic images, `natural` produces more natural, less hyper-real looking images',
|
||||
),
|
||||
quality: z
|
||||
.enum(['hd', 'standard'])
|
||||
.describe('The quality of the generated image. Only `hd` and `standard` are supported.'),
|
||||
size: z
|
||||
.enum(['1024x1024', '1792x1024', '1024x1792'])
|
||||
.describe(
|
||||
'The size of the requested image. Use 1024x1024 (square) as the default, 1792x1024 if the user requests a wide image, and 1024x1792 for full-body portraits. Always include this parameter in the request.',
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
getApiKey() {
|
||||
const apiKey = process.env.DALLE3_API_KEY ?? process.env.DALLE_API_KEY ?? '';
|
||||
if (!apiKey && !this.override) {
|
||||
throw new Error('Missing DALLE_API_KEY environment variable.');
|
||||
}
|
||||
return apiKey;
|
||||
}
|
||||
|
||||
replaceUnwantedChars(inputString) {
|
||||
return inputString
|
||||
.replace(/\r\n|\r|\n/g, ' ')
|
||||
.replace(/"/g, '')
|
||||
.trim();
|
||||
}
|
||||
|
||||
wrapInMarkdown(imageUrl) {
|
||||
return ``;
|
||||
}
|
||||
|
||||
returnValue(value) {
|
||||
if (this.isAgent === true && typeof value === 'string') {
|
||||
return [value, {}];
|
||||
} else if (this.isAgent === true && typeof value === 'object') {
|
||||
return [displayMessage, value];
|
||||
}
|
||||
|
||||
return value;
|
||||
}
|
||||
|
||||
async _call(data) {
|
||||
const { prompt, quality = 'standard', size = '1024x1024', style = 'vivid' } = data;
|
||||
if (!prompt) {
|
||||
throw new Error('Missing required field: prompt');
|
||||
}
|
||||
|
||||
let resp;
|
||||
try {
|
||||
resp = await this.openai.images.generate({
|
||||
model: 'dall-e-3',
|
||||
quality,
|
||||
style,
|
||||
size,
|
||||
prompt: this.replaceUnwantedChars(prompt),
|
||||
n: 1,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error('[DALL-E-3] Problem generating the image:', error);
|
||||
return this
|
||||
.returnValue(`Something went wrong when trying to generate the image. The DALL-E API may be unavailable:
|
||||
Error Message: ${error.message}`);
|
||||
}
|
||||
|
||||
if (!resp) {
|
||||
return this.returnValue(
|
||||
'Something went wrong when trying to generate the image. The DALL-E API may be unavailable',
|
||||
);
|
||||
}
|
||||
|
||||
const theImageUrl = resp.data[0].url;
|
||||
|
||||
if (!theImageUrl) {
|
||||
return this.returnValue(
|
||||
'No image URL returned from OpenAI API. There may be a problem with the API or your configuration.',
|
||||
);
|
||||
}
|
||||
|
||||
if (this.isAgent) {
|
||||
let fetchOptions = {};
|
||||
if (process.env.PROXY) {
|
||||
const proxyAgent = new ProxyAgent(process.env.PROXY);
|
||||
fetchOptions.dispatcher = proxyAgent;
|
||||
}
|
||||
const imageResponse = await fetch(theImageUrl, fetchOptions);
|
||||
const arrayBuffer = await imageResponse.arrayBuffer();
|
||||
const base64 = Buffer.from(arrayBuffer).toString('base64');
|
||||
const content = [
|
||||
{
|
||||
type: ContentTypes.IMAGE_URL,
|
||||
image_url: {
|
||||
url: `data:image/png;base64,${base64}`,
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
const response = [
|
||||
{
|
||||
type: ContentTypes.TEXT,
|
||||
text: displayMessage,
|
||||
},
|
||||
];
|
||||
return [response, { content }];
|
||||
}
|
||||
|
||||
const imageBasename = getImageBasename(theImageUrl);
|
||||
const imageExt = path.extname(imageBasename);
|
||||
|
||||
const extension = imageExt.startsWith('.') ? imageExt.slice(1) : imageExt;
|
||||
const imageName = `img-${uuidv4()}.${extension}`;
|
||||
|
||||
logger.debug('[DALL-E-3]', {
|
||||
imageName,
|
||||
imageBasename,
|
||||
imageExt,
|
||||
extension,
|
||||
theImageUrl,
|
||||
data: resp.data[0],
|
||||
});
|
||||
|
||||
try {
|
||||
const result = await this.processFileURL({
|
||||
URL: theImageUrl,
|
||||
basePath: 'images',
|
||||
userId: this.userId,
|
||||
fileName: imageName,
|
||||
fileStrategy: this.fileStrategy,
|
||||
context: FileContext.image_generation,
|
||||
});
|
||||
|
||||
if (this.returnMetadata) {
|
||||
this.result = result;
|
||||
} else {
|
||||
this.result = this.wrapInMarkdown(result.filepath);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error('Error while saving the image:', error);
|
||||
this.result = `Failed to save the image locally. ${error.message}`;
|
||||
}
|
||||
|
||||
return this.returnValue(this.result);
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = DALLE3;
|
||||
554
api/app/clients/tools/structured/FluxAPI.js
Normal file
554
api/app/clients/tools/structured/FluxAPI.js
Normal file
@ -0,0 +1,554 @@
|
||||
const { z } = require('zod');
|
||||
const axios = require('axios');
|
||||
const fetch = require('node-fetch');
|
||||
const { v4: uuidv4 } = require('uuid');
|
||||
const { Tool } = require('@langchain/core/tools');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { HttpsProxyAgent } = require('https-proxy-agent');
|
||||
const { FileContext, ContentTypes } = require('librechat-data-provider');
|
||||
|
||||
const displayMessage =
|
||||
"Flux displayed an image. All generated images are already plainly visible, so don't repeat the descriptions in detail. Do not list download links as they are available in the UI already. The user may download the images by clicking on them, but do not mention anything about downloading to the user.";
|
||||
|
||||
/**
|
||||
* FluxAPI - A tool for generating high-quality images from text prompts using the Flux API.
|
||||
* Each call generates one image. If multiple images are needed, make multiple consecutive calls with the same or varied prompts.
|
||||
*/
|
||||
class FluxAPI extends Tool {
|
||||
// Pricing constants in USD per image
|
||||
static PRICING = {
|
||||
FLUX_PRO_1_1_ULTRA: -0.06, // /v1/flux-pro-1.1-ultra
|
||||
FLUX_PRO_1_1: -0.04, // /v1/flux-pro-1.1
|
||||
FLUX_PRO: -0.05, // /v1/flux-pro
|
||||
FLUX_DEV: -0.025, // /v1/flux-dev
|
||||
FLUX_PRO_FINETUNED: -0.06, // /v1/flux-pro-finetuned
|
||||
FLUX_PRO_1_1_ULTRA_FINETUNED: -0.07, // /v1/flux-pro-1.1-ultra-finetuned
|
||||
};
|
||||
|
||||
constructor(fields = {}) {
|
||||
super();
|
||||
|
||||
/** @type {boolean} Used to initialize the Tool without necessary variables. */
|
||||
this.override = fields.override ?? false;
|
||||
|
||||
this.userId = fields.userId;
|
||||
this.fileStrategy = fields.fileStrategy;
|
||||
|
||||
/** @type {boolean} **/
|
||||
this.isAgent = fields.isAgent;
|
||||
this.returnMetadata = fields.returnMetadata ?? false;
|
||||
|
||||
if (fields.processFileURL) {
|
||||
/** @type {processFileURL} Necessary for output to contain all image metadata. */
|
||||
this.processFileURL = fields.processFileURL.bind(this);
|
||||
}
|
||||
|
||||
this.apiKey = fields.FLUX_API_KEY || this.getApiKey();
|
||||
|
||||
this.name = 'flux';
|
||||
this.description =
|
||||
'Use Flux to generate images from text descriptions. This tool can generate images and list available finetunes. Each generate call creates one image. For multiple images, make multiple consecutive calls.';
|
||||
|
||||
this.description_for_model = `// Transform any image description into a detailed, high-quality prompt. Never submit a prompt under 3 sentences. Follow these core rules:
|
||||
// 1. ALWAYS enhance basic prompts into 5-10 detailed sentences (e.g., "a cat" becomes: "A close-up photo of a sleek Siamese cat with piercing blue eyes. The cat sits elegantly on a vintage leather armchair, its tail curled gracefully around its paws. Warm afternoon sunlight streams through a nearby window, casting gentle shadows across its face and highlighting the subtle variations in its cream and chocolate-point fur. The background is softly blurred, creating a shallow depth of field that draws attention to the cat's expressive features. The overall composition has a peaceful, contemplative mood with a professional photography style.")
|
||||
// 2. Each prompt MUST be 3-6 descriptive sentences minimum, focusing on visual elements: lighting, composition, mood, and style
|
||||
// Use action: 'list_finetunes' to see available custom models. When using finetunes, use endpoint: '/v1/flux-pro-finetuned' (default) or '/v1/flux-pro-1.1-ultra-finetuned' for higher quality and aspect ratio.`;
|
||||
|
||||
// Add base URL from environment variable with fallback
|
||||
this.baseUrl = process.env.FLUX_API_BASE_URL || 'https://api.us1.bfl.ai';
|
||||
|
||||
// Define the schema for structured input
|
||||
this.schema = z.object({
|
||||
action: z
|
||||
.enum(['generate', 'list_finetunes', 'generate_finetuned'])
|
||||
.default('generate')
|
||||
.describe(
|
||||
'Action to perform: "generate" for image generation, "generate_finetuned" for finetuned model generation, "list_finetunes" to get available custom models',
|
||||
),
|
||||
prompt: z
|
||||
.string()
|
||||
.optional()
|
||||
.describe(
|
||||
'Text prompt for image generation. Required when action is "generate". Not used for list_finetunes.',
|
||||
),
|
||||
width: z
|
||||
.number()
|
||||
.optional()
|
||||
.describe(
|
||||
'Width of the generated image in pixels. Must be a multiple of 32. Default is 1024.',
|
||||
),
|
||||
height: z
|
||||
.number()
|
||||
.optional()
|
||||
.describe(
|
||||
'Height of the generated image in pixels. Must be a multiple of 32. Default is 768.',
|
||||
),
|
||||
prompt_upsampling: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.default(false)
|
||||
.describe('Whether to perform upsampling on the prompt.'),
|
||||
steps: z
|
||||
.number()
|
||||
.int()
|
||||
.optional()
|
||||
.describe('Number of steps to run the model for, a number from 1 to 50. Default is 40.'),
|
||||
seed: z.number().optional().describe('Optional seed for reproducibility.'),
|
||||
safety_tolerance: z
|
||||
.number()
|
||||
.optional()
|
||||
.default(6)
|
||||
.describe(
|
||||
'Tolerance level for input and output moderation. Between 0 and 6, 0 being most strict, 6 being least strict.',
|
||||
),
|
||||
endpoint: z
|
||||
.enum([
|
||||
'/v1/flux-pro-1.1',
|
||||
'/v1/flux-pro',
|
||||
'/v1/flux-dev',
|
||||
'/v1/flux-pro-1.1-ultra',
|
||||
'/v1/flux-pro-finetuned',
|
||||
'/v1/flux-pro-1.1-ultra-finetuned',
|
||||
])
|
||||
.optional()
|
||||
.default('/v1/flux-pro-1.1')
|
||||
.describe('Endpoint to use for image generation.'),
|
||||
raw: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.default(false)
|
||||
.describe(
|
||||
'Generate less processed, more natural-looking images. Only works for /v1/flux-pro-1.1-ultra.',
|
||||
),
|
||||
finetune_id: z.string().optional().describe('ID of the finetuned model to use'),
|
||||
finetune_strength: z
|
||||
.number()
|
||||
.optional()
|
||||
.default(1.1)
|
||||
.describe('Strength of the finetuning effect (typically between 0.1 and 1.2)'),
|
||||
guidance: z.number().optional().default(2.5).describe('Guidance scale for finetuned models'),
|
||||
aspect_ratio: z
|
||||
.string()
|
||||
.optional()
|
||||
.default('16:9')
|
||||
.describe('Aspect ratio for ultra models (e.g., "16:9")'),
|
||||
});
|
||||
}
|
||||
|
||||
getAxiosConfig() {
|
||||
const config = {};
|
||||
if (process.env.PROXY) {
|
||||
config.httpsAgent = new HttpsProxyAgent(process.env.PROXY);
|
||||
}
|
||||
return config;
|
||||
}
|
||||
|
||||
/** @param {Object|string} value */
|
||||
getDetails(value) {
|
||||
if (typeof value === 'string') {
|
||||
return value;
|
||||
}
|
||||
return JSON.stringify(value, null, 2);
|
||||
}
|
||||
|
||||
getApiKey() {
|
||||
const apiKey = process.env.FLUX_API_KEY || '';
|
||||
if (!apiKey && !this.override) {
|
||||
throw new Error('Missing FLUX_API_KEY environment variable.');
|
||||
}
|
||||
return apiKey;
|
||||
}
|
||||
|
||||
wrapInMarkdown(imageUrl) {
|
||||
const serverDomain = process.env.DOMAIN_SERVER || 'http://localhost:3080';
|
||||
return ``;
|
||||
}
|
||||
|
||||
returnValue(value) {
|
||||
if (this.isAgent === true && typeof value === 'string') {
|
||||
return [value, {}];
|
||||
} else if (this.isAgent === true && typeof value === 'object') {
|
||||
if (Array.isArray(value)) {
|
||||
return value;
|
||||
}
|
||||
return [displayMessage, value];
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
async _call(data) {
|
||||
const { action = 'generate', ...imageData } = data;
|
||||
|
||||
// Use provided API key for this request if available, otherwise use default
|
||||
const requestApiKey = this.apiKey || this.getApiKey();
|
||||
|
||||
// Handle list_finetunes action
|
||||
if (action === 'list_finetunes') {
|
||||
return this.getMyFinetunes(requestApiKey);
|
||||
}
|
||||
|
||||
// Handle finetuned generation
|
||||
if (action === 'generate_finetuned') {
|
||||
return this.generateFinetunedImage(imageData, requestApiKey);
|
||||
}
|
||||
|
||||
// For generate action, ensure prompt is provided
|
||||
if (!imageData.prompt) {
|
||||
throw new Error('Missing required field: prompt');
|
||||
}
|
||||
|
||||
let payload = {
|
||||
prompt: imageData.prompt,
|
||||
prompt_upsampling: imageData.prompt_upsampling || false,
|
||||
safety_tolerance: imageData.safety_tolerance || 6,
|
||||
output_format: imageData.output_format || 'png',
|
||||
};
|
||||
|
||||
// Add optional parameters if provided
|
||||
if (imageData.width) {
|
||||
payload.width = imageData.width;
|
||||
}
|
||||
if (imageData.height) {
|
||||
payload.height = imageData.height;
|
||||
}
|
||||
if (imageData.steps) {
|
||||
payload.steps = imageData.steps;
|
||||
}
|
||||
if (imageData.seed !== undefined) {
|
||||
payload.seed = imageData.seed;
|
||||
}
|
||||
if (imageData.raw) {
|
||||
payload.raw = imageData.raw;
|
||||
}
|
||||
|
||||
const generateUrl = `${this.baseUrl}${imageData.endpoint || '/v1/flux-pro'}`;
|
||||
const resultUrl = `${this.baseUrl}/v1/get_result`;
|
||||
|
||||
logger.debug('[FluxAPI] Generating image with payload:', payload);
|
||||
logger.debug('[FluxAPI] Using endpoint:', generateUrl);
|
||||
|
||||
let taskResponse;
|
||||
try {
|
||||
taskResponse = await axios.post(generateUrl, payload, {
|
||||
headers: {
|
||||
'x-key': requestApiKey,
|
||||
'Content-Type': 'application/json',
|
||||
Accept: 'application/json',
|
||||
},
|
||||
...this.getAxiosConfig(),
|
||||
});
|
||||
} catch (error) {
|
||||
const details = this.getDetails(error?.response?.data || error.message);
|
||||
logger.error('[FluxAPI] Error while submitting task:', details);
|
||||
|
||||
return this.returnValue(
|
||||
`Something went wrong when trying to generate the image. The Flux API may be unavailable:
|
||||
Error Message: ${details}`,
|
||||
);
|
||||
}
|
||||
|
||||
const taskId = taskResponse.data.id;
|
||||
|
||||
// Polling for the result
|
||||
let status = 'Pending';
|
||||
let resultData = null;
|
||||
while (status !== 'Ready' && status !== 'Error') {
|
||||
try {
|
||||
// Wait 2 seconds between polls
|
||||
await new Promise((resolve) => setTimeout(resolve, 2000));
|
||||
const resultResponse = await axios.get(resultUrl, {
|
||||
headers: {
|
||||
'x-key': requestApiKey,
|
||||
Accept: 'application/json',
|
||||
},
|
||||
params: { id: taskId },
|
||||
...this.getAxiosConfig(),
|
||||
});
|
||||
status = resultResponse.data.status;
|
||||
|
||||
if (status === 'Ready') {
|
||||
resultData = resultResponse.data.result;
|
||||
break;
|
||||
} else if (status === 'Error') {
|
||||
logger.error('[FluxAPI] Error in task:', resultResponse.data);
|
||||
return this.returnValue('An error occurred during image generation.');
|
||||
}
|
||||
} catch (error) {
|
||||
const details = this.getDetails(error?.response?.data || error.message);
|
||||
logger.error('[FluxAPI] Error while getting result:', details);
|
||||
return this.returnValue('An error occurred while retrieving the image.');
|
||||
}
|
||||
}
|
||||
|
||||
// If no result data
|
||||
if (!resultData || !resultData.sample) {
|
||||
logger.error('[FluxAPI] No image data received from API. Response:', resultData);
|
||||
return this.returnValue('No image data received from Flux API.');
|
||||
}
|
||||
|
||||
// Try saving the image locally
|
||||
const imageUrl = resultData.sample;
|
||||
const imageName = `img-${uuidv4()}.png`;
|
||||
|
||||
if (this.isAgent) {
|
||||
try {
|
||||
// Fetch the image and convert to base64
|
||||
const fetchOptions = {};
|
||||
if (process.env.PROXY) {
|
||||
fetchOptions.agent = new HttpsProxyAgent(process.env.PROXY);
|
||||
}
|
||||
const imageResponse = await fetch(imageUrl, fetchOptions);
|
||||
const arrayBuffer = await imageResponse.arrayBuffer();
|
||||
const base64 = Buffer.from(arrayBuffer).toString('base64');
|
||||
const content = [
|
||||
{
|
||||
type: ContentTypes.IMAGE_URL,
|
||||
image_url: {
|
||||
url: `data:image/png;base64,${base64}`,
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
const response = [
|
||||
{
|
||||
type: ContentTypes.TEXT,
|
||||
text: displayMessage,
|
||||
},
|
||||
];
|
||||
return [response, { content }];
|
||||
} catch (error) {
|
||||
logger.error('Error processing image for agent:', error);
|
||||
return this.returnValue(`Failed to process the image. ${error.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
try {
|
||||
logger.debug('[FluxAPI] Saving image:', imageUrl);
|
||||
const result = await this.processFileURL({
|
||||
fileStrategy: this.fileStrategy,
|
||||
userId: this.userId,
|
||||
URL: imageUrl,
|
||||
fileName: imageName,
|
||||
basePath: 'images',
|
||||
context: FileContext.image_generation,
|
||||
});
|
||||
|
||||
logger.debug('[FluxAPI] Image saved to path:', result.filepath);
|
||||
|
||||
// Calculate cost based on endpoint
|
||||
/**
|
||||
* TODO: Cost handling
|
||||
const endpoint = imageData.endpoint || '/v1/flux-pro';
|
||||
const endpointKey = Object.entries(FluxAPI.PRICING).find(([key, _]) =>
|
||||
endpoint.includes(key.toLowerCase().replace(/_/g, '-')),
|
||||
)?.[0];
|
||||
const cost = FluxAPI.PRICING[endpointKey] || 0;
|
||||
*/
|
||||
this.result = this.returnMetadata ? result : this.wrapInMarkdown(result.filepath);
|
||||
return this.returnValue(this.result);
|
||||
} catch (error) {
|
||||
const details = this.getDetails(error?.message ?? 'No additional error details.');
|
||||
logger.error('Error while saving the image:', details);
|
||||
return this.returnValue(`Failed to save the image locally. ${details}`);
|
||||
}
|
||||
}
|
||||
|
||||
async getMyFinetunes(apiKey = null) {
|
||||
const finetunesUrl = `${this.baseUrl}/v1/my_finetunes`;
|
||||
const detailsUrl = `${this.baseUrl}/v1/finetune_details`;
|
||||
|
||||
try {
|
||||
const headers = {
|
||||
'x-key': apiKey || this.getApiKey(),
|
||||
'Content-Type': 'application/json',
|
||||
Accept: 'application/json',
|
||||
};
|
||||
|
||||
// Get list of finetunes
|
||||
const response = await axios.get(finetunesUrl, {
|
||||
headers,
|
||||
...this.getAxiosConfig(),
|
||||
});
|
||||
const finetunes = response.data.finetunes;
|
||||
|
||||
// Fetch details for each finetune
|
||||
const finetuneDetails = await Promise.all(
|
||||
finetunes.map(async (finetuneId) => {
|
||||
try {
|
||||
const detailResponse = await axios.get(`${detailsUrl}?finetune_id=${finetuneId}`, {
|
||||
headers,
|
||||
...this.getAxiosConfig(),
|
||||
});
|
||||
return {
|
||||
id: finetuneId,
|
||||
...detailResponse.data,
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error(`[FluxAPI] Error fetching details for finetune ${finetuneId}:`, error);
|
||||
return {
|
||||
id: finetuneId,
|
||||
error: 'Failed to fetch details',
|
||||
};
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
if (this.isAgent) {
|
||||
const formattedDetails = JSON.stringify(finetuneDetails, null, 2);
|
||||
return [`Here are the available finetunes:\n${formattedDetails}`, null];
|
||||
}
|
||||
return JSON.stringify(finetuneDetails);
|
||||
} catch (error) {
|
||||
const details = this.getDetails(error?.response?.data || error.message);
|
||||
logger.error('[FluxAPI] Error while getting finetunes:', details);
|
||||
const errorMsg = `Failed to get finetunes: ${details}`;
|
||||
return this.isAgent ? this.returnValue([errorMsg, {}]) : new Error(errorMsg);
|
||||
}
|
||||
}
|
||||
|
||||
async generateFinetunedImage(imageData, requestApiKey) {
|
||||
if (!imageData.prompt) {
|
||||
throw new Error('Missing required field: prompt');
|
||||
}
|
||||
|
||||
if (!imageData.finetune_id) {
|
||||
throw new Error(
|
||||
'Missing required field: finetune_id for finetuned generation. Please supply a finetune_id!',
|
||||
);
|
||||
}
|
||||
|
||||
// Validate endpoint is appropriate for finetuned generation
|
||||
const validFinetunedEndpoints = ['/v1/flux-pro-finetuned', '/v1/flux-pro-1.1-ultra-finetuned'];
|
||||
const endpoint = imageData.endpoint || '/v1/flux-pro-finetuned';
|
||||
|
||||
if (!validFinetunedEndpoints.includes(endpoint)) {
|
||||
throw new Error(
|
||||
`Invalid endpoint for finetuned generation. Must be one of: ${validFinetunedEndpoints.join(', ')}`,
|
||||
);
|
||||
}
|
||||
|
||||
let payload = {
|
||||
prompt: imageData.prompt,
|
||||
prompt_upsampling: imageData.prompt_upsampling || false,
|
||||
safety_tolerance: imageData.safety_tolerance || 6,
|
||||
output_format: imageData.output_format || 'png',
|
||||
finetune_id: imageData.finetune_id,
|
||||
finetune_strength: imageData.finetune_strength || 1.0,
|
||||
guidance: imageData.guidance || 2.5,
|
||||
};
|
||||
|
||||
// Add optional parameters if provided
|
||||
if (imageData.width) {
|
||||
payload.width = imageData.width;
|
||||
}
|
||||
if (imageData.height) {
|
||||
payload.height = imageData.height;
|
||||
}
|
||||
if (imageData.steps) {
|
||||
payload.steps = imageData.steps;
|
||||
}
|
||||
if (imageData.seed !== undefined) {
|
||||
payload.seed = imageData.seed;
|
||||
}
|
||||
if (imageData.raw) {
|
||||
payload.raw = imageData.raw;
|
||||
}
|
||||
|
||||
const generateUrl = `${this.baseUrl}${endpoint}`;
|
||||
const resultUrl = `${this.baseUrl}/v1/get_result`;
|
||||
|
||||
logger.debug('[FluxAPI] Generating finetuned image with payload:', payload);
|
||||
logger.debug('[FluxAPI] Using endpoint:', generateUrl);
|
||||
|
||||
let taskResponse;
|
||||
try {
|
||||
taskResponse = await axios.post(generateUrl, payload, {
|
||||
headers: {
|
||||
'x-key': requestApiKey,
|
||||
'Content-Type': 'application/json',
|
||||
Accept: 'application/json',
|
||||
},
|
||||
...this.getAxiosConfig(),
|
||||
});
|
||||
} catch (error) {
|
||||
const details = this.getDetails(error?.response?.data || error.message);
|
||||
logger.error('[FluxAPI] Error while submitting finetuned task:', details);
|
||||
return this.returnValue(
|
||||
`Something went wrong when trying to generate the finetuned image. The Flux API may be unavailable:
|
||||
Error Message: ${details}`,
|
||||
);
|
||||
}
|
||||
|
||||
const taskId = taskResponse.data.id;
|
||||
|
||||
// Polling for the result
|
||||
let status = 'Pending';
|
||||
let resultData = null;
|
||||
while (status !== 'Ready' && status !== 'Error') {
|
||||
try {
|
||||
// Wait 2 seconds between polls
|
||||
await new Promise((resolve) => setTimeout(resolve, 2000));
|
||||
const resultResponse = await axios.get(resultUrl, {
|
||||
headers: {
|
||||
'x-key': requestApiKey,
|
||||
Accept: 'application/json',
|
||||
},
|
||||
params: { id: taskId },
|
||||
...this.getAxiosConfig(),
|
||||
});
|
||||
status = resultResponse.data.status;
|
||||
|
||||
if (status === 'Ready') {
|
||||
resultData = resultResponse.data.result;
|
||||
break;
|
||||
} else if (status === 'Error') {
|
||||
logger.error('[FluxAPI] Error in finetuned task:', resultResponse.data);
|
||||
return this.returnValue('An error occurred during finetuned image generation.');
|
||||
}
|
||||
} catch (error) {
|
||||
const details = this.getDetails(error?.response?.data || error.message);
|
||||
logger.error('[FluxAPI] Error while getting finetuned result:', details);
|
||||
return this.returnValue('An error occurred while retrieving the finetuned image.');
|
||||
}
|
||||
}
|
||||
|
||||
// If no result data
|
||||
if (!resultData || !resultData.sample) {
|
||||
logger.error('[FluxAPI] No image data received from API. Response:', resultData);
|
||||
return this.returnValue('No image data received from Flux API.');
|
||||
}
|
||||
|
||||
// Try saving the image locally
|
||||
const imageUrl = resultData.sample;
|
||||
const imageName = `img-${uuidv4()}.png`;
|
||||
|
||||
try {
|
||||
logger.debug('[FluxAPI] Saving finetuned image:', imageUrl);
|
||||
const result = await this.processFileURL({
|
||||
fileStrategy: this.fileStrategy,
|
||||
userId: this.userId,
|
||||
URL: imageUrl,
|
||||
fileName: imageName,
|
||||
basePath: 'images',
|
||||
context: FileContext.image_generation,
|
||||
});
|
||||
|
||||
logger.debug('[FluxAPI] Finetuned image saved to path:', result.filepath);
|
||||
|
||||
// Calculate cost based on endpoint
|
||||
const endpointKey = endpoint.includes('ultra')
|
||||
? 'FLUX_PRO_1_1_ULTRA_FINETUNED'
|
||||
: 'FLUX_PRO_FINETUNED';
|
||||
const cost = FluxAPI.PRICING[endpointKey] || 0;
|
||||
// Return the result based on returnMetadata flag
|
||||
this.result = this.returnMetadata ? result : this.wrapInMarkdown(result.filepath);
|
||||
return this.returnValue(this.result);
|
||||
} catch (error) {
|
||||
const details = this.getDetails(error?.message ?? 'No additional error details.');
|
||||
logger.error('Error while saving the finetuned image:', details);
|
||||
return this.returnValue(`Failed to save the finetuned image locally. ${details}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = FluxAPI;
|
||||
595
api/app/clients/tools/structured/GeminiImageGen.js
Normal file
595
api/app/clients/tools/structured/GeminiImageGen.js
Normal file
@ -0,0 +1,595 @@
|
||||
const fs = require('fs');
|
||||
const path = require('path');
|
||||
const sharp = require('sharp');
|
||||
const { v4 } = require('uuid');
|
||||
const { ProxyAgent } = require('undici');
|
||||
const { GoogleGenAI } = require('@google/genai');
|
||||
const { tool } = require('@langchain/core/tools');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const {
|
||||
FileContext,
|
||||
ContentTypes,
|
||||
FileSources,
|
||||
EImageOutputType,
|
||||
} = require('librechat-data-provider');
|
||||
const {
|
||||
geminiToolkit,
|
||||
loadServiceKey,
|
||||
getBalanceConfig,
|
||||
getTransactionsConfig,
|
||||
} = require('@librechat/api');
|
||||
const { getStrategyFunctions } = require('~/server/services/Files/strategies');
|
||||
const { spendTokens } = require('~/models/spendTokens');
|
||||
const { getFiles } = require('~/models/File');
|
||||
|
||||
/**
|
||||
* Configure proxy support for Google APIs
|
||||
* This wraps globalThis.fetch to add a proxy dispatcher only for googleapis.com URLs
|
||||
* This is necessary because @google/genai SDK doesn't support custom fetch or httpOptions.dispatcher
|
||||
*/
|
||||
if (process.env.PROXY) {
|
||||
const originalFetch = globalThis.fetch;
|
||||
const proxyAgent = new ProxyAgent(process.env.PROXY);
|
||||
|
||||
globalThis.fetch = function (url, options = {}) {
|
||||
const urlString = url.toString();
|
||||
if (urlString.includes('googleapis.com')) {
|
||||
options = { ...options, dispatcher: proxyAgent };
|
||||
}
|
||||
return originalFetch.call(this, url, options);
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the default service key file path (consistent with main Google endpoint)
|
||||
* @returns {string} - The default path to the service key file
|
||||
*/
|
||||
function getDefaultServiceKeyPath() {
|
||||
return (
|
||||
process.env.GOOGLE_SERVICE_KEY_FILE || path.join(process.cwd(), 'api', 'data', 'auth.json')
|
||||
);
|
||||
}
|
||||
|
||||
const displayMessage =
|
||||
"Gemini displayed an image. All generated images are already plainly visible, so don't repeat the descriptions in detail. Do not list download links as they are available in the UI already. The user may download the images by clicking on them, but do not mention anything about downloading to the user.";
|
||||
|
||||
/**
|
||||
* Replaces unwanted characters from the input string
|
||||
* @param {string} inputString - The input string to process
|
||||
* @returns {string} - The processed string
|
||||
*/
|
||||
function replaceUnwantedChars(inputString) {
|
||||
return inputString?.replace(/[^\w\s\-_.,!?()]/g, '') || '';
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate and sanitize image format
|
||||
* @param {string} format - The format to validate
|
||||
* @returns {string} - Safe format
|
||||
*/
|
||||
function getSafeFormat(format) {
|
||||
const allowedFormats = ['png', 'jpg', 'jpeg', 'webp', 'gif'];
|
||||
return allowedFormats.includes(format?.toLowerCase()) ? format.toLowerCase() : 'png';
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert image buffer to target format if needed
|
||||
* @param {Buffer} inputBuffer - The input image buffer
|
||||
* @param {string} targetFormat - The target format (png, jpeg, webp)
|
||||
* @returns {Promise<{buffer: Buffer, format: string}>} - Converted buffer and format
|
||||
*/
|
||||
async function convertImageFormat(inputBuffer, targetFormat) {
|
||||
const metadata = await sharp(inputBuffer).metadata();
|
||||
const currentFormat = metadata.format;
|
||||
|
||||
// Normalize format names (jpg -> jpeg)
|
||||
const normalizedTarget = targetFormat === 'jpg' ? 'jpeg' : targetFormat.toLowerCase();
|
||||
const normalizedCurrent = currentFormat === 'jpg' ? 'jpeg' : currentFormat;
|
||||
|
||||
// If already in target format, return as-is
|
||||
if (normalizedCurrent === normalizedTarget) {
|
||||
return { buffer: inputBuffer, format: normalizedTarget };
|
||||
}
|
||||
|
||||
// Convert to target format
|
||||
const convertedBuffer = await sharp(inputBuffer).toFormat(normalizedTarget).toBuffer();
|
||||
return { buffer: convertedBuffer, format: normalizedTarget };
|
||||
}
|
||||
|
||||
/**
|
||||
* Initialize Gemini client (supports both Gemini API and Vertex AI)
|
||||
* Priority: API key (from options, resolved by loadAuthValues) > Vertex AI service account
|
||||
* @param {Object} options - Initialization options
|
||||
* @param {string} [options.GEMINI_API_KEY] - Gemini API key (resolved by loadAuthValues)
|
||||
* @param {string} [options.GOOGLE_KEY] - Google API key (resolved by loadAuthValues)
|
||||
* @returns {Promise<GoogleGenAI>} - The initialized client
|
||||
*/
|
||||
async function initializeGeminiClient(options = {}) {
|
||||
const geminiKey = options.GEMINI_API_KEY;
|
||||
if (geminiKey) {
|
||||
logger.debug('[GeminiImageGen] Using Gemini API with GEMINI_API_KEY');
|
||||
return new GoogleGenAI({ apiKey: geminiKey });
|
||||
}
|
||||
|
||||
const googleKey = options.GOOGLE_KEY;
|
||||
if (googleKey) {
|
||||
logger.debug('[GeminiImageGen] Using Gemini API with GOOGLE_KEY');
|
||||
return new GoogleGenAI({ apiKey: googleKey });
|
||||
}
|
||||
|
||||
// Fall back to Vertex AI with service account
|
||||
logger.debug('[GeminiImageGen] Using Vertex AI with service account');
|
||||
const credentialsPath = getDefaultServiceKeyPath();
|
||||
|
||||
// Use loadServiceKey for consistent loading (supports file paths, JSON strings, base64)
|
||||
const serviceKey = await loadServiceKey(credentialsPath);
|
||||
|
||||
if (!serviceKey || !serviceKey.project_id) {
|
||||
throw new Error(
|
||||
'Gemini Image Generation requires one of: user-provided API key, GEMINI_API_KEY or GOOGLE_KEY env var, or a valid Google service account. ' +
|
||||
`Service account file not found or invalid at: ${credentialsPath}`,
|
||||
);
|
||||
}
|
||||
|
||||
// Set GOOGLE_APPLICATION_CREDENTIALS for any Google Cloud SDK dependencies
|
||||
try {
|
||||
await fs.promises.access(credentialsPath);
|
||||
process.env.GOOGLE_APPLICATION_CREDENTIALS = credentialsPath;
|
||||
} catch {
|
||||
// File doesn't exist, skip setting env var
|
||||
}
|
||||
|
||||
return new GoogleGenAI({
|
||||
vertexai: true,
|
||||
project: serviceKey.project_id,
|
||||
location: process.env.GOOGLE_LOC || process.env.GOOGLE_CLOUD_LOCATION || 'global',
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Save image to local filesystem
|
||||
* @param {string} base64Data - Base64 encoded image data
|
||||
* @param {string} format - Image format
|
||||
* @param {string} userId - User ID
|
||||
* @returns {Promise<string>} - The relative URL
|
||||
*/
|
||||
async function saveImageLocally(base64Data, format, userId) {
|
||||
const safeFormat = getSafeFormat(format);
|
||||
const safeUserId = userId ? path.basename(userId) : 'default';
|
||||
const imageName = `gemini-img-${v4()}.${safeFormat}`;
|
||||
const userDir = path.join(process.cwd(), 'client/public/images', safeUserId);
|
||||
|
||||
await fs.promises.mkdir(userDir, { recursive: true });
|
||||
|
||||
const filePath = path.join(userDir, imageName);
|
||||
await fs.promises.writeFile(filePath, Buffer.from(base64Data, 'base64'));
|
||||
|
||||
logger.debug('[GeminiImageGen] Image saved locally to:', filePath);
|
||||
return `/images/${safeUserId}/${imageName}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Save image to cloud storage
|
||||
* @param {Object} params - Parameters
|
||||
* @returns {Promise<string|null>} - The storage URL or null
|
||||
*/
|
||||
async function saveToCloudStorage({ base64Data, format, processFileURL, fileStrategy, userId }) {
|
||||
if (!processFileURL || !fileStrategy || !userId) {
|
||||
return null;
|
||||
}
|
||||
|
||||
try {
|
||||
const safeFormat = getSafeFormat(format);
|
||||
const safeUserId = path.basename(userId);
|
||||
const dataURL = `data:image/${safeFormat};base64,${base64Data}`;
|
||||
const imageName = `gemini-img-${v4()}.${safeFormat}`;
|
||||
|
||||
const result = await processFileURL({
|
||||
URL: dataURL,
|
||||
basePath: 'images',
|
||||
userId: safeUserId,
|
||||
fileName: imageName,
|
||||
fileStrategy,
|
||||
context: FileContext.image_generation,
|
||||
});
|
||||
|
||||
return result.filepath;
|
||||
} catch (error) {
|
||||
logger.error('[GeminiImageGen] Error saving to cloud storage:', error);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert image files to Gemini inline data format
|
||||
* @param {Object} params - Parameters
|
||||
* @returns {Promise<Array>} - Array of inline data objects
|
||||
*/
|
||||
async function convertImagesToInlineData({ imageFiles, image_ids, req, fileStrategy }) {
|
||||
if (!image_ids || image_ids.length === 0) {
|
||||
return [];
|
||||
}
|
||||
|
||||
const streamMethods = {};
|
||||
const requestFilesMap = Object.fromEntries(imageFiles.map((f) => [f.file_id, { ...f }]));
|
||||
const orderedFiles = new Array(image_ids.length);
|
||||
const idsToFetch = [];
|
||||
const indexOfMissing = Object.create(null);
|
||||
|
||||
for (let i = 0; i < image_ids.length; i++) {
|
||||
const id = image_ids[i];
|
||||
const file = requestFilesMap[id];
|
||||
if (file) {
|
||||
orderedFiles[i] = file;
|
||||
} else {
|
||||
idsToFetch.push(id);
|
||||
indexOfMissing[id] = i;
|
||||
}
|
||||
}
|
||||
|
||||
if (idsToFetch.length && req?.user?.id) {
|
||||
const fetchedFiles = await getFiles(
|
||||
{
|
||||
user: req.user.id,
|
||||
file_id: { $in: idsToFetch },
|
||||
height: { $exists: true },
|
||||
width: { $exists: true },
|
||||
},
|
||||
{},
|
||||
{},
|
||||
);
|
||||
|
||||
for (const file of fetchedFiles) {
|
||||
requestFilesMap[file.file_id] = file;
|
||||
orderedFiles[indexOfMissing[file.file_id]] = file;
|
||||
}
|
||||
}
|
||||
|
||||
const inlineDataArray = [];
|
||||
for (const imageFile of orderedFiles) {
|
||||
if (!imageFile) continue;
|
||||
|
||||
try {
|
||||
const source = imageFile.source || fileStrategy;
|
||||
if (!source) continue;
|
||||
|
||||
let getDownloadStream = streamMethods[source];
|
||||
if (!getDownloadStream) {
|
||||
({ getDownloadStream } = getStrategyFunctions(source));
|
||||
streamMethods[source] = getDownloadStream;
|
||||
}
|
||||
if (!getDownloadStream) continue;
|
||||
|
||||
const stream = await getDownloadStream(req, imageFile.filepath);
|
||||
if (!stream) continue;
|
||||
|
||||
const chunks = [];
|
||||
for await (const chunk of stream) {
|
||||
chunks.push(chunk);
|
||||
}
|
||||
const buffer = Buffer.concat(chunks);
|
||||
const base64Data = buffer.toString('base64');
|
||||
const mimeType = imageFile.type || 'image/png';
|
||||
|
||||
inlineDataArray.push({
|
||||
inlineData: { mimeType, data: base64Data },
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error('[GeminiImageGen] Error processing image:', imageFile.file_id, error);
|
||||
}
|
||||
}
|
||||
|
||||
return inlineDataArray;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check for safety blocks in API response
|
||||
* @param {Object} response - The API response
|
||||
* @returns {Object|null} - Safety block info or null
|
||||
*/
|
||||
function checkForSafetyBlock(response) {
|
||||
if (!response?.candidates?.length) {
|
||||
return { reason: 'NO_CANDIDATES', message: 'No candidates returned' };
|
||||
}
|
||||
|
||||
const candidate = response.candidates[0];
|
||||
const finishReason = candidate.finishReason;
|
||||
|
||||
if (finishReason === 'SAFETY' || finishReason === 'PROHIBITED_CONTENT') {
|
||||
return { reason: finishReason, message: 'Content blocked by safety filters' };
|
||||
}
|
||||
|
||||
if (finishReason === 'RECITATION') {
|
||||
return { reason: finishReason, message: 'Content blocked due to recitation concerns' };
|
||||
}
|
||||
|
||||
if (candidate.safetyRatings) {
|
||||
for (const rating of candidate.safetyRatings) {
|
||||
if (rating.probability === 'HIGH' || rating.blocked === true) {
|
||||
return {
|
||||
reason: 'SAFETY_RATING',
|
||||
message: `Blocked due to ${rating.category}`,
|
||||
category: rating.category,
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Record token usage for balance tracking
|
||||
* @param {Object} params - Parameters
|
||||
* @param {Object} params.usageMetadata - The usage metadata from API response
|
||||
* @param {Object} params.req - The request object
|
||||
* @param {string} params.userId - The user ID
|
||||
* @param {string} params.conversationId - The conversation ID
|
||||
* @param {string} params.model - The model name
|
||||
*/
|
||||
async function recordTokenUsage({ usageMetadata, req, userId, conversationId, model }) {
|
||||
if (!usageMetadata) {
|
||||
logger.debug('[GeminiImageGen] No usage metadata available for balance tracking');
|
||||
return;
|
||||
}
|
||||
|
||||
const appConfig = req?.config;
|
||||
const balance = getBalanceConfig(appConfig);
|
||||
const transactions = getTransactionsConfig(appConfig);
|
||||
|
||||
// Skip if neither balance nor transactions are enabled
|
||||
if (!balance?.enabled && transactions?.enabled === false) {
|
||||
return;
|
||||
}
|
||||
|
||||
const promptTokens = usageMetadata.prompt_token_count || usageMetadata.promptTokenCount || 0;
|
||||
const completionTokens =
|
||||
usageMetadata.candidates_token_count || usageMetadata.candidatesTokenCount || 0;
|
||||
|
||||
if (promptTokens === 0 && completionTokens === 0) {
|
||||
logger.debug('[GeminiImageGen] No tokens to record');
|
||||
return;
|
||||
}
|
||||
|
||||
logger.debug('[GeminiImageGen] Recording token usage:', {
|
||||
promptTokens,
|
||||
completionTokens,
|
||||
model,
|
||||
conversationId,
|
||||
});
|
||||
|
||||
try {
|
||||
await spendTokens(
|
||||
{
|
||||
user: userId,
|
||||
model,
|
||||
conversationId,
|
||||
context: 'image_generation',
|
||||
balance,
|
||||
transactions,
|
||||
},
|
||||
{
|
||||
promptTokens,
|
||||
completionTokens,
|
||||
},
|
||||
);
|
||||
} catch (error) {
|
||||
logger.error('[GeminiImageGen] Error recording token usage:', error);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates Gemini Image Generation tool
|
||||
* @param {Object} fields - Configuration fields
|
||||
* @returns {ReturnType<tool>} - The image generation tool
|
||||
*/
|
||||
function createGeminiImageTool(fields = {}) {
|
||||
const override = fields.override ?? false;
|
||||
|
||||
if (!override && !fields.isAgent) {
|
||||
throw new Error('This tool is only available for agents.');
|
||||
}
|
||||
|
||||
// Skip validation during tool creation - validation happens at runtime in initializeGeminiClient
|
||||
// This allows the tool to be added to agents when using Vertex AI without requiring API keys
|
||||
// The actual credentials check happens when the tool is invoked
|
||||
|
||||
const {
|
||||
req,
|
||||
imageFiles = [],
|
||||
processFileURL,
|
||||
userId,
|
||||
fileStrategy,
|
||||
GEMINI_API_KEY,
|
||||
GOOGLE_KEY,
|
||||
// GEMINI_VERTEX_ENABLED is used for auth validation only (not used in code)
|
||||
// When set as env var, it signals Vertex AI is configured and bypasses API key requirement
|
||||
} = fields;
|
||||
|
||||
const imageOutputType = fields.imageOutputType || EImageOutputType.PNG;
|
||||
|
||||
const geminiImageGenTool = tool(
|
||||
async ({ prompt, image_ids, aspectRatio, imageSize }, _runnableConfig) => {
|
||||
if (!prompt) {
|
||||
throw new Error('Missing required field: prompt');
|
||||
}
|
||||
|
||||
logger.debug('[GeminiImageGen] Generating image with prompt:', prompt?.substring(0, 100));
|
||||
logger.debug('[GeminiImageGen] Options:', { aspectRatio, imageSize });
|
||||
|
||||
// Initialize Gemini client with user-provided credentials
|
||||
let ai;
|
||||
try {
|
||||
ai = await initializeGeminiClient({
|
||||
GEMINI_API_KEY,
|
||||
GOOGLE_KEY,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error('[GeminiImageGen] Failed to initialize client:', error);
|
||||
return [
|
||||
[{ type: ContentTypes.TEXT, text: `Failed to initialize Gemini: ${error.message}` }],
|
||||
{ content: [], file_ids: [] },
|
||||
];
|
||||
}
|
||||
|
||||
// Build request contents
|
||||
const contents = [{ text: replaceUnwantedChars(prompt) }];
|
||||
|
||||
// Add context images if provided
|
||||
if (image_ids?.length > 0) {
|
||||
const contextImages = await convertImagesToInlineData({
|
||||
imageFiles,
|
||||
image_ids,
|
||||
req,
|
||||
fileStrategy,
|
||||
});
|
||||
contents.push(...contextImages);
|
||||
logger.debug('[GeminiImageGen] Added', contextImages.length, 'context images');
|
||||
}
|
||||
|
||||
// Generate image
|
||||
let apiResponse;
|
||||
const geminiModel = process.env.GEMINI_IMAGE_MODEL || 'gemini-2.5-flash-image';
|
||||
try {
|
||||
// Build config with optional imageConfig
|
||||
const config = {
|
||||
responseModalities: ['TEXT', 'IMAGE'],
|
||||
};
|
||||
|
||||
// Add imageConfig if aspectRatio or imageSize is specified
|
||||
// Note: gemini-2.5-flash-image doesn't support imageSize
|
||||
const supportsImageSize = !geminiModel.includes('gemini-2.5-flash-image');
|
||||
if (aspectRatio || (imageSize && supportsImageSize)) {
|
||||
config.imageConfig = {};
|
||||
if (aspectRatio) {
|
||||
config.imageConfig.aspectRatio = aspectRatio;
|
||||
}
|
||||
if (imageSize && supportsImageSize) {
|
||||
config.imageConfig.imageSize = imageSize;
|
||||
}
|
||||
}
|
||||
|
||||
apiResponse = await ai.models.generateContent({
|
||||
model: geminiModel,
|
||||
contents,
|
||||
config,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error('[GeminiImageGen] API error:', error);
|
||||
return [
|
||||
[{ type: ContentTypes.TEXT, text: `Image generation failed: ${error.message}` }],
|
||||
{ content: [], file_ids: [] },
|
||||
];
|
||||
}
|
||||
|
||||
// Check for safety blocks
|
||||
const safetyBlock = checkForSafetyBlock(apiResponse);
|
||||
if (safetyBlock) {
|
||||
logger.warn('[GeminiImageGen] Safety block:', safetyBlock);
|
||||
const errorMsg = 'Image blocked by content safety filters. Please try different content.';
|
||||
return [[{ type: ContentTypes.TEXT, text: errorMsg }], { content: [], file_ids: [] }];
|
||||
}
|
||||
|
||||
const rawImageData = apiResponse.candidates?.[0]?.content?.parts?.find((p) => p.inlineData)
|
||||
?.inlineData?.data;
|
||||
|
||||
if (!rawImageData) {
|
||||
logger.warn('[GeminiImageGen] No image data in response');
|
||||
return [
|
||||
[{ type: ContentTypes.TEXT, text: 'No image was generated. Please try again.' }],
|
||||
{ content: [], file_ids: [] },
|
||||
];
|
||||
}
|
||||
|
||||
const rawBuffer = Buffer.from(rawImageData, 'base64');
|
||||
const { buffer: convertedBuffer, format: outputFormat } = await convertImageFormat(
|
||||
rawBuffer,
|
||||
imageOutputType,
|
||||
);
|
||||
const imageData = convertedBuffer.toString('base64');
|
||||
const mimeType = outputFormat === 'jpeg' ? 'image/jpeg' : `image/${outputFormat}`;
|
||||
|
||||
logger.debug('[GeminiImageGen] Image format:', { outputFormat, mimeType });
|
||||
|
||||
let imageUrl;
|
||||
const useLocalStorage = !fileStrategy || fileStrategy === FileSources.local;
|
||||
|
||||
if (useLocalStorage) {
|
||||
try {
|
||||
imageUrl = await saveImageLocally(imageData, outputFormat, userId);
|
||||
} catch (error) {
|
||||
logger.error('[GeminiImageGen] Local save failed:', error);
|
||||
imageUrl = `data:${mimeType};base64,${imageData}`;
|
||||
}
|
||||
} else {
|
||||
const cloudUrl = await saveToCloudStorage({
|
||||
base64Data: imageData,
|
||||
format: outputFormat,
|
||||
processFileURL,
|
||||
fileStrategy,
|
||||
userId,
|
||||
});
|
||||
|
||||
if (cloudUrl) {
|
||||
imageUrl = cloudUrl;
|
||||
} else {
|
||||
// Fallback to local
|
||||
try {
|
||||
imageUrl = await saveImageLocally(imageData, outputFormat, userId);
|
||||
} catch (_error) {
|
||||
imageUrl = `data:${mimeType};base64,${imageData}`;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.debug('[GeminiImageGen] Image URL:', imageUrl);
|
||||
|
||||
// For the artifact, we need a data URL (same as OpenAI)
|
||||
// The local file save is for persistence, but the response needs a data URL
|
||||
const dataUrl = `data:${mimeType};base64,${imageData}`;
|
||||
|
||||
// Return in content_and_artifact format (same as OpenAI)
|
||||
const file_ids = [v4()];
|
||||
const content = [
|
||||
{
|
||||
type: ContentTypes.IMAGE_URL,
|
||||
image_url: { url: dataUrl },
|
||||
},
|
||||
];
|
||||
|
||||
const textResponse = [
|
||||
{
|
||||
type: ContentTypes.TEXT,
|
||||
text:
|
||||
displayMessage +
|
||||
`\n\ngenerated_image_id: "${file_ids[0]}"` +
|
||||
(image_ids?.length > 0 ? `\nreferenced_image_ids: ["${image_ids.join('", "')}"]` : ''),
|
||||
},
|
||||
];
|
||||
|
||||
// Record token usage for balance tracking (don't await to avoid blocking response)
|
||||
const conversationId = _runnableConfig?.configurable?.thread_id;
|
||||
recordTokenUsage({
|
||||
usageMetadata: apiResponse.usageMetadata,
|
||||
req,
|
||||
userId,
|
||||
conversationId,
|
||||
model: geminiModel,
|
||||
}).catch((error) => {
|
||||
logger.error('[GeminiImageGen] Failed to record token usage:', error);
|
||||
});
|
||||
|
||||
return [textResponse, { content, file_ids }];
|
||||
},
|
||||
{
|
||||
...geminiToolkit.gemini_image_gen,
|
||||
responseFormat: 'content_and_artifact',
|
||||
},
|
||||
);
|
||||
|
||||
return geminiImageGenTool;
|
||||
}
|
||||
|
||||
// Export both for compatibility
|
||||
module.exports = createGeminiImageTool;
|
||||
module.exports.createGeminiImageTool = createGeminiImageTool;
|
||||
72
api/app/clients/tools/structured/GoogleSearch.js
Normal file
72
api/app/clients/tools/structured/GoogleSearch.js
Normal file
@ -0,0 +1,72 @@
|
||||
const { z } = require('zod');
|
||||
const { Tool } = require('@langchain/core/tools');
|
||||
const { getEnvironmentVariable } = require('@langchain/core/utils/env');
|
||||
|
||||
class GoogleSearchResults extends Tool {
|
||||
static lc_name() {
|
||||
return 'google';
|
||||
}
|
||||
|
||||
constructor(fields = {}) {
|
||||
super(fields);
|
||||
this.name = 'google';
|
||||
this.envVarApiKey = 'GOOGLE_SEARCH_API_KEY';
|
||||
this.envVarSearchEngineId = 'GOOGLE_CSE_ID';
|
||||
this.override = fields.override ?? false;
|
||||
this.apiKey = fields[this.envVarApiKey] ?? getEnvironmentVariable(this.envVarApiKey);
|
||||
this.searchEngineId =
|
||||
fields[this.envVarSearchEngineId] ?? getEnvironmentVariable(this.envVarSearchEngineId);
|
||||
|
||||
if (!this.override && (!this.apiKey || !this.searchEngineId)) {
|
||||
throw new Error(
|
||||
`Missing ${this.envVarApiKey} or ${this.envVarSearchEngineId} environment variable.`,
|
||||
);
|
||||
}
|
||||
|
||||
this.kwargs = fields?.kwargs ?? {};
|
||||
this.name = 'google';
|
||||
this.description =
|
||||
'A search engine optimized for comprehensive, accurate, and trusted results. Useful for when you need to answer questions about current events.';
|
||||
|
||||
this.schema = z.object({
|
||||
query: z.string().min(1).describe('The search query string.'),
|
||||
max_results: z
|
||||
.number()
|
||||
.min(1)
|
||||
.max(10)
|
||||
.optional()
|
||||
.describe('The maximum number of search results to return. Defaults to 10.'),
|
||||
// Note: Google API has its own parameters for search customization, adjust as needed.
|
||||
});
|
||||
}
|
||||
|
||||
async _call(input) {
|
||||
const validationResult = this.schema.safeParse(input);
|
||||
if (!validationResult.success) {
|
||||
throw new Error(`Validation failed: ${JSON.stringify(validationResult.error.issues)}`);
|
||||
}
|
||||
|
||||
const { query, max_results = 5 } = validationResult.data;
|
||||
|
||||
const response = await fetch(
|
||||
`https://www.googleapis.com/customsearch/v1?key=${this.apiKey}&cx=${
|
||||
this.searchEngineId
|
||||
}&q=${encodeURIComponent(query)}&num=${max_results}`,
|
||||
{
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
const json = await response.json();
|
||||
if (!response.ok) {
|
||||
throw new Error(`Request failed with status ${response.status}: ${json.error.message}`);
|
||||
}
|
||||
|
||||
return JSON.stringify(json);
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = GoogleSearchResults;
|
||||
413
api/app/clients/tools/structured/OpenAIImageTools.js
Normal file
413
api/app/clients/tools/structured/OpenAIImageTools.js
Normal file
@ -0,0 +1,413 @@
|
||||
const axios = require('axios');
|
||||
const { v4 } = require('uuid');
|
||||
const OpenAI = require('openai');
|
||||
const FormData = require('form-data');
|
||||
const { ProxyAgent } = require('undici');
|
||||
const { tool } = require('@langchain/core/tools');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { HttpsProxyAgent } = require('https-proxy-agent');
|
||||
const { ContentTypes, EImageOutputType } = require('librechat-data-provider');
|
||||
const { logAxiosError, oaiToolkit, extractBaseURL } = require('@librechat/api');
|
||||
const { getStrategyFunctions } = require('~/server/services/Files/strategies');
|
||||
const { getFiles } = require('~/models');
|
||||
|
||||
const displayMessage =
|
||||
"The tool displayed an image. All generated images are already plainly visible, so don't repeat the descriptions in detail. Do not list download links as they are available in the UI already. The user may download the images by clicking on them, but do not mention anything about downloading to the user.";
|
||||
|
||||
/**
|
||||
* Replaces unwanted characters from the input string
|
||||
* @param {string} inputString - The input string to process
|
||||
* @returns {string} - The processed string
|
||||
*/
|
||||
function replaceUnwantedChars(inputString) {
|
||||
return inputString
|
||||
.replace(/\r\n|\r|\n/g, ' ')
|
||||
.replace(/"/g, '')
|
||||
.trim();
|
||||
}
|
||||
|
||||
function returnValue(value) {
|
||||
if (typeof value === 'string') {
|
||||
return [value, {}];
|
||||
} else if (typeof value === 'object') {
|
||||
if (Array.isArray(value)) {
|
||||
return value;
|
||||
}
|
||||
return [displayMessage, value];
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
function createAbortHandler() {
|
||||
return function () {
|
||||
logger.debug('[ImageGenOAI] Image generation aborted');
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates OpenAI Image tools (generation and editing)
|
||||
* @param {Object} fields - Configuration fields
|
||||
* @param {ServerRequest} fields.req - Whether the tool is being used in an agent context
|
||||
* @param {boolean} fields.isAgent - Whether the tool is being used in an agent context
|
||||
* @param {string} fields.IMAGE_GEN_OAI_API_KEY - The OpenAI API key
|
||||
* @param {boolean} [fields.override] - Whether to override the API key check, necessary for app initialization
|
||||
* @param {MongoFile[]} [fields.imageFiles] - The images to be used for editing
|
||||
* @param {string} [fields.imageOutputType] - The image output type configuration
|
||||
* @param {string} [fields.fileStrategy] - The file storage strategy
|
||||
* @returns {Array<ReturnType<tool>>} - Array of image tools
|
||||
*/
|
||||
function createOpenAIImageTools(fields = {}) {
|
||||
/** @type {boolean} Used to initialize the Tool without necessary variables. */
|
||||
const override = fields.override ?? false;
|
||||
/** @type {boolean} */
|
||||
if (!override && !fields.isAgent) {
|
||||
throw new Error('This tool is only available for agents.');
|
||||
}
|
||||
const { req } = fields;
|
||||
const imageOutputType = fields.imageOutputType || EImageOutputType.PNG;
|
||||
const appFileStrategy = fields.fileStrategy;
|
||||
|
||||
const getApiKey = () => {
|
||||
const apiKey = process.env.IMAGE_GEN_OAI_API_KEY ?? '';
|
||||
if (!apiKey && !override) {
|
||||
throw new Error('Missing IMAGE_GEN_OAI_API_KEY environment variable.');
|
||||
}
|
||||
return apiKey;
|
||||
};
|
||||
|
||||
let apiKey = fields.IMAGE_GEN_OAI_API_KEY ?? getApiKey();
|
||||
const closureConfig = { apiKey };
|
||||
|
||||
const imageModel = process.env.IMAGE_GEN_OAI_MODEL || 'gpt-image-1';
|
||||
|
||||
let baseURL = 'https://api.openai.com/v1/';
|
||||
if (!override && process.env.IMAGE_GEN_OAI_BASEURL) {
|
||||
baseURL = extractBaseURL(process.env.IMAGE_GEN_OAI_BASEURL);
|
||||
closureConfig.baseURL = baseURL;
|
||||
}
|
||||
|
||||
// Note: Azure may not yet support the latest image generation models
|
||||
if (
|
||||
!override &&
|
||||
process.env.IMAGE_GEN_OAI_AZURE_API_VERSION &&
|
||||
process.env.IMAGE_GEN_OAI_BASEURL
|
||||
) {
|
||||
baseURL = process.env.IMAGE_GEN_OAI_BASEURL;
|
||||
closureConfig.baseURL = baseURL;
|
||||
closureConfig.defaultQuery = { 'api-version': process.env.IMAGE_GEN_OAI_AZURE_API_VERSION };
|
||||
closureConfig.defaultHeaders = {
|
||||
'api-key': process.env.IMAGE_GEN_OAI_API_KEY,
|
||||
'Content-Type': 'application/json',
|
||||
};
|
||||
closureConfig.apiKey = process.env.IMAGE_GEN_OAI_API_KEY;
|
||||
}
|
||||
|
||||
const imageFiles = fields.imageFiles ?? [];
|
||||
|
||||
/**
|
||||
* Image Generation Tool
|
||||
*/
|
||||
const imageGenTool = tool(
|
||||
async (
|
||||
{
|
||||
prompt,
|
||||
background = 'auto',
|
||||
n = 1,
|
||||
output_compression = 100,
|
||||
quality = 'auto',
|
||||
size = 'auto',
|
||||
},
|
||||
runnableConfig,
|
||||
) => {
|
||||
if (!prompt) {
|
||||
throw new Error('Missing required field: prompt');
|
||||
}
|
||||
const clientConfig = { ...closureConfig };
|
||||
if (process.env.PROXY) {
|
||||
const proxyAgent = new ProxyAgent(process.env.PROXY);
|
||||
clientConfig.fetchOptions = {
|
||||
dispatcher: proxyAgent,
|
||||
};
|
||||
}
|
||||
|
||||
/** @type {OpenAI} */
|
||||
const openai = new OpenAI(clientConfig);
|
||||
let output_format = imageOutputType;
|
||||
if (
|
||||
background === 'transparent' &&
|
||||
output_format !== EImageOutputType.PNG &&
|
||||
output_format !== EImageOutputType.WEBP
|
||||
) {
|
||||
logger.warn(
|
||||
'[ImageGenOAI] Transparent background requires PNG or WebP format, defaulting to PNG',
|
||||
);
|
||||
output_format = EImageOutputType.PNG;
|
||||
}
|
||||
|
||||
let resp;
|
||||
/** @type {AbortSignal} */
|
||||
let derivedSignal = null;
|
||||
/** @type {() => void} */
|
||||
let abortHandler = null;
|
||||
|
||||
try {
|
||||
if (runnableConfig?.signal) {
|
||||
derivedSignal = AbortSignal.any([runnableConfig.signal]);
|
||||
abortHandler = createAbortHandler();
|
||||
derivedSignal.addEventListener('abort', abortHandler, { once: true });
|
||||
}
|
||||
|
||||
resp = await openai.images.generate(
|
||||
{
|
||||
model: imageModel,
|
||||
prompt: replaceUnwantedChars(prompt),
|
||||
n: Math.min(Math.max(1, n), 10),
|
||||
background,
|
||||
output_format,
|
||||
output_compression:
|
||||
output_format === EImageOutputType.WEBP || output_format === EImageOutputType.JPEG
|
||||
? output_compression
|
||||
: undefined,
|
||||
quality,
|
||||
size,
|
||||
},
|
||||
{
|
||||
signal: derivedSignal,
|
||||
},
|
||||
);
|
||||
} catch (error) {
|
||||
const message = '[image_gen_oai] Problem generating the image:';
|
||||
logAxiosError({ error, message });
|
||||
return returnValue(`Something went wrong when trying to generate the image. The OpenAI API may be unavailable:
|
||||
Error Message: ${error.message}`);
|
||||
} finally {
|
||||
if (abortHandler && derivedSignal) {
|
||||
derivedSignal.removeEventListener('abort', abortHandler);
|
||||
}
|
||||
}
|
||||
|
||||
if (!resp) {
|
||||
return returnValue(
|
||||
'Something went wrong when trying to generate the image. The OpenAI API may be unavailable',
|
||||
);
|
||||
}
|
||||
|
||||
// For gpt-image-1, the response contains base64-encoded images
|
||||
// TODO: handle cost in `resp.usage`
|
||||
const base64Image = resp.data[0].b64_json;
|
||||
|
||||
if (!base64Image) {
|
||||
return returnValue(
|
||||
'No image data returned from OpenAI API. There may be a problem with the API or your configuration.',
|
||||
);
|
||||
}
|
||||
|
||||
const content = [
|
||||
{
|
||||
type: ContentTypes.IMAGE_URL,
|
||||
image_url: {
|
||||
url: `data:image/${output_format};base64,${base64Image}`,
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
const file_ids = [v4()];
|
||||
const response = [
|
||||
{
|
||||
type: ContentTypes.TEXT,
|
||||
text: displayMessage + `\n\ngenerated_image_id: "${file_ids[0]}"`,
|
||||
},
|
||||
];
|
||||
return [response, { content, file_ids }];
|
||||
},
|
||||
oaiToolkit.image_gen_oai,
|
||||
);
|
||||
|
||||
/**
|
||||
* Image Editing Tool
|
||||
*/
|
||||
const imageEditTool = tool(
|
||||
async ({ prompt, image_ids, quality = 'auto', size = 'auto' }, runnableConfig) => {
|
||||
if (!prompt) {
|
||||
throw new Error('Missing required field: prompt');
|
||||
}
|
||||
|
||||
const clientConfig = { ...closureConfig };
|
||||
if (process.env.PROXY) {
|
||||
const proxyAgent = new ProxyAgent(process.env.PROXY);
|
||||
clientConfig.fetchOptions = {
|
||||
dispatcher: proxyAgent,
|
||||
};
|
||||
}
|
||||
|
||||
const formData = new FormData();
|
||||
formData.append('model', imageModel);
|
||||
formData.append('prompt', replaceUnwantedChars(prompt));
|
||||
// TODO: `mask` support
|
||||
// TODO: more than 1 image support
|
||||
// formData.append('n', n.toString());
|
||||
formData.append('quality', quality);
|
||||
formData.append('size', size);
|
||||
|
||||
/** @type {Record<FileSources, undefined | NodeStreamDownloader<File>>} */
|
||||
const streamMethods = {};
|
||||
|
||||
const requestFilesMap = Object.fromEntries(imageFiles.map((f) => [f.file_id, { ...f }]));
|
||||
|
||||
const orderedFiles = new Array(image_ids.length);
|
||||
const idsToFetch = [];
|
||||
const indexOfMissing = Object.create(null);
|
||||
|
||||
for (let i = 0; i < image_ids.length; i++) {
|
||||
const id = image_ids[i];
|
||||
const file = requestFilesMap[id];
|
||||
|
||||
if (file) {
|
||||
orderedFiles[i] = file;
|
||||
} else {
|
||||
idsToFetch.push(id);
|
||||
indexOfMissing[id] = i;
|
||||
}
|
||||
}
|
||||
|
||||
if (idsToFetch.length) {
|
||||
const fetchedFiles = await getFiles(
|
||||
{
|
||||
user: req.user.id,
|
||||
file_id: { $in: idsToFetch },
|
||||
height: { $exists: true },
|
||||
width: { $exists: true },
|
||||
},
|
||||
{},
|
||||
{},
|
||||
);
|
||||
|
||||
for (const file of fetchedFiles) {
|
||||
requestFilesMap[file.file_id] = file;
|
||||
orderedFiles[indexOfMissing[file.file_id]] = file;
|
||||
}
|
||||
}
|
||||
for (const imageFile of orderedFiles) {
|
||||
if (!imageFile) {
|
||||
continue;
|
||||
}
|
||||
/** @type {NodeStream<File>} */
|
||||
let stream;
|
||||
/** @type {NodeStreamDownloader<File>} */
|
||||
let getDownloadStream;
|
||||
const source = imageFile.source || appFileStrategy;
|
||||
if (!source) {
|
||||
throw new Error('No source found for image file');
|
||||
}
|
||||
if (streamMethods[source]) {
|
||||
getDownloadStream = streamMethods[source];
|
||||
} else {
|
||||
({ getDownloadStream } = getStrategyFunctions(source));
|
||||
streamMethods[source] = getDownloadStream;
|
||||
}
|
||||
if (!getDownloadStream) {
|
||||
throw new Error(`No download stream method found for source: ${source}`);
|
||||
}
|
||||
stream = await getDownloadStream(req, imageFile.filepath);
|
||||
if (!stream) {
|
||||
throw new Error('Failed to get download stream for image file');
|
||||
}
|
||||
formData.append('image[]', stream, {
|
||||
filename: imageFile.filename,
|
||||
contentType: imageFile.type,
|
||||
});
|
||||
}
|
||||
|
||||
/** @type {import('axios').RawAxiosHeaders} */
|
||||
let headers = {
|
||||
...formData.getHeaders(),
|
||||
};
|
||||
|
||||
if (process.env.IMAGE_GEN_OAI_AZURE_API_VERSION && process.env.IMAGE_GEN_OAI_BASEURL) {
|
||||
headers['api-key'] = apiKey;
|
||||
} else {
|
||||
headers['Authorization'] = `Bearer ${apiKey}`;
|
||||
}
|
||||
|
||||
/** @type {AbortSignal} */
|
||||
let derivedSignal = null;
|
||||
/** @type {() => void} */
|
||||
let abortHandler = null;
|
||||
|
||||
try {
|
||||
if (runnableConfig?.signal) {
|
||||
derivedSignal = AbortSignal.any([runnableConfig.signal]);
|
||||
abortHandler = createAbortHandler();
|
||||
derivedSignal.addEventListener('abort', abortHandler, { once: true });
|
||||
}
|
||||
|
||||
/** @type {import('axios').AxiosRequestConfig} */
|
||||
const axiosConfig = {
|
||||
headers,
|
||||
...clientConfig,
|
||||
signal: derivedSignal,
|
||||
baseURL,
|
||||
};
|
||||
|
||||
if (process.env.PROXY) {
|
||||
axiosConfig.httpsAgent = new HttpsProxyAgent(process.env.PROXY);
|
||||
}
|
||||
|
||||
if (process.env.IMAGE_GEN_OAI_AZURE_API_VERSION && process.env.IMAGE_GEN_OAI_BASEURL) {
|
||||
axiosConfig.params = {
|
||||
'api-version': process.env.IMAGE_GEN_OAI_AZURE_API_VERSION,
|
||||
...axiosConfig.params,
|
||||
};
|
||||
}
|
||||
const response = await axios.post('/images/edits', formData, axiosConfig);
|
||||
|
||||
if (!response.data || !response.data.data || !response.data.data.length) {
|
||||
return returnValue(
|
||||
'No image data returned from OpenAI API. There may be a problem with the API or your configuration.',
|
||||
);
|
||||
}
|
||||
|
||||
const base64Image = response.data.data[0].b64_json;
|
||||
if (!base64Image) {
|
||||
return returnValue(
|
||||
'No image data returned from OpenAI API. There may be a problem with the API or your configuration.',
|
||||
);
|
||||
}
|
||||
|
||||
const content = [
|
||||
{
|
||||
type: ContentTypes.IMAGE_URL,
|
||||
image_url: {
|
||||
url: `data:image/${imageOutputType};base64,${base64Image}`,
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
const file_ids = [v4()];
|
||||
const textResponse = [
|
||||
{
|
||||
type: ContentTypes.TEXT,
|
||||
text:
|
||||
displayMessage +
|
||||
`\n\ngenerated_image_id: "${file_ids[0]}"\nreferenced_image_ids: ["${image_ids.join('", "')}"]`,
|
||||
},
|
||||
];
|
||||
return [textResponse, { content, file_ids }];
|
||||
} catch (error) {
|
||||
const message = '[image_edit_oai] Problem editing the image:';
|
||||
logAxiosError({ error, message });
|
||||
return returnValue(`Something went wrong when trying to edit the image. The OpenAI API may be unavailable:
|
||||
Error Message: ${error.message || 'Unknown error'}`);
|
||||
} finally {
|
||||
if (abortHandler && derivedSignal) {
|
||||
derivedSignal.removeEventListener('abort', abortHandler);
|
||||
}
|
||||
}
|
||||
},
|
||||
oaiToolkit.image_edit_oai,
|
||||
);
|
||||
|
||||
return [imageGenTool, imageEditTool];
|
||||
}
|
||||
|
||||
module.exports = createOpenAIImageTools;
|
||||
317
api/app/clients/tools/structured/OpenWeather.js
Normal file
317
api/app/clients/tools/structured/OpenWeather.js
Normal file
@ -0,0 +1,317 @@
|
||||
const { Tool } = require('@langchain/core/tools');
|
||||
const { z } = require('zod');
|
||||
const { getEnvironmentVariable } = require('@langchain/core/utils/env');
|
||||
const fetch = require('node-fetch');
|
||||
|
||||
/**
|
||||
* Map user-friendly units to OpenWeather units.
|
||||
* Defaults to Celsius if not specified.
|
||||
*/
|
||||
function mapUnitsToOpenWeather(unit) {
|
||||
if (!unit) {
|
||||
return 'metric';
|
||||
} // Default to Celsius
|
||||
switch (unit) {
|
||||
case 'Celsius':
|
||||
return 'metric';
|
||||
case 'Kelvin':
|
||||
return 'standard';
|
||||
case 'Fahrenheit':
|
||||
return 'imperial';
|
||||
default:
|
||||
return 'metric'; // fallback
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Recursively round temperature fields in the API response.
|
||||
*/
|
||||
function roundTemperatures(obj) {
|
||||
const tempKeys = new Set([
|
||||
'temp',
|
||||
'feels_like',
|
||||
'dew_point',
|
||||
'day',
|
||||
'min',
|
||||
'max',
|
||||
'night',
|
||||
'eve',
|
||||
'morn',
|
||||
'afternoon',
|
||||
'morning',
|
||||
'evening',
|
||||
]);
|
||||
|
||||
if (Array.isArray(obj)) {
|
||||
return obj.map((item) => roundTemperatures(item));
|
||||
} else if (obj && typeof obj === 'object') {
|
||||
for (const key of Object.keys(obj)) {
|
||||
const value = obj[key];
|
||||
if (value && typeof value === 'object') {
|
||||
obj[key] = roundTemperatures(value);
|
||||
} else if (typeof value === 'number' && tempKeys.has(key)) {
|
||||
obj[key] = Math.round(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
return obj;
|
||||
}
|
||||
|
||||
class OpenWeather extends Tool {
|
||||
name = 'open_weather';
|
||||
description =
|
||||
'Provides weather data from OpenWeather One Call API 3.0. ' +
|
||||
'Actions: help, current_forecast, timestamp, daily_aggregation, overview. ' +
|
||||
'If lat/lon not provided, specify "city" for geocoding. ' +
|
||||
'Units: "Celsius", "Kelvin", or "Fahrenheit" (default: Celsius). ' +
|
||||
'For timestamp action, use "date" in YYYY-MM-DD format.';
|
||||
|
||||
schema = z.object({
|
||||
action: z.enum(['help', 'current_forecast', 'timestamp', 'daily_aggregation', 'overview']),
|
||||
city: z.string().optional(),
|
||||
lat: z.number().optional(),
|
||||
lon: z.number().optional(),
|
||||
exclude: z.string().optional(),
|
||||
units: z.enum(['Celsius', 'Kelvin', 'Fahrenheit']).optional(),
|
||||
lang: z.string().optional(),
|
||||
date: z.string().optional(), // For timestamp and daily_aggregation
|
||||
tz: z.string().optional(),
|
||||
});
|
||||
|
||||
constructor(fields = {}) {
|
||||
super();
|
||||
this.envVar = 'OPENWEATHER_API_KEY';
|
||||
this.override = fields.override ?? false;
|
||||
this.apiKey = fields[this.envVar] ?? this.getApiKey();
|
||||
}
|
||||
|
||||
getApiKey() {
|
||||
const key = getEnvironmentVariable(this.envVar);
|
||||
if (!key && !this.override) {
|
||||
throw new Error(`Missing ${this.envVar} environment variable.`);
|
||||
}
|
||||
return key;
|
||||
}
|
||||
|
||||
async geocodeCity(city) {
|
||||
const geocodeUrl = `https://api.openweathermap.org/geo/1.0/direct?q=${encodeURIComponent(
|
||||
city,
|
||||
)}&limit=1&appid=${this.apiKey}`;
|
||||
const res = await fetch(geocodeUrl);
|
||||
const data = await res.json();
|
||||
if (!res.ok || !Array.isArray(data) || data.length === 0) {
|
||||
throw new Error(`Could not find coordinates for city: ${city}`);
|
||||
}
|
||||
return { lat: data[0].lat, lon: data[0].lon };
|
||||
}
|
||||
|
||||
convertDateToUnix(dateStr) {
|
||||
const parts = dateStr.split('-');
|
||||
if (parts.length !== 3) {
|
||||
throw new Error('Invalid date format. Expected YYYY-MM-DD.');
|
||||
}
|
||||
const year = parseInt(parts[0], 10);
|
||||
const month = parseInt(parts[1], 10);
|
||||
const day = parseInt(parts[2], 10);
|
||||
if (isNaN(year) || isNaN(month) || isNaN(day)) {
|
||||
throw new Error('Invalid date format. Expected YYYY-MM-DD with valid numbers.');
|
||||
}
|
||||
|
||||
const dateObj = new Date(Date.UTC(year, month - 1, day, 0, 0, 0));
|
||||
if (isNaN(dateObj.getTime())) {
|
||||
throw new Error('Invalid date provided. Cannot parse into a valid date.');
|
||||
}
|
||||
|
||||
return Math.floor(dateObj.getTime() / 1000);
|
||||
}
|
||||
|
||||
async _call(args) {
|
||||
try {
|
||||
const { action, city, lat, lon, exclude, units, lang, date, tz } = args;
|
||||
const owmUnits = mapUnitsToOpenWeather(units);
|
||||
|
||||
if (action === 'help') {
|
||||
return JSON.stringify(
|
||||
{
|
||||
title: 'OpenWeather One Call API 3.0 Help',
|
||||
description: 'Guidance on using the OpenWeather One Call API 3.0.',
|
||||
endpoints: {
|
||||
current_and_forecast: {
|
||||
endpoint: 'data/3.0/onecall',
|
||||
data_provided: [
|
||||
'Current weather',
|
||||
'Minute forecast (1h)',
|
||||
'Hourly forecast (48h)',
|
||||
'Daily forecast (8 days)',
|
||||
'Government weather alerts',
|
||||
],
|
||||
required_params: [['lat', 'lon'], ['city']],
|
||||
optional_params: ['exclude', 'units (Celsius/Kelvin/Fahrenheit)', 'lang'],
|
||||
usage_example: {
|
||||
city: 'Knoxville, Tennessee',
|
||||
units: 'Fahrenheit',
|
||||
lang: 'en',
|
||||
},
|
||||
},
|
||||
weather_for_timestamp: {
|
||||
endpoint: 'data/3.0/onecall/timemachine',
|
||||
data_provided: [
|
||||
'Historical weather (since 1979-01-01)',
|
||||
'Future forecast up to 4 days ahead',
|
||||
],
|
||||
required_params: [
|
||||
['lat', 'lon', 'date (YYYY-MM-DD)'],
|
||||
['city', 'date (YYYY-MM-DD)'],
|
||||
],
|
||||
optional_params: ['units (Celsius/Kelvin/Fahrenheit)', 'lang'],
|
||||
usage_example: {
|
||||
city: 'Knoxville, Tennessee',
|
||||
date: '2020-03-04',
|
||||
units: 'Fahrenheit',
|
||||
lang: 'en',
|
||||
},
|
||||
},
|
||||
daily_aggregation: {
|
||||
endpoint: 'data/3.0/onecall/day_summary',
|
||||
data_provided: [
|
||||
'Aggregated weather data for a specific date (1979-01-02 to 1.5 years ahead)',
|
||||
],
|
||||
required_params: [
|
||||
['lat', 'lon', 'date (YYYY-MM-DD)'],
|
||||
['city', 'date (YYYY-MM-DD)'],
|
||||
],
|
||||
optional_params: ['units (Celsius/Kelvin/Fahrenheit)', 'lang', 'tz'],
|
||||
usage_example: {
|
||||
city: 'Knoxville, Tennessee',
|
||||
date: '2020-03-04',
|
||||
units: 'Celsius',
|
||||
lang: 'en',
|
||||
},
|
||||
},
|
||||
weather_overview: {
|
||||
endpoint: 'data/3.0/onecall/overview',
|
||||
data_provided: ['Human-readable weather summary (today/tomorrow)'],
|
||||
required_params: [['lat', 'lon'], ['city']],
|
||||
optional_params: ['date (YYYY-MM-DD)', 'units (Celsius/Kelvin/Fahrenheit)'],
|
||||
usage_example: {
|
||||
city: 'Knoxville, Tennessee',
|
||||
date: '2024-05-13',
|
||||
units: 'Celsius',
|
||||
},
|
||||
},
|
||||
},
|
||||
notes: [
|
||||
'If lat/lon not provided, you can specify a city name and it will be geocoded.',
|
||||
'For the timestamp action, provide a date in YYYY-MM-DD format instead of a Unix timestamp.',
|
||||
'By default, temperatures are returned in Celsius.',
|
||||
'You can specify units as Celsius, Kelvin, or Fahrenheit.',
|
||||
'All temperatures are rounded to the nearest degree.',
|
||||
],
|
||||
errors: [
|
||||
'400: Bad Request (missing/invalid params)',
|
||||
'401: Unauthorized (check API key)',
|
||||
'404: Not Found (no data or city)',
|
||||
'429: Too many requests',
|
||||
'5xx: Internal error',
|
||||
],
|
||||
},
|
||||
null,
|
||||
2,
|
||||
);
|
||||
}
|
||||
|
||||
let finalLat = lat;
|
||||
let finalLon = lon;
|
||||
|
||||
// If lat/lon not provided but city is given, geocode it
|
||||
if ((finalLat == null || finalLon == null) && city) {
|
||||
const coords = await this.geocodeCity(city);
|
||||
finalLat = coords.lat;
|
||||
finalLon = coords.lon;
|
||||
}
|
||||
|
||||
if (['current_forecast', 'timestamp', 'daily_aggregation', 'overview'].includes(action)) {
|
||||
if (typeof finalLat !== 'number' || typeof finalLon !== 'number') {
|
||||
return "Error: lat and lon are required and must be numbers for this action (or specify 'city').";
|
||||
}
|
||||
}
|
||||
|
||||
const baseUrl = 'https://api.openweathermap.org/data/3.0';
|
||||
let endpoint = '';
|
||||
const params = new URLSearchParams({ appid: this.apiKey, units: owmUnits });
|
||||
|
||||
let dt;
|
||||
if (action === 'timestamp') {
|
||||
if (!date) {
|
||||
return "Error: For timestamp action, a 'date' in YYYY-MM-DD format is required.";
|
||||
}
|
||||
dt = this.convertDateToUnix(date);
|
||||
}
|
||||
|
||||
if (action === 'daily_aggregation' && !date) {
|
||||
return 'Error: date (YYYY-MM-DD) is required for daily_aggregation action.';
|
||||
}
|
||||
|
||||
switch (action) {
|
||||
case 'current_forecast':
|
||||
endpoint = '/onecall';
|
||||
params.append('lat', String(finalLat));
|
||||
params.append('lon', String(finalLon));
|
||||
if (exclude) {
|
||||
params.append('exclude', exclude);
|
||||
}
|
||||
if (lang) {
|
||||
params.append('lang', lang);
|
||||
}
|
||||
break;
|
||||
case 'timestamp':
|
||||
endpoint = '/onecall/timemachine';
|
||||
params.append('lat', String(finalLat));
|
||||
params.append('lon', String(finalLon));
|
||||
params.append('dt', String(dt));
|
||||
if (lang) {
|
||||
params.append('lang', lang);
|
||||
}
|
||||
break;
|
||||
case 'daily_aggregation':
|
||||
endpoint = '/onecall/day_summary';
|
||||
params.append('lat', String(finalLat));
|
||||
params.append('lon', String(finalLon));
|
||||
params.append('date', date);
|
||||
if (lang) {
|
||||
params.append('lang', lang);
|
||||
}
|
||||
if (tz) {
|
||||
params.append('tz', tz);
|
||||
}
|
||||
break;
|
||||
case 'overview':
|
||||
endpoint = '/onecall/overview';
|
||||
params.append('lat', String(finalLat));
|
||||
params.append('lon', String(finalLon));
|
||||
if (date) {
|
||||
params.append('date', date);
|
||||
}
|
||||
break;
|
||||
default:
|
||||
return `Error: Unknown action: ${action}`;
|
||||
}
|
||||
|
||||
const url = `${baseUrl}${endpoint}?${params.toString()}`;
|
||||
const response = await fetch(url);
|
||||
const json = await response.json();
|
||||
if (!response.ok) {
|
||||
return `Error: OpenWeather API request failed with status ${response.status}: ${
|
||||
json.message || JSON.stringify(json)
|
||||
}`;
|
||||
}
|
||||
|
||||
const roundedJson = roundTemperatures(json);
|
||||
return JSON.stringify(roundedJson);
|
||||
} catch (err) {
|
||||
return `Error: ${err.message}`;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = OpenWeather;
|
||||
196
api/app/clients/tools/structured/StableDiffusion.js
Normal file
196
api/app/clients/tools/structured/StableDiffusion.js
Normal file
@ -0,0 +1,196 @@
|
||||
// Generates image using stable diffusion webui's api (automatic1111)
|
||||
const fs = require('fs');
|
||||
const { z } = require('zod');
|
||||
const path = require('path');
|
||||
const axios = require('axios');
|
||||
const sharp = require('sharp');
|
||||
const { v4: uuidv4 } = require('uuid');
|
||||
const { Tool } = require('@langchain/core/tools');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { FileContext, ContentTypes } = require('librechat-data-provider');
|
||||
const { getBasePath } = require('@librechat/api');
|
||||
const paths = require('~/config/paths');
|
||||
|
||||
const displayMessage =
|
||||
"Stable Diffusion displayed an image. All generated images are already plainly visible, so don't repeat the descriptions in detail. Do not list download links as they are available in the UI already. The user may download the images by clicking on them, but do not mention anything about downloading to the user.";
|
||||
|
||||
class StableDiffusionAPI extends Tool {
|
||||
constructor(fields) {
|
||||
super();
|
||||
/** @type {string} User ID */
|
||||
this.userId = fields.userId;
|
||||
/** @type {ServerRequest | undefined} Express Request object, only provided by ToolService */
|
||||
this.req = fields.req;
|
||||
/** @type {boolean} Used to initialize the Tool without necessary variables. */
|
||||
this.override = fields.override ?? false;
|
||||
/** @type {boolean} Necessary for output to contain all image metadata. */
|
||||
this.returnMetadata = fields.returnMetadata ?? false;
|
||||
/** @type {boolean} */
|
||||
this.isAgent = fields.isAgent;
|
||||
if (fields.uploadImageBuffer) {
|
||||
/** @type {uploadImageBuffer} Necessary for output to contain all image metadata. */
|
||||
this.uploadImageBuffer = fields.uploadImageBuffer.bind(this);
|
||||
}
|
||||
|
||||
this.name = 'stable-diffusion';
|
||||
this.url = fields.SD_WEBUI_URL || this.getServerURL();
|
||||
this.description_for_model = `// Generate images and visuals using text.
|
||||
// Guidelines:
|
||||
// - ALWAYS use {{"prompt": "7+ detailed keywords", "negative_prompt": "7+ detailed keywords"}} structure for queries.
|
||||
// - ALWAYS include the markdown url in your final response to show the user: }/images/id.png)
|
||||
// - Visually describe the moods, details, structures, styles, and/or proportions of the image. Remember, the focus is on visual attributes.
|
||||
// - Craft your input by "showing" and not "telling" the imagery. Think in terms of what you'd want to see in a photograph or a painting.
|
||||
// - Here's an example for generating a realistic portrait photo of a man:
|
||||
// "prompt":"photo of a man in black clothes, half body, high detailed skin, coastline, overcast weather, wind, waves, 8k uhd, dslr, soft lighting, high quality, film grain, Fujifilm XT3"
|
||||
// "negative_prompt":"semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, out of frame, low quality, ugly, mutation, deformed"
|
||||
// - Generate images only once per human query unless explicitly requested by the user`;
|
||||
this.description =
|
||||
"You can generate images using text with 'stable-diffusion'. This tool is exclusively for visual content.";
|
||||
this.schema = z.object({
|
||||
prompt: z
|
||||
.string()
|
||||
.describe(
|
||||
'Detailed keywords to describe the subject, using at least 7 keywords to accurately describe the image, separated by comma',
|
||||
),
|
||||
negative_prompt: z
|
||||
.string()
|
||||
.describe(
|
||||
'Keywords we want to exclude from the final image, using at least 7 keywords to accurately describe the image, separated by comma',
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
replaceNewLinesWithSpaces(inputString) {
|
||||
return inputString.replace(/\r\n|\r|\n/g, ' ');
|
||||
}
|
||||
|
||||
getMarkdownImageUrl(imageName) {
|
||||
const imageUrl = path
|
||||
.join(this.relativePath, this.userId, imageName)
|
||||
.replace(/\\/g, '/')
|
||||
.replace('public/', '');
|
||||
return ``;
|
||||
}
|
||||
|
||||
returnValue(value) {
|
||||
if (this.isAgent === true && typeof value === 'string') {
|
||||
return [value, {}];
|
||||
} else if (this.isAgent === true && typeof value === 'object') {
|
||||
return [displayMessage, value];
|
||||
}
|
||||
|
||||
return value;
|
||||
}
|
||||
|
||||
getServerURL() {
|
||||
const url = process.env.SD_WEBUI_URL || '';
|
||||
if (!url && !this.override) {
|
||||
throw new Error('Missing SD_WEBUI_URL environment variable.');
|
||||
}
|
||||
return url;
|
||||
}
|
||||
|
||||
async _call(data) {
|
||||
const url = this.url;
|
||||
const { prompt, negative_prompt } = data;
|
||||
const payload = {
|
||||
prompt,
|
||||
negative_prompt,
|
||||
cfg_scale: 4.5,
|
||||
steps: 22,
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
};
|
||||
let generationResponse;
|
||||
try {
|
||||
generationResponse = await axios.post(`${url}/sdapi/v1/txt2img`, payload);
|
||||
} catch (error) {
|
||||
logger.error('[StableDiffusion] Error while generating image:', error);
|
||||
return 'Error making API request.';
|
||||
}
|
||||
const image = generationResponse.data.images[0];
|
||||
|
||||
/** @type {{ height: number, width: number, seed: number, infotexts: string[] }} */
|
||||
let info = {};
|
||||
try {
|
||||
info = JSON.parse(generationResponse.data.info);
|
||||
} catch (error) {
|
||||
logger.error('[StableDiffusion] Error while getting image metadata:', error);
|
||||
}
|
||||
|
||||
const file_id = uuidv4();
|
||||
const imageName = `${file_id}.png`;
|
||||
const { imageOutput: imageOutputPath, clientPath } = paths;
|
||||
const filepath = path.join(imageOutputPath, this.userId, imageName);
|
||||
this.relativePath = path.relative(clientPath, imageOutputPath);
|
||||
|
||||
if (!fs.existsSync(path.join(imageOutputPath, this.userId))) {
|
||||
fs.mkdirSync(path.join(imageOutputPath, this.userId), { recursive: true });
|
||||
}
|
||||
|
||||
try {
|
||||
if (this.isAgent) {
|
||||
const content = [
|
||||
{
|
||||
type: ContentTypes.IMAGE_URL,
|
||||
image_url: {
|
||||
url: `data:image/png;base64,${image}`,
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
const response = [
|
||||
{
|
||||
type: ContentTypes.TEXT,
|
||||
text: displayMessage,
|
||||
},
|
||||
];
|
||||
return [response, { content }];
|
||||
}
|
||||
|
||||
const buffer = Buffer.from(image.split(',', 1)[0], 'base64');
|
||||
if (this.returnMetadata && this.uploadImageBuffer && this.req) {
|
||||
const file = await this.uploadImageBuffer({
|
||||
req: this.req,
|
||||
context: FileContext.image_generation,
|
||||
resize: false,
|
||||
metadata: {
|
||||
buffer,
|
||||
height: info.height,
|
||||
width: info.width,
|
||||
bytes: Buffer.byteLength(buffer),
|
||||
filename: imageName,
|
||||
type: 'image/png',
|
||||
file_id,
|
||||
},
|
||||
});
|
||||
|
||||
const generationInfo = info.infotexts[0].split('\n').pop();
|
||||
return {
|
||||
...file,
|
||||
prompt,
|
||||
metadata: {
|
||||
negative_prompt,
|
||||
seed: info.seed,
|
||||
info: generationInfo,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
await sharp(buffer)
|
||||
.withMetadata({
|
||||
iptcpng: {
|
||||
parameters: info.infotexts[0],
|
||||
},
|
||||
})
|
||||
.toFile(filepath);
|
||||
this.result = this.getMarkdownImageUrl(imageName);
|
||||
} catch (error) {
|
||||
logger.error('[StableDiffusion] Error while saving the image:', error);
|
||||
}
|
||||
|
||||
return this.returnValue(this.result);
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = StableDiffusionAPI;
|
||||
77
api/app/clients/tools/structured/TavilySearch.js
Normal file
77
api/app/clients/tools/structured/TavilySearch.js
Normal file
@ -0,0 +1,77 @@
|
||||
const { z } = require('zod');
|
||||
const { ProxyAgent, fetch } = require('undici');
|
||||
const { tool } = require('@langchain/core/tools');
|
||||
const { getApiKey } = require('./credentials');
|
||||
|
||||
function createTavilySearchTool(fields = {}) {
|
||||
const envVar = 'TAVILY_API_KEY';
|
||||
const override = fields.override ?? false;
|
||||
const apiKey = fields.apiKey ?? getApiKey(envVar, override);
|
||||
const kwargs = fields?.kwargs ?? {};
|
||||
|
||||
return tool(
|
||||
async (input) => {
|
||||
const { query, ...rest } = input;
|
||||
|
||||
const requestBody = {
|
||||
api_key: apiKey,
|
||||
query,
|
||||
...rest,
|
||||
...kwargs,
|
||||
};
|
||||
|
||||
const fetchOptions = {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
body: JSON.stringify(requestBody),
|
||||
};
|
||||
|
||||
if (process.env.PROXY) {
|
||||
fetchOptions.dispatcher = new ProxyAgent(process.env.PROXY);
|
||||
}
|
||||
|
||||
const response = await fetch('https://api.tavily.com/search', fetchOptions);
|
||||
|
||||
const json = await response.json();
|
||||
if (!response.ok) {
|
||||
throw new Error(`Request failed with status ${response.status}: ${json.error}`);
|
||||
}
|
||||
|
||||
return JSON.stringify(json);
|
||||
},
|
||||
{
|
||||
name: 'tavily_search_results_json',
|
||||
description:
|
||||
'A search engine optimized for comprehensive, accurate, and trusted results. Useful for when you need to answer questions about current events.',
|
||||
schema: z.object({
|
||||
query: z.string().min(1).describe('The search query string.'),
|
||||
max_results: z
|
||||
.number()
|
||||
.min(1)
|
||||
.max(10)
|
||||
.optional()
|
||||
.describe('The maximum number of search results to return. Defaults to 5.'),
|
||||
search_depth: z
|
||||
.enum(['basic', 'advanced'])
|
||||
.optional()
|
||||
.describe(
|
||||
'The depth of the search, affecting result quality and response time (`basic` or `advanced`). Default is basic for quick results and advanced for indepth high quality results but longer response time. Advanced calls equals 2 requests.',
|
||||
),
|
||||
include_images: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.describe(
|
||||
'Whether to include a list of query-related images in the response. Default is False.',
|
||||
),
|
||||
include_answer: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.describe('Whether to include answers in the search results. Default is False.'),
|
||||
}),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
module.exports = createTavilySearchTool;
|
||||
131
api/app/clients/tools/structured/TavilySearchResults.js
Normal file
131
api/app/clients/tools/structured/TavilySearchResults.js
Normal file
@ -0,0 +1,131 @@
|
||||
const { z } = require('zod');
|
||||
const { ProxyAgent, fetch } = require('undici');
|
||||
const { Tool } = require('@langchain/core/tools');
|
||||
const { getEnvironmentVariable } = require('@langchain/core/utils/env');
|
||||
|
||||
class TavilySearchResults extends Tool {
|
||||
static lc_name() {
|
||||
return 'TavilySearchResults';
|
||||
}
|
||||
|
||||
constructor(fields = {}) {
|
||||
super(fields);
|
||||
this.envVar = 'TAVILY_API_KEY';
|
||||
/* Used to initialize the Tool without necessary variables. */
|
||||
this.override = fields.override ?? false;
|
||||
this.apiKey = fields[this.envVar] ?? this.getApiKey();
|
||||
|
||||
this.kwargs = fields?.kwargs ?? {};
|
||||
this.name = 'tavily_search_results_json';
|
||||
this.description =
|
||||
'A search engine optimized for comprehensive, accurate, and trusted results. Useful for when you need to answer questions about current events.';
|
||||
|
||||
this.schema = z.object({
|
||||
query: z.string().min(1).describe('The search query string.'),
|
||||
max_results: z
|
||||
.number()
|
||||
.min(1)
|
||||
.max(10)
|
||||
.optional()
|
||||
.describe('The maximum number of search results to return. Defaults to 5.'),
|
||||
search_depth: z
|
||||
.enum(['basic', 'advanced'])
|
||||
.optional()
|
||||
.describe(
|
||||
'The depth of the search, affecting result quality and response time (`basic` or `advanced`). Default is basic for quick results and advanced for indepth high quality results but longer response time. Advanced calls equals 2 requests.',
|
||||
),
|
||||
include_images: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.describe(
|
||||
'Whether to include a list of query-related images in the response. Default is False.',
|
||||
),
|
||||
include_answer: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.describe('Whether to include answers in the search results. Default is False.'),
|
||||
include_raw_content: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.describe('Whether to include raw content in the search results. Default is False.'),
|
||||
include_domains: z
|
||||
.array(z.string())
|
||||
.optional()
|
||||
.describe('A list of domains to specifically include in the search results.'),
|
||||
exclude_domains: z
|
||||
.array(z.string())
|
||||
.optional()
|
||||
.describe('A list of domains to specifically exclude from the search results.'),
|
||||
topic: z
|
||||
.enum(['general', 'news', 'finance'])
|
||||
.optional()
|
||||
.describe(
|
||||
'The category of the search. Use news ONLY if query SPECIFCALLY mentions the word "news".',
|
||||
),
|
||||
time_range: z
|
||||
.enum(['day', 'week', 'month', 'year', 'd', 'w', 'm', 'y'])
|
||||
.optional()
|
||||
.describe('The time range back from the current date to filter results.'),
|
||||
days: z
|
||||
.number()
|
||||
.min(1)
|
||||
.optional()
|
||||
.describe('Number of days back from the current date to include. Only if topic is news.'),
|
||||
include_image_descriptions: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.describe(
|
||||
'When include_images is true, also add a descriptive text for each image. Default is false.',
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
getApiKey() {
|
||||
const apiKey = getEnvironmentVariable(this.envVar);
|
||||
if (!apiKey && !this.override) {
|
||||
throw new Error(`Missing ${this.envVar} environment variable.`);
|
||||
}
|
||||
return apiKey;
|
||||
}
|
||||
|
||||
async _call(input) {
|
||||
const validationResult = this.schema.safeParse(input);
|
||||
if (!validationResult.success) {
|
||||
throw new Error(`Validation failed: ${JSON.stringify(validationResult.error.issues)}`);
|
||||
}
|
||||
|
||||
const { query, ...rest } = validationResult.data;
|
||||
|
||||
const requestBody = {
|
||||
api_key: this.apiKey,
|
||||
query,
|
||||
...rest,
|
||||
...this.kwargs,
|
||||
};
|
||||
|
||||
const fetchOptions = {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
body: JSON.stringify(requestBody),
|
||||
};
|
||||
|
||||
if (process.env.PROXY) {
|
||||
fetchOptions.dispatcher = new ProxyAgent(process.env.PROXY);
|
||||
}
|
||||
|
||||
const response = await fetch('https://api.tavily.com/search', fetchOptions);
|
||||
|
||||
const json = await response.json();
|
||||
if (!response.ok) {
|
||||
throw new Error(
|
||||
`Request failed with status ${response.status}: ${json?.detail?.error || json?.error}`,
|
||||
);
|
||||
}
|
||||
|
||||
return JSON.stringify(json);
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = TavilySearchResults;
|
||||
88
api/app/clients/tools/structured/TraversaalSearch.js
Normal file
88
api/app/clients/tools/structured/TraversaalSearch.js
Normal file
@ -0,0 +1,88 @@
|
||||
const { z } = require('zod');
|
||||
const { Tool } = require('@langchain/core/tools');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { getEnvironmentVariable } = require('@langchain/core/utils/env');
|
||||
|
||||
/**
|
||||
* Tool for the Traversaal AI search API, Ares.
|
||||
*/
|
||||
class TraversaalSearch extends Tool {
|
||||
static lc_name() {
|
||||
return 'TraversaalSearch';
|
||||
}
|
||||
constructor(fields) {
|
||||
super(fields);
|
||||
this.name = 'traversaal_search';
|
||||
this.description = `An AI search engine optimized for comprehensive, accurate, and trusted results.
|
||||
Useful for when you need to answer questions about current events. Input should be a search query.`;
|
||||
this.description_for_model =
|
||||
'\'Please create a specific sentence for the AI to understand and use as a query to search the web based on the user\'s request. For example, "Find information about the highest mountains in the world." or "Show me the latest news articles about climate change and its impact on polar ice caps."\'';
|
||||
this.schema = z.object({
|
||||
query: z
|
||||
.string()
|
||||
.describe(
|
||||
"A properly written sentence to be interpreted by an AI to search the web according to the user's request.",
|
||||
),
|
||||
});
|
||||
|
||||
this.apiKey = fields?.TRAVERSAAL_API_KEY ?? this.getApiKey();
|
||||
}
|
||||
|
||||
getApiKey() {
|
||||
const apiKey = getEnvironmentVariable('TRAVERSAAL_API_KEY');
|
||||
if (!apiKey && this.override) {
|
||||
throw new Error(
|
||||
'No Traversaal API key found. Either set an environment variable named "TRAVERSAAL_API_KEY" or pass an API key as "apiKey".',
|
||||
);
|
||||
}
|
||||
return apiKey;
|
||||
}
|
||||
|
||||
async _call({ query }, _runManager) {
|
||||
const body = {
|
||||
query: [query],
|
||||
};
|
||||
try {
|
||||
const response = await fetch('https://api-ares.traversaal.ai/live/predict', {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'content-type': 'application/json',
|
||||
'x-api-key': this.apiKey,
|
||||
},
|
||||
body: JSON.stringify({ ...body }),
|
||||
});
|
||||
const json = await response.json();
|
||||
if (!response.ok) {
|
||||
throw new Error(
|
||||
`Request failed with status code ${response.status}: ${json.error ?? json.message}`,
|
||||
);
|
||||
}
|
||||
if (!json.data) {
|
||||
throw new Error('Could not parse Traversaal API results. Please try again.');
|
||||
}
|
||||
|
||||
const baseText = json.data?.response_text ?? '';
|
||||
const sources = json.data?.web_url;
|
||||
const noResponse = 'No response found in Traversaal API results';
|
||||
|
||||
if (!baseText && !sources) {
|
||||
return noResponse;
|
||||
}
|
||||
|
||||
const sourcesText = sources?.length ? '\n\nSources:\n - ' + sources.join('\n - ') : '';
|
||||
|
||||
const result = baseText + sourcesText;
|
||||
|
||||
if (!result) {
|
||||
return noResponse;
|
||||
}
|
||||
|
||||
return result;
|
||||
} catch (error) {
|
||||
logger.error('Traversaal API request failed', error);
|
||||
return `Traversaal API request failed: ${error.message}`;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = TraversaalSearch;
|
||||
95
api/app/clients/tools/structured/Wolfram.js
Normal file
95
api/app/clients/tools/structured/Wolfram.js
Normal file
@ -0,0 +1,95 @@
|
||||
/* eslint-disable no-useless-escape */
|
||||
const { z } = require('zod');
|
||||
const axios = require('axios');
|
||||
const { Tool } = require('@langchain/core/tools');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
|
||||
class WolframAlphaAPI extends Tool {
|
||||
constructor(fields) {
|
||||
super();
|
||||
/* Used to initialize the Tool without necessary variables. */
|
||||
this.override = fields.override ?? false;
|
||||
|
||||
this.name = 'wolfram';
|
||||
this.apiKey = fields.WOLFRAM_APP_ID || this.getAppId();
|
||||
this.description_for_model = `// Access dynamic computation and curated data from WolframAlpha and Wolfram Cloud.
|
||||
// General guidelines:
|
||||
// - Use only getWolframAlphaResults or getWolframCloudResults endpoints.
|
||||
// - Prefer getWolframAlphaResults unless Wolfram Language code should be evaluated.
|
||||
// - Use getWolframAlphaResults for natural-language queries in English; translate non-English queries before sending, then respond in the original language.
|
||||
// - Use getWolframCloudResults for problems solvable with Wolfram Language code.
|
||||
// - Suggest only Wolfram Language for external computation.
|
||||
// - Inform users if information is not from Wolfram endpoints.
|
||||
// - Display image URLs with Image Markdown syntax: . You must prefix the caption brackets with "!".
|
||||
// - ALWAYS use this exponent notation: \`6*10^14\`, NEVER \`6e14\`.
|
||||
// - ALWAYS use {{"input": query}} structure for queries to Wolfram endpoints; \`query\` must ONLY be a single-line string.
|
||||
// - ALWAYS use proper Markdown formatting for all math, scientific, and chemical formulas, symbols, etc.: '$$\n[expression]\n$$' for standalone cases and '\( [expression] \)' when inline.
|
||||
// - Format inline Wolfram Language code with Markdown code formatting.
|
||||
// - Never mention your knowledge cutoff date; Wolfram may return more recent data. getWolframAlphaResults guidelines:
|
||||
// - Understands natural language queries about entities in chemistry, physics, geography, history, art, astronomy, and more.
|
||||
// - Performs mathematical calculations, date and unit conversions, formula solving, etc.
|
||||
// - Convert inputs to simplified keyword queries whenever possible (e.g. convert "how many people live in France" to "France population").
|
||||
// - Use ONLY single-letter variable names, with or without integer subscript (e.g., n, n1, n_1).
|
||||
// - Use named physical constants (e.g., 'speed of light') without numerical substitution.
|
||||
// - Include a space between compound units (e.g., "Ω m" for "ohm*meter").
|
||||
// - To solve for a variable in an equation with units, consider solving a corresponding equation without units; exclude counting units (e.g., books), include genuine units (e.g., kg).
|
||||
// - If data for multiple properties is needed, make separate calls for each property.
|
||||
// - If a Wolfram Alpha result is not relevant to the query:
|
||||
// -- If Wolfram provides multiple 'Assumptions' for a query, choose the more relevant one(s) without explaining the initial result. If you are unsure, ask the user to choose.
|
||||
// -- Re-send the exact same 'input' with NO modifications, and add the 'assumption' parameter, formatted as a list, with the relevant values.
|
||||
// -- ONLY simplify or rephrase the initial query if a more relevant 'Assumption' or other input suggestions are not provided.
|
||||
// -- Do not explain each step unless user input is needed. Proceed directly to making a better API call based on the available assumptions.`;
|
||||
this.description = `WolframAlpha offers computation, math, curated knowledge, and real-time data. It handles natural language queries and performs complex calculations.
|
||||
Follow the guidelines to get the best results.`;
|
||||
this.schema = z.object({
|
||||
input: z.string().describe('Natural language query to WolframAlpha following the guidelines'),
|
||||
});
|
||||
}
|
||||
|
||||
async fetchRawText(url) {
|
||||
try {
|
||||
const response = await axios.get(url, { responseType: 'text' });
|
||||
return response.data;
|
||||
} catch (error) {
|
||||
logger.error('[WolframAlphaAPI] Error fetching raw text:', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
getAppId() {
|
||||
const appId = process.env.WOLFRAM_APP_ID || '';
|
||||
if (!appId && !this.override) {
|
||||
throw new Error('Missing WOLFRAM_APP_ID environment variable.');
|
||||
}
|
||||
return appId;
|
||||
}
|
||||
|
||||
createWolframAlphaURL(query) {
|
||||
// Clean up query
|
||||
const formattedQuery = query.replaceAll(/`/g, '').replaceAll(/\n/g, ' ');
|
||||
const baseURL = 'https://www.wolframalpha.com/api/v1/llm-api';
|
||||
const encodedQuery = encodeURIComponent(formattedQuery);
|
||||
const appId = this.apiKey || this.getAppId();
|
||||
const url = `${baseURL}?input=${encodedQuery}&appid=${appId}`;
|
||||
return url;
|
||||
}
|
||||
|
||||
async _call(data) {
|
||||
try {
|
||||
const { input } = data;
|
||||
const url = this.createWolframAlphaURL(input);
|
||||
const response = await this.fetchRawText(url);
|
||||
return response;
|
||||
} catch (error) {
|
||||
if (error.response && error.response.data) {
|
||||
logger.error('[WolframAlphaAPI] Error data:', error);
|
||||
return error.response.data;
|
||||
} else {
|
||||
logger.error('[WolframAlphaAPI] Error querying Wolfram Alpha', error);
|
||||
return 'There was an error querying Wolfram Alpha.';
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = WolframAlphaAPI;
|
||||
13
api/app/clients/tools/structured/credentials.js
Normal file
13
api/app/clients/tools/structured/credentials.js
Normal file
@ -0,0 +1,13 @@
|
||||
const { getEnvironmentVariable } = require('@langchain/core/utils/env');
|
||||
|
||||
function getApiKey(envVar, override) {
|
||||
const key = getEnvironmentVariable(envVar);
|
||||
if (!key && !override) {
|
||||
throw new Error(`Missing ${envVar} environment variable.`);
|
||||
}
|
||||
return key;
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
getApiKey,
|
||||
};
|
||||
60
api/app/clients/tools/structured/specs/DALLE3-proxy.spec.js
Normal file
60
api/app/clients/tools/structured/specs/DALLE3-proxy.spec.js
Normal file
@ -0,0 +1,60 @@
|
||||
const DALLE3 = require('../DALLE3');
|
||||
const { ProxyAgent } = require('undici');
|
||||
|
||||
jest.mock('tiktoken');
|
||||
const processFileURL = jest.fn();
|
||||
|
||||
describe('DALLE3 Proxy Configuration', () => {
|
||||
let originalEnv;
|
||||
|
||||
beforeAll(() => {
|
||||
originalEnv = { ...process.env };
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
jest.resetModules();
|
||||
process.env = { ...originalEnv };
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
process.env = originalEnv;
|
||||
});
|
||||
|
||||
it('should configure ProxyAgent in fetchOptions.dispatcher when PROXY env is set', () => {
|
||||
// Set proxy environment variable
|
||||
process.env.PROXY = 'http://proxy.example.com:8080';
|
||||
process.env.DALLE_API_KEY = 'test-api-key';
|
||||
|
||||
// Create instance
|
||||
const dalleWithProxy = new DALLE3({ processFileURL });
|
||||
|
||||
// Check that the openai client exists
|
||||
expect(dalleWithProxy.openai).toBeDefined();
|
||||
|
||||
// Check that _options exists and has fetchOptions with a dispatcher
|
||||
expect(dalleWithProxy.openai._options).toBeDefined();
|
||||
expect(dalleWithProxy.openai._options.fetchOptions).toBeDefined();
|
||||
expect(dalleWithProxy.openai._options.fetchOptions.dispatcher).toBeDefined();
|
||||
expect(dalleWithProxy.openai._options.fetchOptions.dispatcher).toBeInstanceOf(ProxyAgent);
|
||||
});
|
||||
|
||||
it('should not configure ProxyAgent when PROXY env is not set', () => {
|
||||
// Ensure PROXY is not set
|
||||
delete process.env.PROXY;
|
||||
process.env.DALLE_API_KEY = 'test-api-key';
|
||||
|
||||
// Create instance
|
||||
const dalleWithoutProxy = new DALLE3({ processFileURL });
|
||||
|
||||
// Check that the openai client exists
|
||||
expect(dalleWithoutProxy.openai).toBeDefined();
|
||||
|
||||
// Check that _options exists but fetchOptions either doesn't exist or doesn't have a dispatcher
|
||||
expect(dalleWithoutProxy.openai._options).toBeDefined();
|
||||
|
||||
// fetchOptions should either not exist or not have a dispatcher
|
||||
if (dalleWithoutProxy.openai._options.fetchOptions) {
|
||||
expect(dalleWithoutProxy.openai._options.fetchOptions.dispatcher).toBeUndefined();
|
||||
}
|
||||
});
|
||||
});
|
||||
216
api/app/clients/tools/structured/specs/DALLE3.spec.js
Normal file
216
api/app/clients/tools/structured/specs/DALLE3.spec.js
Normal file
@ -0,0 +1,216 @@
|
||||
const OpenAI = require('openai');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const DALLE3 = require('../DALLE3');
|
||||
|
||||
jest.mock('openai');
|
||||
jest.mock('@librechat/data-schemas', () => {
|
||||
return {
|
||||
logger: {
|
||||
info: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
debug: jest.fn(),
|
||||
error: jest.fn(),
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
jest.mock('tiktoken', () => {
|
||||
return {
|
||||
encoding_for_model: jest.fn().mockReturnValue({
|
||||
encode: jest.fn(),
|
||||
decode: jest.fn(),
|
||||
}),
|
||||
};
|
||||
});
|
||||
|
||||
const processFileURL = jest.fn();
|
||||
|
||||
const generate = jest.fn();
|
||||
OpenAI.mockImplementation(() => ({
|
||||
images: {
|
||||
generate,
|
||||
},
|
||||
}));
|
||||
|
||||
jest.mock('fs', () => {
|
||||
return {
|
||||
existsSync: jest.fn(),
|
||||
mkdirSync: jest.fn(),
|
||||
promises: {
|
||||
writeFile: jest.fn(),
|
||||
readFile: jest.fn(),
|
||||
unlink: jest.fn(),
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
jest.mock('path', () => {
|
||||
return {
|
||||
resolve: jest.fn(),
|
||||
join: jest.fn(),
|
||||
relative: jest.fn(),
|
||||
extname: jest.fn().mockImplementation((filename) => {
|
||||
return filename.slice(filename.lastIndexOf('.'));
|
||||
}),
|
||||
};
|
||||
});
|
||||
|
||||
describe('DALLE3', () => {
|
||||
let originalEnv;
|
||||
let dalle; // Keep this declaration if you need to use dalle in other tests
|
||||
const mockApiKey = 'mock_api_key';
|
||||
|
||||
beforeAll(() => {
|
||||
// Save the original process.env
|
||||
originalEnv = { ...process.env };
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
// Reset the process.env before each test
|
||||
jest.resetModules();
|
||||
process.env = { ...originalEnv, DALLE_API_KEY: mockApiKey };
|
||||
// Instantiate DALLE3 for tests that do not depend on DALLE3_SYSTEM_PROMPT
|
||||
dalle = new DALLE3({ processFileURL });
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
jest.clearAllMocks();
|
||||
// Restore the original process.env after each test
|
||||
process.env = originalEnv;
|
||||
});
|
||||
|
||||
it('should throw an error if all potential API keys are missing', () => {
|
||||
delete process.env.DALLE3_API_KEY;
|
||||
delete process.env.DALLE_API_KEY;
|
||||
expect(() => new DALLE3()).toThrow('Missing DALLE_API_KEY environment variable.');
|
||||
});
|
||||
|
||||
it('should replace unwanted characters in input string', () => {
|
||||
const input = 'This is a test\nstring with "quotes" and new lines.';
|
||||
const expectedOutput = 'This is a test string with quotes and new lines.';
|
||||
expect(dalle.replaceUnwantedChars(input)).toBe(expectedOutput);
|
||||
});
|
||||
|
||||
it('should generate markdown image URL correctly', () => {
|
||||
const imageName = 'test.png';
|
||||
const markdownImage = dalle.wrapInMarkdown(imageName);
|
||||
expect(markdownImage).toBe('');
|
||||
});
|
||||
|
||||
it('should call OpenAI API with correct parameters', async () => {
|
||||
const mockData = {
|
||||
prompt: 'A test prompt',
|
||||
quality: 'standard',
|
||||
size: '1024x1024',
|
||||
style: 'vivid',
|
||||
};
|
||||
|
||||
const mockResponse = {
|
||||
data: [
|
||||
{
|
||||
url: 'http://example.com/img-test.png',
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
generate.mockResolvedValue(mockResponse);
|
||||
processFileURL.mockResolvedValue({
|
||||
filepath: 'http://example.com/img-test.png',
|
||||
});
|
||||
|
||||
const result = await dalle._call(mockData);
|
||||
|
||||
expect(generate).toHaveBeenCalledWith({
|
||||
model: 'dall-e-3',
|
||||
quality: mockData.quality,
|
||||
style: mockData.style,
|
||||
size: mockData.size,
|
||||
prompt: mockData.prompt,
|
||||
n: 1,
|
||||
});
|
||||
|
||||
expect(result).toContain('![generated image]');
|
||||
});
|
||||
|
||||
it('should use the system prompt if provided', () => {
|
||||
process.env.DALLE3_SYSTEM_PROMPT = 'System prompt for testing';
|
||||
jest.resetModules(); // This will ensure the module is fresh and will read the new env var
|
||||
const DALLE3 = require('../DALLE3'); // Re-require after setting the env var
|
||||
const dalleWithSystemPrompt = new DALLE3();
|
||||
expect(dalleWithSystemPrompt.description_for_model).toBe('System prompt for testing');
|
||||
});
|
||||
|
||||
it('should not use the system prompt if not provided', async () => {
|
||||
delete process.env.DALLE3_SYSTEM_PROMPT;
|
||||
const dalleWithoutSystemPrompt = new DALLE3();
|
||||
expect(dalleWithoutSystemPrompt.description_for_model).not.toBe('System prompt for testing');
|
||||
});
|
||||
|
||||
it('should throw an error if prompt is missing', async () => {
|
||||
const mockData = {
|
||||
quality: 'standard',
|
||||
size: '1024x1024',
|
||||
style: 'vivid',
|
||||
};
|
||||
await expect(dalle._call(mockData)).rejects.toThrow('Missing required field: prompt');
|
||||
});
|
||||
|
||||
it('should log appropriate debug values', async () => {
|
||||
const mockData = {
|
||||
prompt: 'A test prompt',
|
||||
};
|
||||
const mockResponse = {
|
||||
data: [
|
||||
{
|
||||
url: 'http://example.com/invalid-url',
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
generate.mockResolvedValue(mockResponse);
|
||||
await dalle._call(mockData);
|
||||
expect(logger.debug).toHaveBeenCalledWith('[DALL-E-3]', {
|
||||
data: { url: 'http://example.com/invalid-url' },
|
||||
theImageUrl: 'http://example.com/invalid-url',
|
||||
extension: expect.any(String),
|
||||
imageBasename: expect.any(String),
|
||||
imageExt: expect.any(String),
|
||||
imageName: expect.any(String),
|
||||
});
|
||||
});
|
||||
|
||||
it('should log an error and return the image URL if there is an error saving the image', async () => {
|
||||
const mockData = {
|
||||
prompt: 'A test prompt',
|
||||
};
|
||||
const mockResponse = {
|
||||
data: [
|
||||
{
|
||||
url: 'http://example.com/img-test.png',
|
||||
},
|
||||
],
|
||||
};
|
||||
const error = new Error('Error while saving the image');
|
||||
generate.mockResolvedValue(mockResponse);
|
||||
processFileURL.mockRejectedValue(error);
|
||||
const result = await dalle._call(mockData);
|
||||
expect(logger.error).toHaveBeenCalledWith('Error while saving the image:', error);
|
||||
expect(result).toBe('Failed to save the image locally. Error while saving the image');
|
||||
});
|
||||
|
||||
it('should handle error when saving image to Firebase Storage fails', async () => {
|
||||
const mockData = {
|
||||
prompt: 'A test prompt',
|
||||
};
|
||||
const mockImageUrl = 'http://example.com/img-test.png';
|
||||
const mockResponse = { data: [{ url: mockImageUrl }] };
|
||||
const error = new Error('Error while saving to Firebase');
|
||||
generate.mockResolvedValue(mockResponse);
|
||||
processFileURL.mockRejectedValue(error);
|
||||
|
||||
const result = await dalle._call(mockData);
|
||||
|
||||
expect(logger.error).toHaveBeenCalledWith('Error while saving the image:', error);
|
||||
expect(result).toContain('Failed to save the image');
|
||||
});
|
||||
});
|
||||
@ -0,0 +1,125 @@
|
||||
const { ProxyAgent } = require('undici');
|
||||
|
||||
/**
|
||||
* These tests verify the proxy wrapper behavior for GeminiImageGen.
|
||||
* Instead of loading the full module (which has many dependencies),
|
||||
* we directly test the wrapper logic that would be applied.
|
||||
*/
|
||||
describe('GeminiImageGen Proxy Configuration', () => {
|
||||
let originalEnv;
|
||||
let originalFetch;
|
||||
|
||||
beforeAll(() => {
|
||||
originalEnv = { ...process.env };
|
||||
originalFetch = globalThis.fetch;
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
process.env = { ...originalEnv };
|
||||
globalThis.fetch = originalFetch;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
process.env = originalEnv;
|
||||
globalThis.fetch = originalFetch;
|
||||
});
|
||||
|
||||
/**
|
||||
* Simulates the proxy wrapper that GeminiImageGen applies at module load.
|
||||
* This is the same logic from GeminiImageGen.js lines 30-42.
|
||||
*/
|
||||
function applyProxyWrapper() {
|
||||
if (process.env.PROXY) {
|
||||
const _originalFetch = globalThis.fetch;
|
||||
const proxyAgent = new ProxyAgent(process.env.PROXY);
|
||||
|
||||
globalThis.fetch = function (url, options = {}) {
|
||||
const urlString = url.toString();
|
||||
if (urlString.includes('googleapis.com')) {
|
||||
options = { ...options, dispatcher: proxyAgent };
|
||||
}
|
||||
return _originalFetch.call(this, url, options);
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
it('should wrap globalThis.fetch when PROXY env is set', () => {
|
||||
process.env.PROXY = 'http://proxy.example.com:8080';
|
||||
|
||||
const fetchBeforeWrap = globalThis.fetch;
|
||||
|
||||
applyProxyWrapper();
|
||||
|
||||
expect(globalThis.fetch).not.toBe(fetchBeforeWrap);
|
||||
});
|
||||
|
||||
it('should not wrap globalThis.fetch when PROXY env is not set', () => {
|
||||
delete process.env.PROXY;
|
||||
|
||||
const fetchBeforeWrap = globalThis.fetch;
|
||||
|
||||
applyProxyWrapper();
|
||||
|
||||
expect(globalThis.fetch).toBe(fetchBeforeWrap);
|
||||
});
|
||||
|
||||
it('should add dispatcher to googleapis.com URLs', async () => {
|
||||
process.env.PROXY = 'http://proxy.example.com:8080';
|
||||
|
||||
let capturedOptions = null;
|
||||
const mockFetch = jest.fn((url, options) => {
|
||||
capturedOptions = options;
|
||||
return Promise.resolve({ ok: true });
|
||||
});
|
||||
globalThis.fetch = mockFetch;
|
||||
|
||||
applyProxyWrapper();
|
||||
|
||||
await globalThis.fetch('https://generativelanguage.googleapis.com/v1/models', {});
|
||||
|
||||
expect(capturedOptions).toBeDefined();
|
||||
expect(capturedOptions.dispatcher).toBeInstanceOf(ProxyAgent);
|
||||
});
|
||||
|
||||
it('should not add dispatcher to non-googleapis.com URLs', async () => {
|
||||
process.env.PROXY = 'http://proxy.example.com:8080';
|
||||
|
||||
let capturedOptions = null;
|
||||
const mockFetch = jest.fn((url, options) => {
|
||||
capturedOptions = options;
|
||||
return Promise.resolve({ ok: true });
|
||||
});
|
||||
globalThis.fetch = mockFetch;
|
||||
|
||||
applyProxyWrapper();
|
||||
|
||||
await globalThis.fetch('https://api.openai.com/v1/images', {});
|
||||
|
||||
expect(capturedOptions).toBeDefined();
|
||||
expect(capturedOptions.dispatcher).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should preserve existing options when adding dispatcher', async () => {
|
||||
process.env.PROXY = 'http://proxy.example.com:8080';
|
||||
|
||||
let capturedOptions = null;
|
||||
const mockFetch = jest.fn((url, options) => {
|
||||
capturedOptions = options;
|
||||
return Promise.resolve({ ok: true });
|
||||
});
|
||||
globalThis.fetch = mockFetch;
|
||||
|
||||
applyProxyWrapper();
|
||||
|
||||
const customHeaders = { 'X-Custom-Header': 'test' };
|
||||
await globalThis.fetch('https://aiplatform.googleapis.com/v1/models', {
|
||||
headers: customHeaders,
|
||||
method: 'POST',
|
||||
});
|
||||
|
||||
expect(capturedOptions).toBeDefined();
|
||||
expect(capturedOptions.dispatcher).toBeInstanceOf(ProxyAgent);
|
||||
expect(capturedOptions.headers).toEqual(customHeaders);
|
||||
expect(capturedOptions.method).toBe('POST');
|
||||
});
|
||||
});
|
||||
50
api/app/clients/tools/structured/specs/GoogleSearch.spec.js
Normal file
50
api/app/clients/tools/structured/specs/GoogleSearch.spec.js
Normal file
@ -0,0 +1,50 @@
|
||||
const GoogleSearch = require('../GoogleSearch');
|
||||
|
||||
jest.mock('node-fetch');
|
||||
jest.mock('@langchain/core/utils/env');
|
||||
|
||||
describe('GoogleSearch', () => {
|
||||
let originalEnv;
|
||||
const mockApiKey = 'mock_api';
|
||||
const mockSearchEngineId = 'mock_search_engine_id';
|
||||
|
||||
beforeAll(() => {
|
||||
originalEnv = { ...process.env };
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
jest.resetModules();
|
||||
process.env = {
|
||||
...originalEnv,
|
||||
GOOGLE_SEARCH_API_KEY: mockApiKey,
|
||||
GOOGLE_CSE_ID: mockSearchEngineId,
|
||||
};
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
jest.clearAllMocks();
|
||||
process.env = originalEnv;
|
||||
});
|
||||
|
||||
it('should use mockApiKey and mockSearchEngineId when environment variables are not set', () => {
|
||||
const instance = new GoogleSearch({
|
||||
GOOGLE_SEARCH_API_KEY: mockApiKey,
|
||||
GOOGLE_CSE_ID: mockSearchEngineId,
|
||||
});
|
||||
expect(instance.apiKey).toBe(mockApiKey);
|
||||
expect(instance.searchEngineId).toBe(mockSearchEngineId);
|
||||
});
|
||||
|
||||
it('should throw an error if GOOGLE_SEARCH_API_KEY or GOOGLE_CSE_ID is missing', () => {
|
||||
delete process.env.GOOGLE_SEARCH_API_KEY;
|
||||
expect(() => new GoogleSearch()).toThrow(
|
||||
'Missing GOOGLE_SEARCH_API_KEY or GOOGLE_CSE_ID environment variable.',
|
||||
);
|
||||
|
||||
process.env.GOOGLE_SEARCH_API_KEY = mockApiKey;
|
||||
delete process.env.GOOGLE_CSE_ID;
|
||||
expect(() => new GoogleSearch()).toThrow(
|
||||
'Missing GOOGLE_SEARCH_API_KEY or GOOGLE_CSE_ID environment variable.',
|
||||
);
|
||||
});
|
||||
});
|
||||
@ -0,0 +1,84 @@
|
||||
const { fetch, ProxyAgent } = require('undici');
|
||||
const TavilySearchResults = require('../TavilySearchResults');
|
||||
|
||||
jest.mock('undici');
|
||||
jest.mock('@langchain/core/utils/env');
|
||||
|
||||
describe('TavilySearchResults', () => {
|
||||
let originalEnv;
|
||||
const mockApiKey = 'mock_api_key';
|
||||
|
||||
beforeAll(() => {
|
||||
originalEnv = { ...process.env };
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
jest.resetModules();
|
||||
jest.clearAllMocks();
|
||||
process.env = {
|
||||
...originalEnv,
|
||||
TAVILY_API_KEY: mockApiKey,
|
||||
};
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
process.env = originalEnv;
|
||||
});
|
||||
|
||||
it('should throw an error if TAVILY_API_KEY is missing', () => {
|
||||
delete process.env.TAVILY_API_KEY;
|
||||
expect(() => new TavilySearchResults()).toThrow('Missing TAVILY_API_KEY environment variable.');
|
||||
});
|
||||
|
||||
it('should use mockApiKey when TAVILY_API_KEY is not set in the environment', () => {
|
||||
const instance = new TavilySearchResults({
|
||||
TAVILY_API_KEY: mockApiKey,
|
||||
});
|
||||
expect(instance.apiKey).toBe(mockApiKey);
|
||||
});
|
||||
|
||||
describe('proxy support', () => {
|
||||
const mockResponse = {
|
||||
ok: true,
|
||||
json: jest.fn().mockResolvedValue({ results: [] }),
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
fetch.mockResolvedValue(mockResponse);
|
||||
});
|
||||
|
||||
it('should use ProxyAgent when PROXY env var is set', async () => {
|
||||
const proxyUrl = 'http://proxy.example.com:8080';
|
||||
process.env.PROXY = proxyUrl;
|
||||
|
||||
const mockProxyAgent = { type: 'proxy-agent' };
|
||||
ProxyAgent.mockImplementation(() => mockProxyAgent);
|
||||
|
||||
const instance = new TavilySearchResults({ TAVILY_API_KEY: mockApiKey });
|
||||
await instance._call({ query: 'test query' });
|
||||
|
||||
expect(ProxyAgent).toHaveBeenCalledWith(proxyUrl);
|
||||
expect(fetch).toHaveBeenCalledWith(
|
||||
'https://api.tavily.com/search',
|
||||
expect.objectContaining({
|
||||
dispatcher: mockProxyAgent,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should not use ProxyAgent when PROXY env var is not set', async () => {
|
||||
delete process.env.PROXY;
|
||||
|
||||
const instance = new TavilySearchResults({ TAVILY_API_KEY: mockApiKey });
|
||||
await instance._call({ query: 'test query' });
|
||||
|
||||
expect(ProxyAgent).not.toHaveBeenCalled();
|
||||
expect(fetch).toHaveBeenCalledWith(
|
||||
'https://api.tavily.com/search',
|
||||
expect.not.objectContaining({
|
||||
dispatcher: expect.anything(),
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
@ -0,0 +1,224 @@
|
||||
// __tests__/openWeather.integration.test.js
|
||||
const OpenWeather = require('../OpenWeather');
|
||||
|
||||
describe('OpenWeather Tool (Integration Test)', () => {
|
||||
let tool;
|
||||
|
||||
beforeAll(() => {
|
||||
tool = new OpenWeather({ override: true });
|
||||
console.log('API Key present:', !!process.env.OPENWEATHER_API_KEY);
|
||||
});
|
||||
|
||||
test('current_forecast with a real API key returns current weather', async () => {
|
||||
// Check if API key is available
|
||||
if (!process.env.OPENWEATHER_API_KEY) {
|
||||
console.warn('Skipping real API test, no OPENWEATHER_API_KEY found.');
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const result = await tool.call({
|
||||
action: 'current_forecast',
|
||||
city: 'London',
|
||||
units: 'Celsius',
|
||||
});
|
||||
|
||||
console.log('Raw API response:', result);
|
||||
|
||||
const parsed = JSON.parse(result);
|
||||
expect(parsed).toHaveProperty('current');
|
||||
expect(typeof parsed.current.temp).toBe('number');
|
||||
} catch (error) {
|
||||
console.error('Test failed with error:', error);
|
||||
throw error;
|
||||
}
|
||||
});
|
||||
|
||||
test('timestamp action with real API key returns historical data', async () => {
|
||||
if (!process.env.OPENWEATHER_API_KEY) {
|
||||
console.warn('Skipping real API test, no OPENWEATHER_API_KEY found.');
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
// Use a date from yesterday to ensure data availability
|
||||
const yesterday = new Date();
|
||||
yesterday.setDate(yesterday.getDate() - 1);
|
||||
const dateStr = yesterday.toISOString().split('T')[0];
|
||||
|
||||
const result = await tool.call({
|
||||
action: 'timestamp',
|
||||
city: 'London',
|
||||
date: dateStr,
|
||||
units: 'Celsius',
|
||||
});
|
||||
|
||||
console.log('Timestamp API response:', result);
|
||||
|
||||
const parsed = JSON.parse(result);
|
||||
expect(parsed).toHaveProperty('data');
|
||||
expect(Array.isArray(parsed.data)).toBe(true);
|
||||
expect(parsed.data[0]).toHaveProperty('temp');
|
||||
} catch (error) {
|
||||
console.error('Timestamp test failed with error:', error);
|
||||
throw error;
|
||||
}
|
||||
});
|
||||
|
||||
test('daily_aggregation action with real API key returns aggregated data', async () => {
|
||||
if (!process.env.OPENWEATHER_API_KEY) {
|
||||
console.warn('Skipping real API test, no OPENWEATHER_API_KEY found.');
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
// Use yesterday's date for aggregation
|
||||
const yesterday = new Date();
|
||||
yesterday.setDate(yesterday.getDate() - 1);
|
||||
const dateStr = yesterday.toISOString().split('T')[0];
|
||||
|
||||
const result = await tool.call({
|
||||
action: 'daily_aggregation',
|
||||
city: 'London',
|
||||
date: dateStr,
|
||||
units: 'Celsius',
|
||||
});
|
||||
|
||||
console.log('Daily aggregation API response:', result);
|
||||
|
||||
const parsed = JSON.parse(result);
|
||||
expect(parsed).toHaveProperty('temperature');
|
||||
expect(parsed.temperature).toHaveProperty('morning');
|
||||
expect(parsed.temperature).toHaveProperty('afternoon');
|
||||
expect(parsed.temperature).toHaveProperty('evening');
|
||||
} catch (error) {
|
||||
console.error('Daily aggregation test failed with error:', error);
|
||||
throw error;
|
||||
}
|
||||
});
|
||||
|
||||
test('overview action with real API key returns weather summary', async () => {
|
||||
if (!process.env.OPENWEATHER_API_KEY) {
|
||||
console.warn('Skipping real API test, no OPENWEATHER_API_KEY found.');
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const result = await tool.call({
|
||||
action: 'overview',
|
||||
city: 'London',
|
||||
units: 'Celsius',
|
||||
});
|
||||
|
||||
console.log('Overview API response:', result);
|
||||
|
||||
const parsed = JSON.parse(result);
|
||||
expect(parsed).toHaveProperty('weather_overview');
|
||||
expect(typeof parsed.weather_overview).toBe('string');
|
||||
expect(parsed.weather_overview.length).toBeGreaterThan(0);
|
||||
expect(parsed).toHaveProperty('date');
|
||||
expect(parsed).toHaveProperty('units');
|
||||
expect(parsed.units).toBe('metric');
|
||||
} catch (error) {
|
||||
console.error('Overview test failed with error:', error);
|
||||
throw error;
|
||||
}
|
||||
});
|
||||
|
||||
test('different temperature units return correct values', async () => {
|
||||
if (!process.env.OPENWEATHER_API_KEY) {
|
||||
console.warn('Skipping real API test, no OPENWEATHER_API_KEY found.');
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
// Test Celsius
|
||||
let result = await tool.call({
|
||||
action: 'current_forecast',
|
||||
city: 'London',
|
||||
units: 'Celsius',
|
||||
});
|
||||
let parsed = JSON.parse(result);
|
||||
const celsiusTemp = parsed.current.temp;
|
||||
|
||||
// Test Kelvin
|
||||
result = await tool.call({
|
||||
action: 'current_forecast',
|
||||
city: 'London',
|
||||
units: 'Kelvin',
|
||||
});
|
||||
parsed = JSON.parse(result);
|
||||
const kelvinTemp = parsed.current.temp;
|
||||
|
||||
// Test Fahrenheit
|
||||
result = await tool.call({
|
||||
action: 'current_forecast',
|
||||
city: 'London',
|
||||
units: 'Fahrenheit',
|
||||
});
|
||||
parsed = JSON.parse(result);
|
||||
const fahrenheitTemp = parsed.current.temp;
|
||||
|
||||
// Verify temperature conversions are roughly correct
|
||||
// K = C + 273.15
|
||||
// F = (C * 9/5) + 32
|
||||
const celsiusToKelvin = Math.round(celsiusTemp + 273.15);
|
||||
const celsiusToFahrenheit = Math.round((celsiusTemp * 9) / 5 + 32);
|
||||
|
||||
console.log('Temperature comparisons:', {
|
||||
celsius: celsiusTemp,
|
||||
kelvin: kelvinTemp,
|
||||
fahrenheit: fahrenheitTemp,
|
||||
calculatedKelvin: celsiusToKelvin,
|
||||
calculatedFahrenheit: celsiusToFahrenheit,
|
||||
});
|
||||
|
||||
// Allow for some rounding differences
|
||||
expect(Math.abs(kelvinTemp - celsiusToKelvin)).toBeLessThanOrEqual(1);
|
||||
expect(Math.abs(fahrenheitTemp - celsiusToFahrenheit)).toBeLessThanOrEqual(1);
|
||||
} catch (error) {
|
||||
console.error('Temperature units test failed with error:', error);
|
||||
throw error;
|
||||
}
|
||||
});
|
||||
|
||||
test('language parameter returns localized data', async () => {
|
||||
if (!process.env.OPENWEATHER_API_KEY) {
|
||||
console.warn('Skipping real API test, no OPENWEATHER_API_KEY found.');
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
// Test with English
|
||||
let result = await tool.call({
|
||||
action: 'current_forecast',
|
||||
city: 'Paris',
|
||||
units: 'Celsius',
|
||||
lang: 'en',
|
||||
});
|
||||
let parsed = JSON.parse(result);
|
||||
const englishDescription = parsed.current.weather[0].description;
|
||||
|
||||
// Test with French
|
||||
result = await tool.call({
|
||||
action: 'current_forecast',
|
||||
city: 'Paris',
|
||||
units: 'Celsius',
|
||||
lang: 'fr',
|
||||
});
|
||||
parsed = JSON.parse(result);
|
||||
const frenchDescription = parsed.current.weather[0].description;
|
||||
|
||||
console.log('Language comparison:', {
|
||||
english: englishDescription,
|
||||
french: frenchDescription,
|
||||
});
|
||||
|
||||
// Verify descriptions are different (indicating translation worked)
|
||||
expect(englishDescription).not.toBe(frenchDescription);
|
||||
} catch (error) {
|
||||
console.error('Language test failed with error:', error);
|
||||
throw error;
|
||||
}
|
||||
});
|
||||
});
|
||||
358
api/app/clients/tools/structured/specs/openweather.test.js
Normal file
358
api/app/clients/tools/structured/specs/openweather.test.js
Normal file
@ -0,0 +1,358 @@
|
||||
// __tests__/openweather.test.js
|
||||
const OpenWeather = require('../OpenWeather');
|
||||
const fetch = require('node-fetch');
|
||||
|
||||
// Mock environment variable
|
||||
process.env.OPENWEATHER_API_KEY = 'test-api-key';
|
||||
|
||||
// Mock the fetch function globally
|
||||
jest.mock('node-fetch', () => jest.fn());
|
||||
|
||||
describe('OpenWeather Tool', () => {
|
||||
let tool;
|
||||
|
||||
beforeAll(() => {
|
||||
tool = new OpenWeather();
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
fetch.mockReset();
|
||||
});
|
||||
|
||||
test('action=help returns help instructions', async () => {
|
||||
const result = await tool.call({
|
||||
action: 'help',
|
||||
});
|
||||
|
||||
expect(typeof result).toBe('string');
|
||||
const parsed = JSON.parse(result);
|
||||
expect(parsed.title).toBe('OpenWeather One Call API 3.0 Help');
|
||||
});
|
||||
|
||||
test('current_forecast with a city and successful geocoding + forecast', async () => {
|
||||
// Mock geocoding response
|
||||
fetch.mockImplementationOnce((url) => {
|
||||
if (url.includes('geo/1.0/direct')) {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => [{ lat: 35.9606, lon: -83.9207 }],
|
||||
});
|
||||
}
|
||||
return Promise.reject('Unexpected fetch call for geocoding');
|
||||
});
|
||||
|
||||
// Mock forecast response
|
||||
fetch.mockImplementationOnce(() =>
|
||||
Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => ({
|
||||
current: { temp: 293.15, feels_like: 295.15 },
|
||||
daily: [{ temp: { day: 293.15, night: 283.15 } }],
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
const result = await tool.call({
|
||||
action: 'current_forecast',
|
||||
city: 'Knoxville, Tennessee',
|
||||
units: 'Kelvin',
|
||||
});
|
||||
|
||||
const parsed = JSON.parse(result);
|
||||
expect(parsed.current.temp).toBe(293);
|
||||
expect(parsed.current.feels_like).toBe(295);
|
||||
expect(parsed.daily[0].temp.day).toBe(293);
|
||||
expect(parsed.daily[0].temp.night).toBe(283);
|
||||
});
|
||||
|
||||
test('timestamp action with valid date returns mocked historical data', async () => {
|
||||
// Mock geocoding response
|
||||
fetch.mockImplementationOnce((url) => {
|
||||
if (url.includes('geo/1.0/direct')) {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => [{ lat: 35.9606, lon: -83.9207 }],
|
||||
});
|
||||
}
|
||||
return Promise.reject('Unexpected fetch call for geocoding');
|
||||
});
|
||||
|
||||
// Mock historical weather response
|
||||
fetch.mockImplementationOnce(() =>
|
||||
Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => ({
|
||||
data: [
|
||||
{
|
||||
dt: 1583280000,
|
||||
temp: 283.15,
|
||||
feels_like: 280.15,
|
||||
humidity: 75,
|
||||
weather: [{ description: 'clear sky' }],
|
||||
},
|
||||
],
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
const result = await tool.call({
|
||||
action: 'timestamp',
|
||||
city: 'Knoxville, Tennessee',
|
||||
date: '2020-03-04',
|
||||
units: 'Kelvin',
|
||||
});
|
||||
|
||||
const parsed = JSON.parse(result);
|
||||
expect(parsed.data[0].temp).toBe(283);
|
||||
expect(parsed.data[0].feels_like).toBe(280);
|
||||
});
|
||||
|
||||
test('daily_aggregation action returns aggregated weather data', async () => {
|
||||
// Mock geocoding response
|
||||
fetch.mockImplementationOnce((url) => {
|
||||
if (url.includes('geo/1.0/direct')) {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => [{ lat: 35.9606, lon: -83.9207 }],
|
||||
});
|
||||
}
|
||||
return Promise.reject('Unexpected fetch call for geocoding');
|
||||
});
|
||||
|
||||
// Mock daily aggregation response
|
||||
fetch.mockImplementationOnce(() =>
|
||||
Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => ({
|
||||
date: '2020-03-04',
|
||||
temperature: {
|
||||
morning: 283.15,
|
||||
afternoon: 293.15,
|
||||
evening: 288.15,
|
||||
},
|
||||
humidity: {
|
||||
morning: 75,
|
||||
afternoon: 60,
|
||||
evening: 70,
|
||||
},
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
const result = await tool.call({
|
||||
action: 'daily_aggregation',
|
||||
city: 'Knoxville, Tennessee',
|
||||
date: '2020-03-04',
|
||||
units: 'Kelvin',
|
||||
});
|
||||
|
||||
const parsed = JSON.parse(result);
|
||||
expect(parsed.temperature.morning).toBe(283);
|
||||
expect(parsed.temperature.afternoon).toBe(293);
|
||||
expect(parsed.temperature.evening).toBe(288);
|
||||
});
|
||||
|
||||
test('overview action returns weather summary', async () => {
|
||||
// Mock geocoding response
|
||||
fetch.mockImplementationOnce((url) => {
|
||||
if (url.includes('geo/1.0/direct')) {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => [{ lat: 35.9606, lon: -83.9207 }],
|
||||
});
|
||||
}
|
||||
return Promise.reject('Unexpected fetch call for geocoding');
|
||||
});
|
||||
|
||||
// Mock overview response
|
||||
fetch.mockImplementationOnce(() =>
|
||||
Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => ({
|
||||
date: '2024-01-07',
|
||||
lat: 35.9606,
|
||||
lon: -83.9207,
|
||||
tz: '+00:00',
|
||||
units: 'metric',
|
||||
weather_overview:
|
||||
'Currently, the temperature is 2°C with a real feel of -2°C. The sky is clear with moderate wind.',
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
const result = await tool.call({
|
||||
action: 'overview',
|
||||
city: 'Knoxville, Tennessee',
|
||||
units: 'Celsius',
|
||||
});
|
||||
|
||||
const parsed = JSON.parse(result);
|
||||
expect(parsed).toHaveProperty('weather_overview');
|
||||
expect(typeof parsed.weather_overview).toBe('string');
|
||||
expect(parsed.weather_overview.length).toBeGreaterThan(0);
|
||||
expect(parsed).toHaveProperty('date');
|
||||
expect(parsed).toHaveProperty('units');
|
||||
expect(parsed.units).toBe('metric');
|
||||
});
|
||||
|
||||
test('temperature units are correctly converted', async () => {
|
||||
// Mock geocoding response for all three calls
|
||||
const geocodingMock = Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => [{ lat: 35.9606, lon: -83.9207 }],
|
||||
});
|
||||
|
||||
// Mock weather response for Kelvin
|
||||
const kelvinMock = Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => ({
|
||||
current: { temp: 293.15 },
|
||||
}),
|
||||
});
|
||||
|
||||
// Mock weather response for Celsius
|
||||
const celsiusMock = Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => ({
|
||||
current: { temp: 20 },
|
||||
}),
|
||||
});
|
||||
|
||||
// Mock weather response for Fahrenheit
|
||||
const fahrenheitMock = Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => ({
|
||||
current: { temp: 68 },
|
||||
}),
|
||||
});
|
||||
|
||||
// Test Kelvin
|
||||
fetch.mockImplementationOnce(() => geocodingMock).mockImplementationOnce(() => kelvinMock);
|
||||
|
||||
let result = await tool.call({
|
||||
action: 'current_forecast',
|
||||
city: 'Knoxville, Tennessee',
|
||||
units: 'Kelvin',
|
||||
});
|
||||
let parsed = JSON.parse(result);
|
||||
expect(parsed.current.temp).toBe(293);
|
||||
|
||||
// Test Celsius
|
||||
fetch.mockImplementationOnce(() => geocodingMock).mockImplementationOnce(() => celsiusMock);
|
||||
|
||||
result = await tool.call({
|
||||
action: 'current_forecast',
|
||||
city: 'Knoxville, Tennessee',
|
||||
units: 'Celsius',
|
||||
});
|
||||
parsed = JSON.parse(result);
|
||||
expect(parsed.current.temp).toBe(20);
|
||||
|
||||
// Test Fahrenheit
|
||||
fetch.mockImplementationOnce(() => geocodingMock).mockImplementationOnce(() => fahrenheitMock);
|
||||
|
||||
result = await tool.call({
|
||||
action: 'current_forecast',
|
||||
city: 'Knoxville, Tennessee',
|
||||
units: 'Fahrenheit',
|
||||
});
|
||||
parsed = JSON.parse(result);
|
||||
expect(parsed.current.temp).toBe(68);
|
||||
});
|
||||
|
||||
test('timestamp action without a date returns an error message', async () => {
|
||||
const result = await tool.call({
|
||||
action: 'timestamp',
|
||||
lat: 35.9606,
|
||||
lon: -83.9207,
|
||||
});
|
||||
expect(result).toMatch(
|
||||
/Error: For timestamp action, a 'date' in YYYY-MM-DD format is required./,
|
||||
);
|
||||
});
|
||||
|
||||
test('daily_aggregation action without a date returns an error message', async () => {
|
||||
const result = await tool.call({
|
||||
action: 'daily_aggregation',
|
||||
lat: 35.9606,
|
||||
lon: -83.9207,
|
||||
});
|
||||
expect(result).toMatch(/Error: date \(YYYY-MM-DD\) is required for daily_aggregation action./);
|
||||
});
|
||||
|
||||
test('unknown action returns an error due to schema validation', async () => {
|
||||
await expect(
|
||||
tool.call({
|
||||
action: 'unknown_action',
|
||||
}),
|
||||
).rejects.toThrow(/Received tool input did not match expected schema/);
|
||||
});
|
||||
|
||||
test('geocoding failure returns a descriptive error', async () => {
|
||||
fetch.mockImplementationOnce(() =>
|
||||
Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => [],
|
||||
}),
|
||||
);
|
||||
|
||||
const result = await tool.call({
|
||||
action: 'current_forecast',
|
||||
city: 'NowhereCity',
|
||||
});
|
||||
expect(result).toMatch(/Error: Could not find coordinates for city: NowhereCity/);
|
||||
});
|
||||
|
||||
test('API request failure returns an error', async () => {
|
||||
// Mock geocoding success
|
||||
fetch.mockImplementationOnce(() =>
|
||||
Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => [{ lat: 35.9606, lon: -83.9207 }],
|
||||
}),
|
||||
);
|
||||
|
||||
// Mock weather request failure
|
||||
fetch.mockImplementationOnce(() =>
|
||||
Promise.resolve({
|
||||
ok: false,
|
||||
status: 404,
|
||||
json: async () => ({ message: 'Not found' }),
|
||||
}),
|
||||
);
|
||||
|
||||
const result = await tool.call({
|
||||
action: 'current_forecast',
|
||||
city: 'Knoxville, Tennessee',
|
||||
});
|
||||
expect(result).toMatch(/Error: OpenWeather API request failed with status 404: Not found/);
|
||||
});
|
||||
|
||||
test('invalid date format returns an error', async () => {
|
||||
// Mock geocoding response first
|
||||
fetch.mockImplementationOnce((url) => {
|
||||
if (url.includes('geo/1.0/direct')) {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => [{ lat: 35.9606, lon: -83.9207 }],
|
||||
});
|
||||
}
|
||||
return Promise.reject('Unexpected fetch call for geocoding');
|
||||
});
|
||||
|
||||
// Mock timestamp API response
|
||||
fetch.mockImplementationOnce((url) => {
|
||||
if (url.includes('onecall/timemachine')) {
|
||||
throw new Error('Invalid date format. Expected YYYY-MM-DD.');
|
||||
}
|
||||
return Promise.reject('Unexpected fetch call');
|
||||
});
|
||||
|
||||
const result = await tool.call({
|
||||
action: 'timestamp',
|
||||
city: 'Knoxville, Tennessee',
|
||||
date: '03-04-2020', // Wrong format
|
||||
});
|
||||
expect(result).toMatch(/Error: Invalid date format. Expected YYYY-MM-DD./);
|
||||
});
|
||||
});
|
||||
196
api/app/clients/tools/util/fileSearch.js
Normal file
196
api/app/clients/tools/util/fileSearch.js
Normal file
@ -0,0 +1,196 @@
|
||||
const { z } = require('zod');
|
||||
const axios = require('axios');
|
||||
const { tool } = require('@langchain/core/tools');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { generateShortLivedToken } = require('@librechat/api');
|
||||
const { Tools, EToolResources } = require('librechat-data-provider');
|
||||
const { filterFilesByAgentAccess } = require('~/server/services/Files/permissions');
|
||||
const { getFiles } = require('~/models');
|
||||
|
||||
/**
|
||||
*
|
||||
* @param {Object} options
|
||||
* @param {ServerRequest} options.req
|
||||
* @param {Agent['tool_resources']} options.tool_resources
|
||||
* @param {string} [options.agentId] - The agent ID for file access control
|
||||
* @returns {Promise<{
|
||||
* files: Array<{ file_id: string; filename: string }>,
|
||||
* toolContext: string
|
||||
* }>}
|
||||
*/
|
||||
const primeFiles = async (options) => {
|
||||
const { tool_resources, req, agentId } = options;
|
||||
const file_ids = tool_resources?.[EToolResources.file_search]?.file_ids ?? [];
|
||||
const agentResourceIds = new Set(file_ids);
|
||||
const resourceFiles = tool_resources?.[EToolResources.file_search]?.files ?? [];
|
||||
|
||||
// Get all files first
|
||||
const allFiles = (await getFiles({ file_id: { $in: file_ids } }, null, { text: 0 })) ?? [];
|
||||
|
||||
// Filter by access if user and agent are provided
|
||||
let dbFiles;
|
||||
if (req?.user?.id && agentId) {
|
||||
dbFiles = await filterFilesByAgentAccess({
|
||||
files: allFiles,
|
||||
userId: req.user.id,
|
||||
role: req.user.role,
|
||||
agentId,
|
||||
});
|
||||
} else {
|
||||
dbFiles = allFiles;
|
||||
}
|
||||
|
||||
dbFiles = dbFiles.concat(resourceFiles);
|
||||
|
||||
let toolContext = `- Note: Semantic search is available through the ${Tools.file_search} tool but no files are currently loaded. Request the user to upload documents to search through.`;
|
||||
|
||||
const files = [];
|
||||
for (let i = 0; i < dbFiles.length; i++) {
|
||||
const file = dbFiles[i];
|
||||
if (!file) {
|
||||
continue;
|
||||
}
|
||||
if (i === 0) {
|
||||
toolContext = `- Note: Use the ${Tools.file_search} tool to find relevant information within:`;
|
||||
}
|
||||
toolContext += `\n\t- ${file.filename}${
|
||||
agentResourceIds.has(file.file_id) ? '' : ' (just attached by user)'
|
||||
}`;
|
||||
files.push({
|
||||
file_id: file.file_id,
|
||||
filename: file.filename,
|
||||
});
|
||||
}
|
||||
|
||||
return { files, toolContext };
|
||||
};
|
||||
|
||||
/**
|
||||
*
|
||||
* @param {Object} options
|
||||
* @param {string} options.userId
|
||||
* @param {Array<{ file_id: string; filename: string }>} options.files
|
||||
* @param {string} [options.entity_id]
|
||||
* @param {boolean} [options.fileCitations=false] - Whether to include citation instructions
|
||||
* @returns
|
||||
*/
|
||||
const createFileSearchTool = async ({ userId, files, entity_id, fileCitations = false }) => {
|
||||
return tool(
|
||||
async ({ query }) => {
|
||||
if (files.length === 0) {
|
||||
return ['No files to search. Instruct the user to add files for the search.', undefined];
|
||||
}
|
||||
const jwtToken = generateShortLivedToken(userId);
|
||||
if (!jwtToken) {
|
||||
return ['There was an error authenticating the file search request.', undefined];
|
||||
}
|
||||
|
||||
/**
|
||||
* @param {import('librechat-data-provider').TFile} file
|
||||
* @returns {{ file_id: string, query: string, k: number, entity_id?: string }}
|
||||
*/
|
||||
const createQueryBody = (file) => {
|
||||
const body = {
|
||||
file_id: file.file_id,
|
||||
query,
|
||||
k: 5,
|
||||
};
|
||||
if (!entity_id) {
|
||||
return body;
|
||||
}
|
||||
body.entity_id = entity_id;
|
||||
logger.debug(`[${Tools.file_search}] RAG API /query body`, body);
|
||||
return body;
|
||||
};
|
||||
|
||||
const queryPromises = files.map((file) =>
|
||||
axios
|
||||
.post(`${process.env.RAG_API_URL}/query`, createQueryBody(file), {
|
||||
headers: {
|
||||
Authorization: `Bearer ${jwtToken}`,
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
})
|
||||
.catch((error) => {
|
||||
logger.error('Error encountered in `file_search` while querying file:', error);
|
||||
return null;
|
||||
}),
|
||||
);
|
||||
|
||||
const results = await Promise.all(queryPromises);
|
||||
const validResults = results.filter((result) => result !== null);
|
||||
|
||||
if (validResults.length === 0) {
|
||||
return ['No results found or errors occurred while searching the files.', undefined];
|
||||
}
|
||||
|
||||
const formattedResults = validResults
|
||||
.flatMap((result, fileIndex) =>
|
||||
result.data.map(([docInfo, distance]) => ({
|
||||
filename: docInfo.metadata.source.split('/').pop(),
|
||||
content: docInfo.page_content,
|
||||
distance,
|
||||
file_id: files[fileIndex]?.file_id,
|
||||
page: docInfo.metadata.page || null,
|
||||
})),
|
||||
)
|
||||
.sort((a, b) => a.distance - b.distance)
|
||||
.slice(0, 10);
|
||||
|
||||
if (formattedResults.length === 0) {
|
||||
return [
|
||||
'No content found in the files. The files may not have been processed correctly or you may need to refine your query.',
|
||||
undefined,
|
||||
];
|
||||
}
|
||||
|
||||
const formattedString = formattedResults
|
||||
.map(
|
||||
(result, index) =>
|
||||
`File: ${result.filename}${
|
||||
fileCitations ? `\nAnchor: \\ue202turn0file${index} (${result.filename})` : ''
|
||||
}\nRelevance: ${(1.0 - result.distance).toFixed(4)}\nContent: ${result.content}\n`,
|
||||
)
|
||||
.join('\n---\n');
|
||||
|
||||
const sources = formattedResults.map((result) => ({
|
||||
type: 'file',
|
||||
fileId: result.file_id,
|
||||
content: result.content,
|
||||
fileName: result.filename,
|
||||
relevance: 1.0 - result.distance,
|
||||
pages: result.page ? [result.page] : [],
|
||||
pageRelevance: result.page ? { [result.page]: 1.0 - result.distance } : {},
|
||||
}));
|
||||
|
||||
return [formattedString, { [Tools.file_search]: { sources, fileCitations } }];
|
||||
},
|
||||
{
|
||||
name: Tools.file_search,
|
||||
responseFormat: 'content_and_artifact',
|
||||
description: `Performs semantic search across attached "${Tools.file_search}" documents using natural language queries. This tool analyzes the content of uploaded files to find relevant information, quotes, and passages that best match your query. Use this to extract specific information or find relevant sections within the available documents.${
|
||||
fileCitations
|
||||
? `
|
||||
|
||||
**CITE FILE SEARCH RESULTS:**
|
||||
Use the EXACT anchor markers shown below (copy them verbatim) immediately after statements derived from file content. Reference the filename in your text:
|
||||
- File citation: "The document.pdf states that... \\ue202turn0file0"
|
||||
- Page reference: "According to report.docx... \\ue202turn0file1"
|
||||
- Multi-file: "Multiple sources confirm... \\ue200\\ue202turn0file0\\ue202turn0file1\\ue201"
|
||||
|
||||
**CRITICAL:** Output these escape sequences EXACTLY as shown (e.g., \\ue202turn0file0). Do NOT substitute with other characters like † or similar symbols.
|
||||
**ALWAYS mention the filename in your text before the citation marker. NEVER use markdown links or footnotes.**`
|
||||
: ''
|
||||
}`,
|
||||
schema: z.object({
|
||||
query: z
|
||||
.string()
|
||||
.describe(
|
||||
"A natural language query to search for relevant information in the files. Be specific and use keywords related to the information you're looking for. The query will be used for semantic similarity matching against the file contents.",
|
||||
),
|
||||
}),
|
||||
},
|
||||
);
|
||||
};
|
||||
|
||||
module.exports = { createFileSearchTool, primeFiles };
|
||||
504
api/app/clients/tools/util/handleTools.js
Normal file
504
api/app/clients/tools/util/handleTools.js
Normal file
@ -0,0 +1,504 @@
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const {
|
||||
EnvVar,
|
||||
Calculator,
|
||||
createSearchTool,
|
||||
createCodeExecutionTool,
|
||||
} = require('@librechat/agents');
|
||||
const {
|
||||
checkAccess,
|
||||
createSafeUser,
|
||||
mcpToolPattern,
|
||||
loadWebSearchAuth,
|
||||
buildImageToolContext,
|
||||
} = require('@librechat/api');
|
||||
const { getMCPServersRegistry } = require('~/config');
|
||||
const {
|
||||
Tools,
|
||||
Constants,
|
||||
Permissions,
|
||||
EToolResources,
|
||||
PermissionTypes,
|
||||
replaceSpecialVars,
|
||||
} = require('librechat-data-provider');
|
||||
const {
|
||||
availableTools,
|
||||
manifestToolMap,
|
||||
// Basic Tools
|
||||
GoogleSearchAPI,
|
||||
// Structured Tools
|
||||
DALLE3,
|
||||
FluxAPI,
|
||||
OpenWeather,
|
||||
StructuredSD,
|
||||
StructuredACS,
|
||||
TraversaalSearch,
|
||||
StructuredWolfram,
|
||||
TavilySearchResults,
|
||||
createGeminiImageTool,
|
||||
createOpenAIImageTools,
|
||||
} = require('../');
|
||||
const { primeFiles: primeCodeFiles } = require('~/server/services/Files/Code/process');
|
||||
const { createFileSearchTool, primeFiles: primeSearchFiles } = require('./fileSearch');
|
||||
const { getUserPluginAuthValue } = require('~/server/services/PluginService');
|
||||
const { createMCPTool, createMCPTools } = require('~/server/services/MCP');
|
||||
const { loadAuthValues } = require('~/server/services/Tools/credentials');
|
||||
const { getMCPServerTools } = require('~/server/services/Config');
|
||||
const { getRoleByName } = require('~/models/Role');
|
||||
|
||||
/**
|
||||
* Validates the availability and authentication of tools for a user based on environment variables or user-specific plugin authentication values.
|
||||
* Tools without required authentication or with valid authentication are considered valid.
|
||||
*
|
||||
* @param {Object} user The user object for whom to validate tool access.
|
||||
* @param {Array<string>} tools An array of tool identifiers to validate. Defaults to an empty array.
|
||||
* @returns {Promise<Array<string>>} A promise that resolves to an array of valid tool identifiers.
|
||||
*/
|
||||
const validateTools = async (user, tools = []) => {
|
||||
try {
|
||||
const validToolsSet = new Set(tools);
|
||||
const availableToolsToValidate = availableTools.filter((tool) =>
|
||||
validToolsSet.has(tool.pluginKey),
|
||||
);
|
||||
|
||||
/**
|
||||
* Validates the credentials for a given auth field or set of alternate auth fields for a tool.
|
||||
* If valid admin or user authentication is found, the function returns early. Otherwise, it removes the tool from the set of valid tools.
|
||||
*
|
||||
* @param {string} authField The authentication field or fields (separated by "||" for alternates) to validate.
|
||||
* @param {string} toolName The identifier of the tool being validated.
|
||||
*/
|
||||
const validateCredentials = async (authField, toolName) => {
|
||||
const fields = authField.split('||');
|
||||
for (const field of fields) {
|
||||
const adminAuth = process.env[field];
|
||||
if (adminAuth && adminAuth.length > 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
let userAuth = null;
|
||||
try {
|
||||
userAuth = await getUserPluginAuthValue(user, field);
|
||||
} catch (err) {
|
||||
if (field === fields[fields.length - 1] && !userAuth) {
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
if (userAuth && userAuth.length > 0) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
validToolsSet.delete(toolName);
|
||||
};
|
||||
|
||||
for (const tool of availableToolsToValidate) {
|
||||
if (!tool.authConfig || tool.authConfig.length === 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
for (const auth of tool.authConfig) {
|
||||
await validateCredentials(auth.authField, tool.pluginKey);
|
||||
}
|
||||
}
|
||||
|
||||
return Array.from(validToolsSet.values());
|
||||
} catch (err) {
|
||||
logger.error('[validateTools] There was a problem validating tools', err);
|
||||
throw new Error(err);
|
||||
}
|
||||
};
|
||||
|
||||
/** @typedef {typeof import('@langchain/core/tools').Tool} ToolConstructor */
|
||||
/** @typedef {import('@langchain/core/tools').Tool} Tool */
|
||||
|
||||
/**
|
||||
* Initializes a tool with authentication values for the given user, supporting alternate authentication fields.
|
||||
* Authentication fields can have alternates separated by "||", and the first defined variable will be used.
|
||||
*
|
||||
* @param {string} userId The user ID for which the tool is being loaded.
|
||||
* @param {Array<string>} authFields Array of strings representing the authentication fields. Supports alternate fields delimited by "||".
|
||||
* @param {ToolConstructor} ToolConstructor The constructor function for the tool to be initialized.
|
||||
* @param {Object} options Optional parameters to be passed to the tool constructor alongside authentication values.
|
||||
* @returns {() => Promise<Tool>} An Async function that, when called, asynchronously initializes and returns an instance of the tool with authentication.
|
||||
*/
|
||||
const loadToolWithAuth = (userId, authFields, ToolConstructor, options = {}) => {
|
||||
return async function () {
|
||||
const authValues = await loadAuthValues({ userId, authFields });
|
||||
return new ToolConstructor({ ...options, ...authValues, userId });
|
||||
};
|
||||
};
|
||||
|
||||
/**
|
||||
* @param {string} toolKey
|
||||
* @returns {Array<string>}
|
||||
*/
|
||||
const getAuthFields = (toolKey) => {
|
||||
return manifestToolMap[toolKey]?.authConfig.map((auth) => auth.authField) ?? [];
|
||||
};
|
||||
|
||||
/**
|
||||
*
|
||||
* @param {object} params
|
||||
* @param {string} params.user
|
||||
* @param {Record<string, Record<string, string>>} [object.userMCPAuthMap]
|
||||
* @param {AbortSignal} [object.signal]
|
||||
* @param {Pick<Agent, 'id' | 'provider' | 'model'>} [params.agent]
|
||||
* @param {string} [params.model]
|
||||
* @param {EModelEndpoint} [params.endpoint]
|
||||
* @param {LoadToolOptions} [params.options]
|
||||
* @param {boolean} [params.useSpecs]
|
||||
* @param {Array<string>} params.tools
|
||||
* @param {boolean} [params.functions]
|
||||
* @param {boolean} [params.returnMap]
|
||||
* @param {AppConfig['webSearch']} [params.webSearch]
|
||||
* @param {AppConfig['fileStrategy']} [params.fileStrategy]
|
||||
* @param {AppConfig['imageOutputType']} [params.imageOutputType]
|
||||
* @returns {Promise<{ loadedTools: Tool[], toolContextMap: Object<string, any> } | Record<string,Tool>>}
|
||||
*/
|
||||
const loadTools = async ({
|
||||
user,
|
||||
agent,
|
||||
model,
|
||||
signal,
|
||||
endpoint,
|
||||
userMCPAuthMap,
|
||||
tools = [],
|
||||
options = {},
|
||||
functions = true,
|
||||
returnMap = false,
|
||||
webSearch,
|
||||
fileStrategy,
|
||||
imageOutputType,
|
||||
}) => {
|
||||
const toolConstructors = {
|
||||
flux: FluxAPI,
|
||||
calculator: Calculator,
|
||||
google: GoogleSearchAPI,
|
||||
open_weather: OpenWeather,
|
||||
wolfram: StructuredWolfram,
|
||||
'stable-diffusion': StructuredSD,
|
||||
'azure-ai-search': StructuredACS,
|
||||
traversaal_search: TraversaalSearch,
|
||||
tavily_search_results_json: TavilySearchResults,
|
||||
};
|
||||
|
||||
const customConstructors = {
|
||||
image_gen_oai: async (toolContextMap) => {
|
||||
const authFields = getAuthFields('image_gen_oai');
|
||||
const authValues = await loadAuthValues({ userId: user, authFields });
|
||||
const imageFiles = options.tool_resources?.[EToolResources.image_edit]?.files ?? [];
|
||||
const toolContext = buildImageToolContext({
|
||||
imageFiles,
|
||||
toolName: `${EToolResources.image_edit}_oai`,
|
||||
contextDescription: 'image editing',
|
||||
});
|
||||
if (toolContext) {
|
||||
toolContextMap.image_edit_oai = toolContext;
|
||||
}
|
||||
return createOpenAIImageTools({
|
||||
...authValues,
|
||||
isAgent: !!agent,
|
||||
req: options.req,
|
||||
imageOutputType,
|
||||
fileStrategy,
|
||||
imageFiles,
|
||||
});
|
||||
},
|
||||
gemini_image_gen: async (toolContextMap) => {
|
||||
const authFields = getAuthFields('gemini_image_gen');
|
||||
const authValues = await loadAuthValues({ userId: user, authFields });
|
||||
const imageFiles = options.tool_resources?.[EToolResources.image_edit]?.files ?? [];
|
||||
const toolContext = buildImageToolContext({
|
||||
imageFiles,
|
||||
toolName: 'gemini_image_gen',
|
||||
contextDescription: 'image context',
|
||||
});
|
||||
if (toolContext) {
|
||||
toolContextMap.gemini_image_gen = toolContext;
|
||||
}
|
||||
return createGeminiImageTool({
|
||||
...authValues,
|
||||
isAgent: !!agent,
|
||||
req: options.req,
|
||||
imageFiles,
|
||||
processFileURL: options.processFileURL,
|
||||
userId: user,
|
||||
fileStrategy,
|
||||
});
|
||||
},
|
||||
};
|
||||
|
||||
const requestedTools = {};
|
||||
|
||||
if (functions === true) {
|
||||
toolConstructors.dalle = DALLE3;
|
||||
}
|
||||
|
||||
/** @type {ImageGenOptions} */
|
||||
const imageGenOptions = {
|
||||
isAgent: !!agent,
|
||||
req: options.req,
|
||||
fileStrategy,
|
||||
processFileURL: options.processFileURL,
|
||||
returnMetadata: options.returnMetadata,
|
||||
uploadImageBuffer: options.uploadImageBuffer,
|
||||
};
|
||||
|
||||
const toolOptions = {
|
||||
flux: imageGenOptions,
|
||||
dalle: imageGenOptions,
|
||||
'stable-diffusion': imageGenOptions,
|
||||
gemini_image_gen: imageGenOptions,
|
||||
};
|
||||
|
||||
/** @type {Record<string, string>} */
|
||||
const toolContextMap = {};
|
||||
const requestedMCPTools = {};
|
||||
|
||||
for (const tool of tools) {
|
||||
if (tool === Tools.execute_code) {
|
||||
requestedTools[tool] = async () => {
|
||||
const authValues = await loadAuthValues({
|
||||
userId: user,
|
||||
authFields: [EnvVar.CODE_API_KEY],
|
||||
});
|
||||
const codeApiKey = authValues[EnvVar.CODE_API_KEY];
|
||||
const { files, toolContext } = await primeCodeFiles(
|
||||
{
|
||||
...options,
|
||||
agentId: agent?.id,
|
||||
},
|
||||
codeApiKey,
|
||||
);
|
||||
if (toolContext) {
|
||||
toolContextMap[tool] = toolContext;
|
||||
}
|
||||
const CodeExecutionTool = createCodeExecutionTool({
|
||||
user_id: user,
|
||||
files,
|
||||
...authValues,
|
||||
});
|
||||
CodeExecutionTool.apiKey = codeApiKey;
|
||||
return CodeExecutionTool;
|
||||
};
|
||||
continue;
|
||||
} else if (tool === Tools.file_search) {
|
||||
requestedTools[tool] = async () => {
|
||||
const { files, toolContext } = await primeSearchFiles({
|
||||
...options,
|
||||
agentId: agent?.id,
|
||||
});
|
||||
if (toolContext) {
|
||||
toolContextMap[tool] = toolContext;
|
||||
}
|
||||
|
||||
/** @type {boolean | undefined} Check if user has FILE_CITATIONS permission */
|
||||
let fileCitations;
|
||||
if (fileCitations == null && options.req?.user != null) {
|
||||
try {
|
||||
fileCitations = await checkAccess({
|
||||
user: options.req.user,
|
||||
permissionType: PermissionTypes.FILE_CITATIONS,
|
||||
permissions: [Permissions.USE],
|
||||
getRoleByName,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error('[handleTools] FILE_CITATIONS permission check failed:', error);
|
||||
fileCitations = false;
|
||||
}
|
||||
}
|
||||
|
||||
return createFileSearchTool({
|
||||
userId: user,
|
||||
files,
|
||||
entity_id: agent?.id,
|
||||
fileCitations,
|
||||
});
|
||||
};
|
||||
continue;
|
||||
} else if (tool === Tools.web_search) {
|
||||
const result = await loadWebSearchAuth({
|
||||
userId: user,
|
||||
loadAuthValues,
|
||||
webSearchConfig: webSearch,
|
||||
});
|
||||
const { onSearchResults, onGetHighlights } = options?.[Tools.web_search] ?? {};
|
||||
requestedTools[tool] = async () => {
|
||||
toolContextMap[tool] = `# \`${tool}\`:
|
||||
Current Date & Time: ${replaceSpecialVars({ text: '{{iso_datetime}}' })}
|
||||
|
||||
**Execute immediately without preface.** After search, provide a brief summary addressing the query directly, then structure your response with clear Markdown formatting (## headers, lists, tables). Cite sources properly, tailor tone to query type, and provide comprehensive details.
|
||||
|
||||
**CITATION FORMAT - UNICODE ESCAPE SEQUENCES ONLY:**
|
||||
Use these EXACT escape sequences (copy verbatim): \\ue202 (before each anchor), \\ue200 (group start), \\ue201 (group end), \\ue203 (highlight start), \\ue204 (highlight end)
|
||||
|
||||
Anchor pattern: \\ue202turn{N}{type}{index} where N=turn number, type=search|news|image|ref, index=0,1,2...
|
||||
|
||||
**Examples (copy these exactly):**
|
||||
- Single: "Statement.\\ue202turn0search0"
|
||||
- Multiple: "Statement.\\ue202turn0search0\\ue202turn0news1"
|
||||
- Group: "Statement. \\ue200\\ue202turn0search0\\ue202turn0news1\\ue201"
|
||||
- Highlight: "\\ue203Cited text.\\ue204\\ue202turn0search0"
|
||||
- Image: "See photo\\ue202turn0image0."
|
||||
|
||||
**CRITICAL:** Output escape sequences EXACTLY as shown. Do NOT substitute with † or other symbols. Place anchors AFTER punctuation. Cite every non-obvious fact/quote. NEVER use markdown links, [1], footnotes, or HTML tags.`.trim();
|
||||
return createSearchTool({
|
||||
...result.authResult,
|
||||
onSearchResults,
|
||||
onGetHighlights,
|
||||
logger,
|
||||
});
|
||||
};
|
||||
continue;
|
||||
} else if (tool && mcpToolPattern.test(tool)) {
|
||||
const [toolName, serverName] = tool.split(Constants.mcp_delimiter);
|
||||
if (toolName === Constants.mcp_server) {
|
||||
/** Placeholder used for UI purposes */
|
||||
continue;
|
||||
}
|
||||
const serverConfig = serverName
|
||||
? await getMCPServersRegistry().getServerConfig(serverName, user)
|
||||
: null;
|
||||
if (!serverConfig) {
|
||||
logger.warn(
|
||||
`MCP server "${serverName}" for "${toolName}" tool is not configured${agent?.id != null && agent.id ? ` but attached to "${agent.id}"` : ''}`,
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if (toolName === Constants.mcp_all) {
|
||||
requestedMCPTools[serverName] = [
|
||||
{
|
||||
type: 'all',
|
||||
serverName,
|
||||
config: serverConfig,
|
||||
},
|
||||
];
|
||||
continue;
|
||||
}
|
||||
|
||||
requestedMCPTools[serverName] = requestedMCPTools[serverName] || [];
|
||||
requestedMCPTools[serverName].push({
|
||||
type: 'single',
|
||||
toolKey: tool,
|
||||
serverName,
|
||||
config: serverConfig,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
if (customConstructors[tool]) {
|
||||
requestedTools[tool] = async () => customConstructors[tool](toolContextMap);
|
||||
continue;
|
||||
}
|
||||
|
||||
if (toolConstructors[tool]) {
|
||||
const options = toolOptions[tool] || {};
|
||||
const toolInstance = loadToolWithAuth(
|
||||
user,
|
||||
getAuthFields(tool),
|
||||
toolConstructors[tool],
|
||||
options,
|
||||
);
|
||||
requestedTools[tool] = toolInstance;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
if (returnMap) {
|
||||
return requestedTools;
|
||||
}
|
||||
|
||||
const toolPromises = [];
|
||||
for (const tool of tools) {
|
||||
const validTool = requestedTools[tool];
|
||||
if (validTool) {
|
||||
toolPromises.push(
|
||||
validTool().catch((error) => {
|
||||
logger.error(`Error loading tool ${tool}:`, error);
|
||||
return null;
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
const loadedTools = (await Promise.all(toolPromises)).flatMap((plugin) => plugin || []);
|
||||
const mcpToolPromises = [];
|
||||
/** MCP server tools are initialized sequentially by server */
|
||||
let index = -1;
|
||||
const failedMCPServers = new Set();
|
||||
const safeUser = createSafeUser(options.req?.user);
|
||||
for (const [serverName, toolConfigs] of Object.entries(requestedMCPTools)) {
|
||||
index++;
|
||||
/** @type {LCAvailableTools} */
|
||||
let availableTools;
|
||||
for (const config of toolConfigs) {
|
||||
try {
|
||||
if (failedMCPServers.has(serverName)) {
|
||||
continue;
|
||||
}
|
||||
const mcpParams = {
|
||||
index,
|
||||
signal,
|
||||
user: safeUser,
|
||||
userMCPAuthMap,
|
||||
res: options.res,
|
||||
streamId: options.req?._resumableStreamId || null,
|
||||
model: agent?.model ?? model,
|
||||
serverName: config.serverName,
|
||||
provider: agent?.provider ?? endpoint,
|
||||
config: config.config,
|
||||
};
|
||||
|
||||
if (config.type === 'all' && toolConfigs.length === 1) {
|
||||
/** Handle async loading for single 'all' tool config */
|
||||
mcpToolPromises.push(
|
||||
createMCPTools(mcpParams).catch((error) => {
|
||||
logger.error(`Error loading ${serverName} tools:`, error);
|
||||
return null;
|
||||
}),
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if (!availableTools) {
|
||||
try {
|
||||
availableTools = await getMCPServerTools(safeUser.id, serverName);
|
||||
} catch (error) {
|
||||
logger.error(`Error fetching available tools for MCP server ${serverName}:`, error);
|
||||
}
|
||||
}
|
||||
|
||||
/** Handle synchronous loading */
|
||||
const mcpTool =
|
||||
config.type === 'all'
|
||||
? await createMCPTools(mcpParams)
|
||||
: await createMCPTool({
|
||||
...mcpParams,
|
||||
availableTools,
|
||||
toolKey: config.toolKey,
|
||||
});
|
||||
|
||||
if (Array.isArray(mcpTool)) {
|
||||
loadedTools.push(...mcpTool);
|
||||
} else if (mcpTool) {
|
||||
loadedTools.push(mcpTool);
|
||||
} else {
|
||||
failedMCPServers.add(serverName);
|
||||
logger.warn(
|
||||
`MCP tool creation failed for "${config.toolKey}", server may be unavailable or unauthenticated.`,
|
||||
);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error(`Error loading MCP tool for server ${serverName}:`, error);
|
||||
}
|
||||
}
|
||||
}
|
||||
loadedTools.push(...(await Promise.all(mcpToolPromises)).flatMap((plugin) => plugin || []));
|
||||
return { loadedTools, toolContextMap };
|
||||
};
|
||||
|
||||
module.exports = {
|
||||
loadToolWithAuth,
|
||||
validateTools,
|
||||
loadTools,
|
||||
};
|
||||
286
api/app/clients/tools/util/handleTools.test.js
Normal file
286
api/app/clients/tools/util/handleTools.test.js
Normal file
@ -0,0 +1,286 @@
|
||||
const mongoose = require('mongoose');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
|
||||
const mockPluginService = {
|
||||
updateUserPluginAuth: jest.fn(),
|
||||
deleteUserPluginAuth: jest.fn(),
|
||||
getUserPluginAuthValue: jest.fn(),
|
||||
};
|
||||
|
||||
jest.mock('~/server/services/PluginService', () => mockPluginService);
|
||||
|
||||
jest.mock('~/server/services/Config', () => ({
|
||||
getAppConfig: jest.fn().mockResolvedValue({
|
||||
// Default app config for tool tests
|
||||
paths: { uploads: '/tmp' },
|
||||
fileStrategy: 'local',
|
||||
filteredTools: [],
|
||||
includedTools: [],
|
||||
}),
|
||||
getCachedTools: jest.fn().mockResolvedValue({
|
||||
// Default cached tools for tests
|
||||
dalle: {
|
||||
type: 'function',
|
||||
function: {
|
||||
name: 'dalle',
|
||||
description: 'DALL-E image generation',
|
||||
parameters: {},
|
||||
},
|
||||
},
|
||||
}),
|
||||
}));
|
||||
|
||||
const { Calculator } = require('@librechat/agents');
|
||||
|
||||
const { User } = require('~/db/models');
|
||||
const PluginService = require('~/server/services/PluginService');
|
||||
const { validateTools, loadTools, loadToolWithAuth } = require('./handleTools');
|
||||
const { StructuredSD, availableTools, DALLE3 } = require('../');
|
||||
|
||||
describe('Tool Handlers', () => {
|
||||
let mongoServer;
|
||||
let fakeUser;
|
||||
const pluginKey = 'dalle';
|
||||
const pluginKey2 = 'wolfram';
|
||||
const ToolClass = DALLE3;
|
||||
const initialTools = [pluginKey, pluginKey2];
|
||||
const mockCredential = 'mock-credential';
|
||||
const mainPlugin = availableTools.find((tool) => tool.pluginKey === pluginKey);
|
||||
const authConfigs = mainPlugin.authConfig;
|
||||
|
||||
beforeAll(async () => {
|
||||
mongoServer = await MongoMemoryServer.create();
|
||||
const mongoUri = mongoServer.getUri();
|
||||
await mongoose.connect(mongoUri);
|
||||
|
||||
const userAuthValues = {};
|
||||
mockPluginService.getUserPluginAuthValue.mockImplementation((userId, authField) => {
|
||||
return userAuthValues[`${userId}-${authField}`];
|
||||
});
|
||||
mockPluginService.updateUserPluginAuth.mockImplementation(
|
||||
(userId, authField, _pluginKey, credential) => {
|
||||
const fields = authField.split('||');
|
||||
fields.forEach((field) => {
|
||||
userAuthValues[`${userId}-${field}`] = credential;
|
||||
});
|
||||
},
|
||||
);
|
||||
|
||||
fakeUser = new User({
|
||||
name: 'Fake User',
|
||||
username: 'fakeuser',
|
||||
email: 'fakeuser@example.com',
|
||||
emailVerified: false,
|
||||
// file deepcode ignore NoHardcodedPasswords/test: fake value
|
||||
password: 'fakepassword123',
|
||||
avatar: '',
|
||||
provider: 'local',
|
||||
role: 'USER',
|
||||
googleId: null,
|
||||
plugins: [],
|
||||
refreshToken: [],
|
||||
});
|
||||
await fakeUser.save();
|
||||
for (const authConfig of authConfigs) {
|
||||
await PluginService.updateUserPluginAuth(
|
||||
fakeUser._id,
|
||||
authConfig.authField,
|
||||
pluginKey,
|
||||
mockCredential,
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await mongoose.disconnect();
|
||||
await mongoServer.stop();
|
||||
});
|
||||
|
||||
beforeEach(async () => {
|
||||
// Clear mocks but not the database since we need the user to persist
|
||||
jest.clearAllMocks();
|
||||
|
||||
// Reset the mock implementations
|
||||
const userAuthValues = {};
|
||||
mockPluginService.getUserPluginAuthValue.mockImplementation((userId, authField) => {
|
||||
return userAuthValues[`${userId}-${authField}`];
|
||||
});
|
||||
mockPluginService.updateUserPluginAuth.mockImplementation(
|
||||
(userId, authField, _pluginKey, credential) => {
|
||||
const fields = authField.split('||');
|
||||
fields.forEach((field) => {
|
||||
userAuthValues[`${userId}-${field}`] = credential;
|
||||
});
|
||||
},
|
||||
);
|
||||
|
||||
// Re-add the auth configs for the user
|
||||
for (const authConfig of authConfigs) {
|
||||
await PluginService.updateUserPluginAuth(
|
||||
fakeUser._id,
|
||||
authConfig.authField,
|
||||
pluginKey,
|
||||
mockCredential,
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
describe('validateTools', () => {
|
||||
it('returns valid tools given input tools and user authentication', async () => {
|
||||
const validTools = await validateTools(fakeUser._id, initialTools);
|
||||
expect(validTools).toBeDefined();
|
||||
expect(validTools.some((tool) => tool === pluginKey)).toBeTruthy();
|
||||
expect(validTools.length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it('removes tools without valid credentials from the validTools array', async () => {
|
||||
const validTools = await validateTools(fakeUser._id, initialTools);
|
||||
expect(validTools.some((tool) => tool.pluginKey === pluginKey2)).toBeFalsy();
|
||||
});
|
||||
|
||||
it('returns an empty array when no authenticated tools are provided', async () => {
|
||||
const validTools = await validateTools(fakeUser._id, []);
|
||||
expect(validTools).toEqual([]);
|
||||
});
|
||||
|
||||
it('should validate a tool from an Environment Variable', async () => {
|
||||
const plugin = availableTools.find((tool) => tool.pluginKey === pluginKey2);
|
||||
const authConfigs = plugin.authConfig;
|
||||
for (const authConfig of authConfigs) {
|
||||
process.env[authConfig.authField] = mockCredential;
|
||||
}
|
||||
const validTools = await validateTools(fakeUser._id, [pluginKey2]);
|
||||
expect(validTools.length).toEqual(1);
|
||||
for (const authConfig of authConfigs) {
|
||||
delete process.env[authConfig.authField];
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('loadTools', () => {
|
||||
let toolFunctions;
|
||||
let loadTool1;
|
||||
let loadTool2;
|
||||
let loadTool3;
|
||||
const sampleTools = [...initialTools, 'calculator'];
|
||||
let ToolClass2 = Calculator;
|
||||
let remainingTools = availableTools.filter(
|
||||
(tool) => sampleTools.indexOf(tool.pluginKey) === -1,
|
||||
);
|
||||
|
||||
beforeAll(async () => {
|
||||
const toolMap = await loadTools({
|
||||
user: fakeUser._id,
|
||||
tools: sampleTools,
|
||||
returnMap: true,
|
||||
useSpecs: true,
|
||||
});
|
||||
toolFunctions = toolMap;
|
||||
loadTool1 = toolFunctions[sampleTools[0]];
|
||||
loadTool2 = toolFunctions[sampleTools[1]];
|
||||
loadTool3 = toolFunctions[sampleTools[2]];
|
||||
});
|
||||
|
||||
let originalEnv;
|
||||
|
||||
beforeEach(() => {
|
||||
originalEnv = process.env;
|
||||
process.env = { ...originalEnv };
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
process.env = originalEnv;
|
||||
});
|
||||
|
||||
it('returns the expected load functions for requested tools', async () => {
|
||||
expect(loadTool1).toBeDefined();
|
||||
expect(loadTool2).toBeDefined();
|
||||
expect(loadTool3).toBeDefined();
|
||||
|
||||
for (const tool of remainingTools) {
|
||||
expect(toolFunctions[tool.pluginKey]).toBeUndefined();
|
||||
}
|
||||
});
|
||||
|
||||
it('should initialize an authenticated tool or one without authentication', async () => {
|
||||
const authTool = await loadTool1();
|
||||
const tool = await loadTool3();
|
||||
expect(authTool).toBeInstanceOf(ToolClass);
|
||||
expect(tool).toBeInstanceOf(ToolClass2);
|
||||
});
|
||||
|
||||
it('should initialize an authenticated tool with primary auth field', async () => {
|
||||
process.env.DALLE3_API_KEY = 'mocked_api_key';
|
||||
const initToolFunction = loadToolWithAuth(
|
||||
'userId',
|
||||
['DALLE3_API_KEY||DALLE_API_KEY'],
|
||||
ToolClass,
|
||||
);
|
||||
const authTool = await initToolFunction();
|
||||
|
||||
expect(authTool).toBeInstanceOf(ToolClass);
|
||||
expect(mockPluginService.getUserPluginAuthValue).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should initialize an authenticated tool with alternate auth field when primary is missing', async () => {
|
||||
delete process.env.DALLE3_API_KEY; // Ensure the primary key is not set
|
||||
process.env.DALLE_API_KEY = 'mocked_alternate_api_key';
|
||||
const initToolFunction = loadToolWithAuth(
|
||||
'userId',
|
||||
['DALLE3_API_KEY||DALLE_API_KEY'],
|
||||
ToolClass,
|
||||
);
|
||||
const authTool = await initToolFunction();
|
||||
|
||||
expect(authTool).toBeInstanceOf(ToolClass);
|
||||
expect(mockPluginService.getUserPluginAuthValue).toHaveBeenCalledTimes(1);
|
||||
expect(mockPluginService.getUserPluginAuthValue).toHaveBeenCalledWith(
|
||||
'userId',
|
||||
'DALLE3_API_KEY',
|
||||
true,
|
||||
);
|
||||
});
|
||||
|
||||
it('should fallback to getUserPluginAuthValue when env vars are missing', async () => {
|
||||
mockPluginService.updateUserPluginAuth('userId', 'DALLE_API_KEY', 'dalle', 'mocked_api_key');
|
||||
const initToolFunction = loadToolWithAuth(
|
||||
'userId',
|
||||
['DALLE3_API_KEY||DALLE_API_KEY'],
|
||||
ToolClass,
|
||||
);
|
||||
const authTool = await initToolFunction();
|
||||
|
||||
expect(authTool).toBeInstanceOf(ToolClass);
|
||||
expect(mockPluginService.getUserPluginAuthValue).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('should throw an error for an unauthenticated tool', async () => {
|
||||
try {
|
||||
await loadTool2();
|
||||
} catch (error) {
|
||||
expect(error).toBeDefined();
|
||||
}
|
||||
});
|
||||
it('returns an empty object when no tools are requested', async () => {
|
||||
toolFunctions = await loadTools({
|
||||
user: fakeUser._id,
|
||||
returnMap: true,
|
||||
useSpecs: true,
|
||||
});
|
||||
expect(toolFunctions).toEqual({});
|
||||
});
|
||||
it('should return the StructuredTool version when using functions', async () => {
|
||||
process.env.SD_WEBUI_URL = mockCredential;
|
||||
toolFunctions = await loadTools({
|
||||
user: fakeUser._id,
|
||||
tools: ['stable-diffusion'],
|
||||
functions: true,
|
||||
returnMap: true,
|
||||
useSpecs: true,
|
||||
});
|
||||
const structuredTool = await toolFunctions['stable-diffusion']();
|
||||
expect(structuredTool).toBeInstanceOf(StructuredSD);
|
||||
delete process.env.SD_WEBUI_URL;
|
||||
});
|
||||
});
|
||||
});
|
||||
6
api/app/clients/tools/util/index.js
Normal file
6
api/app/clients/tools/util/index.js
Normal file
@ -0,0 +1,6 @@
|
||||
const { validateTools, loadTools } = require('./handleTools');
|
||||
|
||||
module.exports = {
|
||||
validateTools,
|
||||
loadTools,
|
||||
};
|
||||
5
api/app/index.js
Normal file
5
api/app/index.js
Normal file
@ -0,0 +1,5 @@
|
||||
const clients = require('./clients');
|
||||
|
||||
module.exports = {
|
||||
...clients,
|
||||
};
|
||||
86
api/cache/banViolation.js
vendored
Normal file
86
api/cache/banViolation.js
vendored
Normal file
@ -0,0 +1,86 @@
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { isEnabled, math } = require('@librechat/api');
|
||||
const { ViolationTypes } = require('librechat-data-provider');
|
||||
const { deleteAllUserSessions } = require('~/models');
|
||||
const { removePorts } = require('~/server/utils');
|
||||
const getLogStores = require('./getLogStores');
|
||||
|
||||
const { BAN_VIOLATIONS, BAN_INTERVAL } = process.env ?? {};
|
||||
const interval = math(BAN_INTERVAL, 20);
|
||||
|
||||
/**
|
||||
* Bans a user based on violation criteria.
|
||||
*
|
||||
* If the user's violation count is a multiple of the BAN_INTERVAL, the user will be banned.
|
||||
* The duration of the ban is determined by the BAN_DURATION environment variable.
|
||||
* If BAN_DURATION is not set or invalid, the user will not be banned.
|
||||
* Sessions will be deleted and the refreshToken cookie will be cleared even with
|
||||
* an invalid or nill duration, which is a "soft" ban; the user can remain active until
|
||||
* access token expiry.
|
||||
*
|
||||
* @async
|
||||
* @param {Object} req - Express request object containing user information.
|
||||
* @param {Object} res - Express response object.
|
||||
* @param {Object} errorMessage - Object containing user violation details.
|
||||
* @param {string} errorMessage.type - Type of the violation.
|
||||
* @param {string} errorMessage.user_id - ID of the user who committed the violation.
|
||||
* @param {number} errorMessage.violation_count - Number of violations committed by the user.
|
||||
*
|
||||
* @returns {Promise<void>}
|
||||
*
|
||||
*/
|
||||
const banViolation = async (req, res, errorMessage) => {
|
||||
if (!isEnabled(BAN_VIOLATIONS)) {
|
||||
return;
|
||||
}
|
||||
if (!errorMessage) {
|
||||
return;
|
||||
}
|
||||
|
||||
const { type, user_id, prev_count, violation_count } = errorMessage;
|
||||
|
||||
const prevThreshold = Math.floor(prev_count / interval);
|
||||
const currentThreshold = Math.floor(violation_count / interval);
|
||||
|
||||
if (prevThreshold >= currentThreshold) {
|
||||
return;
|
||||
}
|
||||
|
||||
await deleteAllUserSessions({ userId: user_id });
|
||||
|
||||
/** Clear OpenID session tokens if present */
|
||||
if (req.session?.openidTokens) {
|
||||
delete req.session.openidTokens;
|
||||
}
|
||||
|
||||
res.clearCookie('refreshToken');
|
||||
res.clearCookie('openid_access_token');
|
||||
res.clearCookie('openid_user_id');
|
||||
res.clearCookie('token_provider');
|
||||
|
||||
const banLogs = getLogStores(ViolationTypes.BAN);
|
||||
const duration = errorMessage.duration || banLogs.opts.ttl;
|
||||
if (duration <= 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
req.ip = removePorts(req);
|
||||
logger.info(
|
||||
`[BAN] Banning user ${user_id} ${req.ip ? `@ ${req.ip} ` : ''}for ${
|
||||
duration / 1000 / 60
|
||||
} minutes`,
|
||||
);
|
||||
|
||||
const expiresAt = Date.now() + duration;
|
||||
await banLogs.set(user_id, { type, violation_count, duration, expiresAt });
|
||||
if (req.ip) {
|
||||
await banLogs.set(req.ip, { type, user_id, violation_count, duration, expiresAt });
|
||||
}
|
||||
|
||||
errorMessage.ban = true;
|
||||
errorMessage.ban_duration = duration;
|
||||
|
||||
return;
|
||||
};
|
||||
|
||||
module.exports = banViolation;
|
||||
136
api/cache/banViolation.spec.js
vendored
Normal file
136
api/cache/banViolation.spec.js
vendored
Normal file
@ -0,0 +1,136 @@
|
||||
const mongoose = require('mongoose');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
const banViolation = require('./banViolation');
|
||||
|
||||
// Mock deleteAllUserSessions since we're testing ban logic, not session deletion
|
||||
jest.mock('~/models', () => ({
|
||||
...jest.requireActual('~/models'),
|
||||
deleteAllUserSessions: jest.fn().mockResolvedValue(true),
|
||||
}));
|
||||
|
||||
describe('banViolation', () => {
|
||||
let mongoServer;
|
||||
let req, res, errorMessage;
|
||||
|
||||
beforeAll(async () => {
|
||||
mongoServer = await MongoMemoryServer.create();
|
||||
const mongoUri = mongoServer.getUri();
|
||||
await mongoose.connect(mongoUri);
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await mongoose.disconnect();
|
||||
await mongoServer.stop();
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
req = {
|
||||
ip: '127.0.0.1',
|
||||
cookies: {
|
||||
refreshToken: 'someToken',
|
||||
},
|
||||
};
|
||||
res = {
|
||||
clearCookie: jest.fn(),
|
||||
};
|
||||
errorMessage = {
|
||||
type: 'someViolation',
|
||||
user_id: new mongoose.Types.ObjectId().toString(), // Use valid ObjectId
|
||||
prev_count: 0,
|
||||
violation_count: 0,
|
||||
};
|
||||
process.env.BAN_VIOLATIONS = 'true';
|
||||
process.env.BAN_DURATION = '7200000'; // 2 hours in ms
|
||||
process.env.BAN_INTERVAL = '20';
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
it('should not ban if BAN_VIOLATIONS are not enabled', async () => {
|
||||
process.env.BAN_VIOLATIONS = 'false';
|
||||
await banViolation(req, res, errorMessage);
|
||||
expect(errorMessage.ban).toBeFalsy();
|
||||
});
|
||||
|
||||
it('should not ban if errorMessage is not provided', async () => {
|
||||
await banViolation(req, res, null);
|
||||
expect(errorMessage.ban).toBeFalsy();
|
||||
});
|
||||
|
||||
it('[1/3] should ban if violation_count crosses the interval threshold: 19 -> 39', async () => {
|
||||
errorMessage.prev_count = 19;
|
||||
errorMessage.violation_count = 39;
|
||||
await banViolation(req, res, errorMessage);
|
||||
expect(errorMessage.ban).toBeTruthy();
|
||||
});
|
||||
|
||||
it('[2/3] should ban if violation_count crosses the interval threshold: 19 -> 20', async () => {
|
||||
errorMessage.prev_count = 19;
|
||||
errorMessage.violation_count = 20;
|
||||
await banViolation(req, res, errorMessage);
|
||||
expect(errorMessage.ban).toBeTruthy();
|
||||
});
|
||||
|
||||
const randomValueAbove = Math.floor(20 + Math.random() * 100);
|
||||
it(`[3/3] should ban if violation_count crosses the interval threshold: 19 -> ${randomValueAbove}`, async () => {
|
||||
errorMessage.prev_count = 19;
|
||||
errorMessage.violation_count = randomValueAbove;
|
||||
await banViolation(req, res, errorMessage);
|
||||
expect(errorMessage.ban).toBeTruthy();
|
||||
});
|
||||
|
||||
it('should handle invalid BAN_INTERVAL and default to 20', async () => {
|
||||
process.env.BAN_INTERVAL = 'invalid';
|
||||
errorMessage.prev_count = 19;
|
||||
errorMessage.violation_count = 39;
|
||||
await banViolation(req, res, errorMessage);
|
||||
expect(errorMessage.ban).toBeTruthy();
|
||||
});
|
||||
|
||||
it('should ban if BAN_DURATION is invalid as default is 2 hours', async () => {
|
||||
process.env.BAN_DURATION = 'invalid';
|
||||
errorMessage.prev_count = 19;
|
||||
errorMessage.violation_count = 39;
|
||||
await banViolation(req, res, errorMessage);
|
||||
expect(errorMessage.ban).toBeTruthy();
|
||||
});
|
||||
|
||||
it('should not ban if BAN_DURATION is 0 but should clear cookies', async () => {
|
||||
process.env.BAN_DURATION = '0';
|
||||
errorMessage.prev_count = 19;
|
||||
errorMessage.violation_count = 39;
|
||||
await banViolation(req, res, errorMessage);
|
||||
expect(res.clearCookie).toHaveBeenCalledWith('refreshToken');
|
||||
});
|
||||
|
||||
it('should not ban if violation_count does not change', async () => {
|
||||
errorMessage.prev_count = 0;
|
||||
errorMessage.violation_count = 0;
|
||||
await banViolation(req, res, errorMessage);
|
||||
expect(errorMessage.ban).toBeFalsy();
|
||||
});
|
||||
|
||||
it('[1/2] should not ban if violation_count does not cross the interval threshold: 0 -> 19', async () => {
|
||||
errorMessage.prev_count = 0;
|
||||
errorMessage.violation_count = 19;
|
||||
await banViolation(req, res, errorMessage);
|
||||
expect(errorMessage.ban).toBeFalsy();
|
||||
});
|
||||
|
||||
const randomValueUnder = Math.floor(1 + Math.random() * 19);
|
||||
it(`[2/2] should not ban if violation_count does not cross the interval threshold: 0 -> ${randomValueUnder}`, async () => {
|
||||
errorMessage.prev_count = 0;
|
||||
errorMessage.violation_count = randomValueUnder;
|
||||
await banViolation(req, res, errorMessage);
|
||||
expect(errorMessage.ban).toBeFalsy();
|
||||
});
|
||||
|
||||
it('[EDGE CASE] should not ban if violation_count is lower', async () => {
|
||||
errorMessage.prev_count = 0;
|
||||
errorMessage.violation_count = -10;
|
||||
await banViolation(req, res, errorMessage);
|
||||
expect(errorMessage.ban).toBeFalsy();
|
||||
});
|
||||
});
|
||||
49
api/cache/clearPendingReq.js
vendored
Normal file
49
api/cache/clearPendingReq.js
vendored
Normal file
@ -0,0 +1,49 @@
|
||||
const { isEnabled } = require('@librechat/api');
|
||||
const { Time, CacheKeys } = require('librechat-data-provider');
|
||||
const getLogStores = require('./getLogStores');
|
||||
|
||||
const { USE_REDIS, LIMIT_CONCURRENT_MESSAGES } = process.env ?? {};
|
||||
|
||||
/**
|
||||
* Clear or decrement pending requests from the cache.
|
||||
* Checks the environmental variable LIMIT_CONCURRENT_MESSAGES;
|
||||
* if the rule is enabled ('true'), it either decrements the count of pending requests
|
||||
* or deletes the key if the count is less than or equal to 1.
|
||||
*
|
||||
* @module clearPendingReq
|
||||
* @requires ./getLogStores
|
||||
* @requires ../server/utils
|
||||
* @requires process
|
||||
*
|
||||
* @async
|
||||
* @function
|
||||
* @param {Object} params - The parameters object.
|
||||
* @param {string} params.userId - The user ID for which the pending requests are to be cleared or decremented.
|
||||
* @param {Object} [params.cache] - An optional cache object to use. If not provided, a default cache will be fetched using getLogStores.
|
||||
* @returns {Promise<void>} A promise that either decrements the 'pendingRequests' count, deletes the key from the store, or resolves with no value.
|
||||
*/
|
||||
const clearPendingReq = async ({ userId, cache: _cache }) => {
|
||||
if (!userId) {
|
||||
return;
|
||||
} else if (!isEnabled(LIMIT_CONCURRENT_MESSAGES)) {
|
||||
return;
|
||||
}
|
||||
|
||||
const namespace = CacheKeys.PENDING_REQ;
|
||||
const cache = _cache ?? getLogStores(namespace);
|
||||
|
||||
if (!cache) {
|
||||
return;
|
||||
}
|
||||
|
||||
const key = `${isEnabled(USE_REDIS) ? namespace : ''}:${userId ?? ''}`;
|
||||
const currentReq = +((await cache.get(key)) ?? 0);
|
||||
|
||||
if (currentReq && currentReq >= 1) {
|
||||
await cache.set(key, currentReq - 1, Time.ONE_MINUTE);
|
||||
} else {
|
||||
await cache.delete(key);
|
||||
}
|
||||
};
|
||||
|
||||
module.exports = clearPendingReq;
|
||||
227
api/cache/getLogStores.js
vendored
Normal file
227
api/cache/getLogStores.js
vendored
Normal file
@ -0,0 +1,227 @@
|
||||
const { Keyv } = require('keyv');
|
||||
const { Time, CacheKeys, ViolationTypes } = require('librechat-data-provider');
|
||||
const {
|
||||
logFile,
|
||||
keyvMongo,
|
||||
cacheConfig,
|
||||
sessionCache,
|
||||
standardCache,
|
||||
violationCache,
|
||||
} = require('@librechat/api');
|
||||
|
||||
const namespaces = {
|
||||
[ViolationTypes.GENERAL]: new Keyv({ store: logFile, namespace: 'violations' }),
|
||||
[ViolationTypes.LOGINS]: violationCache(ViolationTypes.LOGINS),
|
||||
[ViolationTypes.CONCURRENT]: violationCache(ViolationTypes.CONCURRENT),
|
||||
[ViolationTypes.NON_BROWSER]: violationCache(ViolationTypes.NON_BROWSER),
|
||||
[ViolationTypes.MESSAGE_LIMIT]: violationCache(ViolationTypes.MESSAGE_LIMIT),
|
||||
[ViolationTypes.REGISTRATIONS]: violationCache(ViolationTypes.REGISTRATIONS),
|
||||
[ViolationTypes.TOKEN_BALANCE]: violationCache(ViolationTypes.TOKEN_BALANCE),
|
||||
[ViolationTypes.TTS_LIMIT]: violationCache(ViolationTypes.TTS_LIMIT),
|
||||
[ViolationTypes.STT_LIMIT]: violationCache(ViolationTypes.STT_LIMIT),
|
||||
[ViolationTypes.CONVO_ACCESS]: violationCache(ViolationTypes.CONVO_ACCESS),
|
||||
[ViolationTypes.TOOL_CALL_LIMIT]: violationCache(ViolationTypes.TOOL_CALL_LIMIT),
|
||||
[ViolationTypes.FILE_UPLOAD_LIMIT]: violationCache(ViolationTypes.FILE_UPLOAD_LIMIT),
|
||||
[ViolationTypes.VERIFY_EMAIL_LIMIT]: violationCache(ViolationTypes.VERIFY_EMAIL_LIMIT),
|
||||
[ViolationTypes.RESET_PASSWORD_LIMIT]: violationCache(ViolationTypes.RESET_PASSWORD_LIMIT),
|
||||
[ViolationTypes.ILLEGAL_MODEL_REQUEST]: violationCache(ViolationTypes.ILLEGAL_MODEL_REQUEST),
|
||||
[ViolationTypes.BAN]: new Keyv({
|
||||
store: keyvMongo,
|
||||
namespace: CacheKeys.BANS,
|
||||
ttl: cacheConfig.BAN_DURATION,
|
||||
}),
|
||||
|
||||
[CacheKeys.OPENID_SESSION]: sessionCache(CacheKeys.OPENID_SESSION),
|
||||
[CacheKeys.SAML_SESSION]: sessionCache(CacheKeys.SAML_SESSION),
|
||||
|
||||
[CacheKeys.ROLES]: standardCache(CacheKeys.ROLES),
|
||||
[CacheKeys.APP_CONFIG]: standardCache(CacheKeys.APP_CONFIG),
|
||||
[CacheKeys.CONFIG_STORE]: standardCache(CacheKeys.CONFIG_STORE),
|
||||
[CacheKeys.PENDING_REQ]: standardCache(CacheKeys.PENDING_REQ),
|
||||
[CacheKeys.ENCODED_DOMAINS]: new Keyv({ store: keyvMongo, namespace: CacheKeys.ENCODED_DOMAINS }),
|
||||
[CacheKeys.ABORT_KEYS]: standardCache(CacheKeys.ABORT_KEYS, Time.TEN_MINUTES),
|
||||
[CacheKeys.TOKEN_CONFIG]: standardCache(CacheKeys.TOKEN_CONFIG, Time.THIRTY_MINUTES),
|
||||
[CacheKeys.GEN_TITLE]: standardCache(CacheKeys.GEN_TITLE, Time.TWO_MINUTES),
|
||||
[CacheKeys.S3_EXPIRY_INTERVAL]: standardCache(CacheKeys.S3_EXPIRY_INTERVAL, Time.THIRTY_MINUTES),
|
||||
[CacheKeys.MODEL_QUERIES]: standardCache(CacheKeys.MODEL_QUERIES),
|
||||
[CacheKeys.AUDIO_RUNS]: standardCache(CacheKeys.AUDIO_RUNS, Time.TEN_MINUTES),
|
||||
[CacheKeys.MESSAGES]: standardCache(CacheKeys.MESSAGES, Time.ONE_MINUTE),
|
||||
[CacheKeys.FLOWS]: standardCache(CacheKeys.FLOWS, Time.ONE_MINUTE * 3),
|
||||
[CacheKeys.OPENID_EXCHANGED_TOKENS]: standardCache(
|
||||
CacheKeys.OPENID_EXCHANGED_TOKENS,
|
||||
Time.TEN_MINUTES,
|
||||
),
|
||||
};
|
||||
|
||||
/**
|
||||
* Gets all cache stores that have TTL configured
|
||||
* @returns {Keyv[]}
|
||||
*/
|
||||
function getTTLStores() {
|
||||
return Object.values(namespaces).filter(
|
||||
(store) =>
|
||||
store instanceof Keyv &&
|
||||
parseInt(store.opts?.ttl ?? '0') > 0 &&
|
||||
!store.opts?.store?.constructor?.name?.includes('Redis'), // Only include non-Redis stores
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Clears entries older than the cache's TTL
|
||||
* @param {Keyv} cache
|
||||
*/
|
||||
async function clearExpiredFromCache(cache) {
|
||||
if (!cache?.opts?.store?.entries) {
|
||||
return;
|
||||
}
|
||||
|
||||
const ttl = cache.opts.ttl;
|
||||
if (!ttl) {
|
||||
return;
|
||||
}
|
||||
|
||||
const expiryTime = Date.now() - ttl;
|
||||
let cleared = 0;
|
||||
|
||||
// Get all keys first to avoid modification during iteration
|
||||
const keys = Array.from(cache.opts.store.keys());
|
||||
|
||||
for (const key of keys) {
|
||||
try {
|
||||
const raw = cache.opts.store.get(key);
|
||||
if (!raw) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const data = cache.opts.deserialize(raw);
|
||||
// Check if the entry is older than TTL
|
||||
if (data?.expires && data.expires <= expiryTime) {
|
||||
const deleted = await cache.opts.store.delete(key);
|
||||
if (!deleted) {
|
||||
cacheConfig.DEBUG_MEMORY_CACHE &&
|
||||
console.warn(`[Cache] Error deleting entry: ${key} from ${cache.opts.namespace}`);
|
||||
continue;
|
||||
}
|
||||
cleared++;
|
||||
}
|
||||
} catch (error) {
|
||||
cacheConfig.DEBUG_MEMORY_CACHE &&
|
||||
console.log(`[Cache] Error processing entry from ${cache.opts.namespace}:`, error);
|
||||
const deleted = await cache.opts.store.delete(key);
|
||||
if (!deleted) {
|
||||
cacheConfig.DEBUG_MEMORY_CACHE &&
|
||||
console.warn(`[Cache] Error deleting entry: ${key} from ${cache.opts.namespace}`);
|
||||
continue;
|
||||
}
|
||||
cleared++;
|
||||
}
|
||||
}
|
||||
|
||||
if (cleared > 0) {
|
||||
cacheConfig.DEBUG_MEMORY_CACHE &&
|
||||
console.log(
|
||||
`[Cache] Cleared ${cleared} entries older than ${ttl}ms from ${cache.opts.namespace}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
const auditCache = () => {
|
||||
const ttlStores = getTTLStores();
|
||||
console.log('[Cache] Starting audit');
|
||||
|
||||
ttlStores.forEach((store) => {
|
||||
if (!store?.opts?.store?.entries) {
|
||||
return;
|
||||
}
|
||||
|
||||
console.log(`[Cache] ${store.opts.namespace} entries:`, {
|
||||
count: store.opts.store.size,
|
||||
ttl: store.opts.ttl,
|
||||
keys: Array.from(store.opts.store.keys()),
|
||||
entriesWithTimestamps: Array.from(store.opts.store.entries()).map(([key, value]) => ({
|
||||
key,
|
||||
value,
|
||||
})),
|
||||
});
|
||||
});
|
||||
};
|
||||
|
||||
/**
|
||||
* Clears expired entries from all TTL-enabled stores
|
||||
*/
|
||||
async function clearAllExpiredFromCache() {
|
||||
const ttlStores = getTTLStores();
|
||||
await Promise.all(ttlStores.map((store) => clearExpiredFromCache(store)));
|
||||
|
||||
// Force garbage collection if available (Node.js with --expose-gc flag)
|
||||
if (global.gc) {
|
||||
global.gc();
|
||||
}
|
||||
}
|
||||
|
||||
if (!cacheConfig.USE_REDIS && !cacheConfig.CI) {
|
||||
/** @type {Set<NodeJS.Timeout>} */
|
||||
const cleanupIntervals = new Set();
|
||||
|
||||
// Clear expired entries every 30 seconds
|
||||
const cleanup = setInterval(() => {
|
||||
clearAllExpiredFromCache();
|
||||
}, Time.THIRTY_SECONDS);
|
||||
|
||||
cleanupIntervals.add(cleanup);
|
||||
|
||||
if (cacheConfig.DEBUG_MEMORY_CACHE) {
|
||||
const monitor = setInterval(() => {
|
||||
const ttlStores = getTTLStores();
|
||||
const memory = process.memoryUsage();
|
||||
const totalSize = ttlStores.reduce((sum, store) => sum + (store.opts?.store?.size ?? 0), 0);
|
||||
|
||||
console.log('[Cache] Memory usage:', {
|
||||
heapUsed: `${(memory.heapUsed / 1024 / 1024).toFixed(2)} MB`,
|
||||
heapTotal: `${(memory.heapTotal / 1024 / 1024).toFixed(2)} MB`,
|
||||
rss: `${(memory.rss / 1024 / 1024).toFixed(2)} MB`,
|
||||
external: `${(memory.external / 1024 / 1024).toFixed(2)} MB`,
|
||||
totalCacheEntries: totalSize,
|
||||
});
|
||||
|
||||
auditCache();
|
||||
}, Time.ONE_MINUTE);
|
||||
|
||||
cleanupIntervals.add(monitor);
|
||||
}
|
||||
|
||||
const dispose = () => {
|
||||
cacheConfig.DEBUG_MEMORY_CACHE && console.log('[Cache] Cleaning up and shutting down...');
|
||||
cleanupIntervals.forEach((interval) => clearInterval(interval));
|
||||
cleanupIntervals.clear();
|
||||
|
||||
// One final cleanup before exit
|
||||
clearAllExpiredFromCache().then(() => {
|
||||
cacheConfig.DEBUG_MEMORY_CACHE && console.log('[Cache] Final cleanup completed');
|
||||
process.exit(0);
|
||||
});
|
||||
};
|
||||
|
||||
// Handle various termination signals
|
||||
process.on('SIGTERM', dispose);
|
||||
process.on('SIGINT', dispose);
|
||||
process.on('SIGQUIT', dispose);
|
||||
process.on('SIGHUP', dispose);
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the keyv cache specified by type.
|
||||
* If an invalid type is passed, an error will be thrown.
|
||||
*
|
||||
* @param {string} key - The key for the namespace to access
|
||||
* @returns {Keyv} - If a valid key is passed, returns an object containing the cache store of the specified key.
|
||||
* @throws Will throw an error if an invalid key is passed.
|
||||
*/
|
||||
const getLogStores = (key) => {
|
||||
if (!key || !namespaces[key]) {
|
||||
throw new Error(`Invalid store key: ${key}`);
|
||||
}
|
||||
return namespaces[key];
|
||||
};
|
||||
|
||||
module.exports = getLogStores;
|
||||
4
api/cache/index.js
vendored
Normal file
4
api/cache/index.js
vendored
Normal file
@ -0,0 +1,4 @@
|
||||
const getLogStores = require('./getLogStores');
|
||||
const logViolation = require('./logViolation');
|
||||
|
||||
module.exports = { getLogStores, logViolation };
|
||||
40
api/cache/logViolation.js
vendored
Normal file
40
api/cache/logViolation.js
vendored
Normal file
@ -0,0 +1,40 @@
|
||||
const { isEnabled } = require('@librechat/api');
|
||||
const { ViolationTypes } = require('librechat-data-provider');
|
||||
const getLogStores = require('./getLogStores');
|
||||
const banViolation = require('./banViolation');
|
||||
|
||||
/**
|
||||
* Logs the violation.
|
||||
*
|
||||
* @param {Object} req - Express request object containing user information.
|
||||
* @param {Object} res - Express response object.
|
||||
* @param {string} type - The type of violation.
|
||||
* @param {Object} errorMessage - The error message to log.
|
||||
* @param {number | string} [score=1] - The severity of the violation. Defaults to 1
|
||||
*/
|
||||
const logViolation = async (req, res, type, errorMessage, score = 1) => {
|
||||
const userId = req.user?.id ?? req.user?._id;
|
||||
if (!userId) {
|
||||
return;
|
||||
}
|
||||
const logs = getLogStores(ViolationTypes.GENERAL);
|
||||
const violationLogs = getLogStores(type);
|
||||
const key = isEnabled(process.env.USE_REDIS) ? `${type}:${userId}` : userId;
|
||||
|
||||
const userViolations = (await violationLogs.get(key)) ?? 0;
|
||||
const violationCount = +userViolations + +score;
|
||||
await violationLogs.set(key, violationCount);
|
||||
|
||||
errorMessage.user_id = userId;
|
||||
errorMessage.prev_count = userViolations;
|
||||
errorMessage.violation_count = violationCount;
|
||||
errorMessage.date = new Date().toISOString();
|
||||
|
||||
await banViolation(req, res, errorMessage);
|
||||
const userLogs = (await logs.get(key)) ?? [];
|
||||
userLogs.push(errorMessage);
|
||||
delete errorMessage.user_id;
|
||||
await logs.set(key, userLogs);
|
||||
};
|
||||
|
||||
module.exports = logViolation;
|
||||
38
api/config/index.js
Normal file
38
api/config/index.js
Normal file
@ -0,0 +1,38 @@
|
||||
const { EventSource } = require('eventsource');
|
||||
const { Time } = require('librechat-data-provider');
|
||||
const {
|
||||
MCPManager,
|
||||
FlowStateManager,
|
||||
MCPServersRegistry,
|
||||
OAuthReconnectionManager,
|
||||
} = require('@librechat/api');
|
||||
const logger = require('./winston');
|
||||
|
||||
global.EventSource = EventSource;
|
||||
|
||||
/** @type {MCPManager} */
|
||||
let flowManager = null;
|
||||
|
||||
/**
|
||||
* @param {Keyv} flowsCache
|
||||
* @returns {FlowStateManager}
|
||||
*/
|
||||
function getFlowStateManager(flowsCache) {
|
||||
if (!flowManager) {
|
||||
flowManager = new FlowStateManager(flowsCache, {
|
||||
ttl: Time.ONE_MINUTE * 3,
|
||||
});
|
||||
}
|
||||
return flowManager;
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
logger,
|
||||
createMCPServersRegistry: MCPServersRegistry.createInstance,
|
||||
getMCPServersRegistry: MCPServersRegistry.getInstance,
|
||||
createMCPManager: MCPManager.createInstance,
|
||||
getMCPManager: MCPManager.getInstance,
|
||||
getFlowStateManager,
|
||||
createOAuthReconnectionManager: OAuthReconnectionManager.createInstance,
|
||||
getOAuthReconnectionManager: OAuthReconnectionManager.getInstance,
|
||||
};
|
||||
102
api/config/meiliLogger.js
Normal file
102
api/config/meiliLogger.js
Normal file
@ -0,0 +1,102 @@
|
||||
const path = require('path');
|
||||
const fs = require('fs');
|
||||
const winston = require('winston');
|
||||
require('winston-daily-rotate-file');
|
||||
|
||||
/**
|
||||
* Determine the log directory.
|
||||
* Priority:
|
||||
* 1. LIBRECHAT_LOG_DIR environment variable (allows user override)
|
||||
* 2. /app/logs if running in Docker (bind-mounted with correct permissions)
|
||||
* 3. api/logs relative to this file (local development)
|
||||
*/
|
||||
const getLogDir = () => {
|
||||
if (process.env.LIBRECHAT_LOG_DIR) {
|
||||
return process.env.LIBRECHAT_LOG_DIR;
|
||||
}
|
||||
|
||||
// Check if running in Docker container (cwd is /app)
|
||||
if (process.cwd() === '/app') {
|
||||
const dockerLogDir = '/app/logs';
|
||||
// Ensure the directory exists
|
||||
if (!fs.existsSync(dockerLogDir)) {
|
||||
fs.mkdirSync(dockerLogDir, { recursive: true });
|
||||
}
|
||||
return dockerLogDir;
|
||||
}
|
||||
|
||||
// Local development: use api/logs relative to this file
|
||||
return path.join(__dirname, '..', 'logs');
|
||||
};
|
||||
|
||||
const logDir = getLogDir();
|
||||
|
||||
const { NODE_ENV, DEBUG_LOGGING = false } = process.env;
|
||||
|
||||
const useDebugLogging =
|
||||
(typeof DEBUG_LOGGING === 'string' && DEBUG_LOGGING?.toLowerCase() === 'true') ||
|
||||
DEBUG_LOGGING === true;
|
||||
|
||||
const levels = {
|
||||
error: 0,
|
||||
warn: 1,
|
||||
info: 2,
|
||||
http: 3,
|
||||
verbose: 4,
|
||||
debug: 5,
|
||||
activity: 6,
|
||||
silly: 7,
|
||||
};
|
||||
|
||||
winston.addColors({
|
||||
info: 'green', // fontStyle color
|
||||
warn: 'italic yellow',
|
||||
error: 'red',
|
||||
debug: 'blue',
|
||||
});
|
||||
|
||||
const level = () => {
|
||||
const env = NODE_ENV || 'development';
|
||||
const isDevelopment = env === 'development';
|
||||
return isDevelopment ? 'debug' : 'warn';
|
||||
};
|
||||
|
||||
const fileFormat = winston.format.combine(
|
||||
winston.format.timestamp({ format: 'YYYY-MM-DD HH:mm:ss' }),
|
||||
winston.format.errors({ stack: true }),
|
||||
winston.format.splat(),
|
||||
);
|
||||
|
||||
const logLevel = useDebugLogging ? 'debug' : 'error';
|
||||
const transports = [
|
||||
new winston.transports.DailyRotateFile({
|
||||
level: logLevel,
|
||||
filename: `${logDir}/meiliSync-%DATE%.log`,
|
||||
datePattern: 'YYYY-MM-DD',
|
||||
zippedArchive: true,
|
||||
maxSize: '20m',
|
||||
maxFiles: '14d',
|
||||
format: fileFormat,
|
||||
}),
|
||||
];
|
||||
|
||||
const consoleFormat = winston.format.combine(
|
||||
winston.format.colorize({ all: true }),
|
||||
winston.format.timestamp({ format: 'YYYY-MM-DD HH:mm:ss' }),
|
||||
winston.format.printf((info) => `${info.timestamp} ${info.level}: ${info.message}`),
|
||||
);
|
||||
|
||||
transports.push(
|
||||
new winston.transports.Console({
|
||||
level: 'info',
|
||||
format: consoleFormat,
|
||||
}),
|
||||
);
|
||||
|
||||
const logger = winston.createLogger({
|
||||
level: level(),
|
||||
levels,
|
||||
transports,
|
||||
});
|
||||
|
||||
module.exports = logger;
|
||||
232
api/config/parsers.js
Normal file
232
api/config/parsers.js
Normal file
@ -0,0 +1,232 @@
|
||||
const { klona } = require('klona');
|
||||
const winston = require('winston');
|
||||
const traverse = require('traverse');
|
||||
|
||||
const SPLAT_SYMBOL = Symbol.for('splat');
|
||||
const MESSAGE_SYMBOL = Symbol.for('message');
|
||||
const CONSOLE_JSON_STRING_LENGTH = parseInt(process.env.CONSOLE_JSON_STRING_LENGTH) || 255;
|
||||
const DEBUG_MESSAGE_LENGTH = parseInt(process.env.DEBUG_MESSAGE_LENGTH) || 150;
|
||||
|
||||
const sensitiveKeys = [
|
||||
/^(sk-)[^\s]+/, // OpenAI API key pattern
|
||||
/(Bearer )[^\s]+/, // Header: Bearer token pattern
|
||||
/(api-key:? )[^\s]+/, // Header: API key pattern
|
||||
/(key=)[^\s]+/, // URL query param: sensitive key pattern (Google)
|
||||
];
|
||||
|
||||
/**
|
||||
* Determines if a given value string is sensitive and returns matching regex patterns.
|
||||
*
|
||||
* @param {string} valueStr - The value string to check.
|
||||
* @returns {Array<RegExp>} An array of regex patterns that match the value string.
|
||||
*/
|
||||
function getMatchingSensitivePatterns(valueStr) {
|
||||
if (valueStr) {
|
||||
// Filter and return all regex patterns that match the value string
|
||||
return sensitiveKeys.filter((regex) => regex.test(valueStr));
|
||||
}
|
||||
return [];
|
||||
}
|
||||
|
||||
/**
|
||||
* Redacts sensitive information from a console message and trims it to a specified length if provided.
|
||||
* @param {string} str - The console message to be redacted.
|
||||
* @param {number} [trimLength] - The optional length at which to trim the redacted message.
|
||||
* @returns {string} - The redacted and optionally trimmed console message.
|
||||
*/
|
||||
function redactMessage(str, trimLength) {
|
||||
if (!str) {
|
||||
return '';
|
||||
}
|
||||
|
||||
const patterns = getMatchingSensitivePatterns(str);
|
||||
patterns.forEach((pattern) => {
|
||||
str = str.replace(pattern, '$1[REDACTED]');
|
||||
});
|
||||
|
||||
if (trimLength !== undefined && str.length > trimLength) {
|
||||
return `${str.substring(0, trimLength)}...`;
|
||||
}
|
||||
|
||||
return str;
|
||||
}
|
||||
|
||||
/**
|
||||
* Redacts sensitive information from log messages if the log level is 'error'.
|
||||
* Note: Intentionally mutates the object.
|
||||
* @param {Object} info - The log information object.
|
||||
* @returns {Object} - The modified log information object.
|
||||
*/
|
||||
const redactFormat = winston.format((info) => {
|
||||
if (info.level === 'error') {
|
||||
info.message = redactMessage(info.message);
|
||||
if (info[MESSAGE_SYMBOL]) {
|
||||
info[MESSAGE_SYMBOL] = redactMessage(info[MESSAGE_SYMBOL]);
|
||||
}
|
||||
}
|
||||
return info;
|
||||
});
|
||||
|
||||
/**
|
||||
* Truncates long strings, especially base64 image data, within log messages.
|
||||
*
|
||||
* @param {any} value - The value to be inspected and potentially truncated.
|
||||
* @param {number} [length] - The length at which to truncate the value. Default: 100.
|
||||
* @returns {any} - The truncated or original value.
|
||||
*/
|
||||
const truncateLongStrings = (value, length = 100) => {
|
||||
if (typeof value === 'string') {
|
||||
return value.length > length ? value.substring(0, length) + '... [truncated]' : value;
|
||||
}
|
||||
|
||||
return value;
|
||||
};
|
||||
|
||||
/**
|
||||
* An array mapping function that truncates long strings (objects converted to JSON strings).
|
||||
* @param {any} item - The item to be condensed.
|
||||
* @returns {any} - The condensed item.
|
||||
*/
|
||||
const condenseArray = (item) => {
|
||||
if (typeof item === 'string') {
|
||||
return truncateLongStrings(JSON.stringify(item));
|
||||
} else if (typeof item === 'object') {
|
||||
return truncateLongStrings(JSON.stringify(item));
|
||||
}
|
||||
return item;
|
||||
};
|
||||
|
||||
/**
|
||||
* Formats log messages for debugging purposes.
|
||||
* - Truncates long strings within log messages.
|
||||
* - Condenses arrays by truncating long strings and objects as strings within array items.
|
||||
* - Redacts sensitive information from log messages if the log level is 'error'.
|
||||
* - Converts log information object to a formatted string.
|
||||
*
|
||||
* @param {Object} options - The options for formatting log messages.
|
||||
* @param {string} options.level - The log level.
|
||||
* @param {string} options.message - The log message.
|
||||
* @param {string} options.timestamp - The timestamp of the log message.
|
||||
* @param {Object} options.metadata - Additional metadata associated with the log message.
|
||||
* @returns {string} - The formatted log message.
|
||||
*/
|
||||
const debugTraverse = winston.format.printf(({ level, message, timestamp, ...metadata }) => {
|
||||
if (!message) {
|
||||
return `${timestamp} ${level}`;
|
||||
}
|
||||
|
||||
if (!message?.trim || typeof message !== 'string') {
|
||||
return `${timestamp} ${level}: ${JSON.stringify(message)}`;
|
||||
}
|
||||
|
||||
let msg = `${timestamp} ${level}: ${truncateLongStrings(message?.trim(), DEBUG_MESSAGE_LENGTH)}`;
|
||||
try {
|
||||
if (level !== 'debug') {
|
||||
return msg;
|
||||
}
|
||||
|
||||
if (!metadata) {
|
||||
return msg;
|
||||
}
|
||||
|
||||
const debugValue = metadata[SPLAT_SYMBOL]?.[0];
|
||||
|
||||
if (!debugValue) {
|
||||
return msg;
|
||||
}
|
||||
|
||||
if (debugValue && Array.isArray(debugValue)) {
|
||||
msg += `\n${JSON.stringify(debugValue.map(condenseArray))}`;
|
||||
return msg;
|
||||
}
|
||||
|
||||
if (typeof debugValue !== 'object') {
|
||||
return (msg += ` ${debugValue}`);
|
||||
}
|
||||
|
||||
msg += '\n{';
|
||||
|
||||
const copy = klona(metadata);
|
||||
traverse(copy).forEach(function (value) {
|
||||
if (typeof this?.key === 'symbol') {
|
||||
return;
|
||||
}
|
||||
|
||||
let _parentKey = '';
|
||||
const parent = this.parent;
|
||||
|
||||
if (typeof parent?.key !== 'symbol' && parent?.key) {
|
||||
_parentKey = parent.key;
|
||||
}
|
||||
|
||||
const parentKey = `${parent && parent.notRoot ? _parentKey + '.' : ''}`;
|
||||
|
||||
const tabs = `${parent && parent.notRoot ? ' ' : ' '}`;
|
||||
|
||||
const currentKey = this?.key ?? 'unknown';
|
||||
|
||||
if (this.isLeaf && typeof value === 'string') {
|
||||
const truncatedText = truncateLongStrings(value);
|
||||
msg += `\n${tabs}${parentKey}${currentKey}: ${JSON.stringify(truncatedText)},`;
|
||||
} else if (this.notLeaf && Array.isArray(value) && value.length > 0) {
|
||||
const currentMessage = `\n${tabs}// ${value.length} ${currentKey.replace(/s$/, '')}(s)`;
|
||||
this.update(currentMessage, true);
|
||||
msg += currentMessage;
|
||||
const stringifiedArray = value.map(condenseArray);
|
||||
msg += `\n${tabs}${parentKey}${currentKey}: [${stringifiedArray}],`;
|
||||
} else if (this.isLeaf && typeof value === 'function') {
|
||||
msg += `\n${tabs}${parentKey}${currentKey}: function,`;
|
||||
} else if (this.isLeaf) {
|
||||
msg += `\n${tabs}${parentKey}${currentKey}: ${value},`;
|
||||
}
|
||||
});
|
||||
|
||||
msg += '\n}';
|
||||
return msg;
|
||||
} catch (e) {
|
||||
return (msg += `\n[LOGGER PARSING ERROR] ${e.message}`);
|
||||
}
|
||||
});
|
||||
|
||||
const jsonTruncateFormat = winston.format((info) => {
|
||||
const truncateLongStrings = (str, maxLength) => {
|
||||
return str.length > maxLength ? str.substring(0, maxLength) + '...' : str;
|
||||
};
|
||||
|
||||
const seen = new WeakSet();
|
||||
|
||||
const truncateObject = (obj) => {
|
||||
if (typeof obj !== 'object' || obj === null) {
|
||||
return obj;
|
||||
}
|
||||
|
||||
// Handle circular references
|
||||
if (seen.has(obj)) {
|
||||
return '[Circular]';
|
||||
}
|
||||
seen.add(obj);
|
||||
|
||||
if (Array.isArray(obj)) {
|
||||
return obj.map((item) => truncateObject(item));
|
||||
}
|
||||
|
||||
const newObj = {};
|
||||
Object.entries(obj).forEach(([key, value]) => {
|
||||
if (typeof value === 'string') {
|
||||
newObj[key] = truncateLongStrings(value, CONSOLE_JSON_STRING_LENGTH);
|
||||
} else {
|
||||
newObj[key] = truncateObject(value);
|
||||
}
|
||||
});
|
||||
return newObj;
|
||||
};
|
||||
|
||||
return truncateObject(info);
|
||||
});
|
||||
|
||||
module.exports = {
|
||||
redactFormat,
|
||||
redactMessage,
|
||||
debugTraverse,
|
||||
jsonTruncateFormat,
|
||||
};
|
||||
14
api/config/paths.js
Normal file
14
api/config/paths.js
Normal file
@ -0,0 +1,14 @@
|
||||
const path = require('path');
|
||||
|
||||
module.exports = {
|
||||
root: path.resolve(__dirname, '..', '..'),
|
||||
uploads: path.resolve(__dirname, '..', '..', 'uploads'),
|
||||
clientPath: path.resolve(__dirname, '..', '..', 'client'),
|
||||
dist: path.resolve(__dirname, '..', '..', 'client', 'dist'),
|
||||
publicPath: path.resolve(__dirname, '..', '..', 'client', 'public'),
|
||||
fonts: path.resolve(__dirname, '..', '..', 'client', 'public', 'fonts'),
|
||||
assets: path.resolve(__dirname, '..', '..', 'client', 'public', 'assets'),
|
||||
imageOutput: path.resolve(__dirname, '..', '..', 'client', 'public', 'images'),
|
||||
structuredTools: path.resolve(__dirname, '..', 'app', 'clients', 'tools', 'structured'),
|
||||
pluginManifest: path.resolve(__dirname, '..', 'app', 'clients', 'tools', 'manifest.json'),
|
||||
};
|
||||
159
api/config/winston.js
Normal file
159
api/config/winston.js
Normal file
@ -0,0 +1,159 @@
|
||||
const path = require('path');
|
||||
const fs = require('fs');
|
||||
const winston = require('winston');
|
||||
require('winston-daily-rotate-file');
|
||||
const { redactFormat, redactMessage, debugTraverse, jsonTruncateFormat } = require('./parsers');
|
||||
|
||||
/**
|
||||
* Determine the log directory.
|
||||
* Priority:
|
||||
* 1. LIBRECHAT_LOG_DIR environment variable (allows user override)
|
||||
* 2. /app/logs if running in Docker (bind-mounted with correct permissions)
|
||||
* 3. api/logs relative to this file (local development)
|
||||
*/
|
||||
const getLogDir = () => {
|
||||
if (process.env.LIBRECHAT_LOG_DIR) {
|
||||
return process.env.LIBRECHAT_LOG_DIR;
|
||||
}
|
||||
|
||||
// Check if running in Docker container (cwd is /app)
|
||||
if (process.cwd() === '/app') {
|
||||
const dockerLogDir = '/app/logs';
|
||||
// Ensure the directory exists
|
||||
if (!fs.existsSync(dockerLogDir)) {
|
||||
fs.mkdirSync(dockerLogDir, { recursive: true });
|
||||
}
|
||||
return dockerLogDir;
|
||||
}
|
||||
|
||||
// Local development: use api/logs relative to this file
|
||||
return path.join(__dirname, '..', 'logs');
|
||||
};
|
||||
|
||||
const logDir = getLogDir();
|
||||
|
||||
const { NODE_ENV, DEBUG_LOGGING = true, CONSOLE_JSON = false, DEBUG_CONSOLE = false } = process.env;
|
||||
|
||||
const useConsoleJson =
|
||||
(typeof CONSOLE_JSON === 'string' && CONSOLE_JSON?.toLowerCase() === 'true') ||
|
||||
CONSOLE_JSON === true;
|
||||
|
||||
const useDebugConsole =
|
||||
(typeof DEBUG_CONSOLE === 'string' && DEBUG_CONSOLE?.toLowerCase() === 'true') ||
|
||||
DEBUG_CONSOLE === true;
|
||||
|
||||
const useDebugLogging =
|
||||
(typeof DEBUG_LOGGING === 'string' && DEBUG_LOGGING?.toLowerCase() === 'true') ||
|
||||
DEBUG_LOGGING === true;
|
||||
|
||||
const levels = {
|
||||
error: 0,
|
||||
warn: 1,
|
||||
info: 2,
|
||||
http: 3,
|
||||
verbose: 4,
|
||||
debug: 5,
|
||||
activity: 6,
|
||||
silly: 7,
|
||||
};
|
||||
|
||||
winston.addColors({
|
||||
info: 'green', // fontStyle color
|
||||
warn: 'italic yellow',
|
||||
error: 'red',
|
||||
debug: 'blue',
|
||||
});
|
||||
|
||||
const level = () => {
|
||||
const env = NODE_ENV || 'development';
|
||||
const isDevelopment = env === 'development';
|
||||
return isDevelopment ? 'debug' : 'warn';
|
||||
};
|
||||
|
||||
const fileFormat = winston.format.combine(
|
||||
redactFormat(),
|
||||
winston.format.timestamp({ format: () => new Date().toISOString() }),
|
||||
winston.format.errors({ stack: true }),
|
||||
winston.format.splat(),
|
||||
// redactErrors(),
|
||||
);
|
||||
|
||||
const transports = [
|
||||
new winston.transports.DailyRotateFile({
|
||||
level: 'error',
|
||||
filename: `${logDir}/error-%DATE%.log`,
|
||||
datePattern: 'YYYY-MM-DD',
|
||||
zippedArchive: true,
|
||||
maxSize: '20m',
|
||||
maxFiles: '14d',
|
||||
format: fileFormat,
|
||||
}),
|
||||
];
|
||||
|
||||
if (useDebugLogging) {
|
||||
transports.push(
|
||||
new winston.transports.DailyRotateFile({
|
||||
level: 'debug',
|
||||
filename: `${logDir}/debug-%DATE%.log`,
|
||||
datePattern: 'YYYY-MM-DD',
|
||||
zippedArchive: true,
|
||||
maxSize: '20m',
|
||||
maxFiles: '14d',
|
||||
format: winston.format.combine(fileFormat, debugTraverse),
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
const consoleFormat = winston.format.combine(
|
||||
redactFormat(),
|
||||
winston.format.colorize({ all: true }),
|
||||
winston.format.timestamp({ format: 'YYYY-MM-DD HH:mm:ss' }),
|
||||
// redactErrors(),
|
||||
winston.format.printf((info) => {
|
||||
const message = `${info.timestamp} ${info.level}: ${info.message}`;
|
||||
if (info.level.includes('error')) {
|
||||
return redactMessage(message);
|
||||
}
|
||||
|
||||
return message;
|
||||
}),
|
||||
);
|
||||
|
||||
// Determine console log level
|
||||
let consoleLogLevel = 'info';
|
||||
if (useDebugConsole) {
|
||||
consoleLogLevel = 'debug';
|
||||
}
|
||||
|
||||
if (useDebugConsole) {
|
||||
transports.push(
|
||||
new winston.transports.Console({
|
||||
level: consoleLogLevel,
|
||||
format: useConsoleJson
|
||||
? winston.format.combine(fileFormat, jsonTruncateFormat(), winston.format.json())
|
||||
: winston.format.combine(fileFormat, debugTraverse),
|
||||
}),
|
||||
);
|
||||
} else if (useConsoleJson) {
|
||||
transports.push(
|
||||
new winston.transports.Console({
|
||||
level: consoleLogLevel,
|
||||
format: winston.format.combine(fileFormat, jsonTruncateFormat(), winston.format.json()),
|
||||
}),
|
||||
);
|
||||
} else {
|
||||
transports.push(
|
||||
new winston.transports.Console({
|
||||
level: consoleLogLevel,
|
||||
format: consoleFormat,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
const logger = winston.createLogger({
|
||||
level: level(),
|
||||
levels,
|
||||
transports,
|
||||
});
|
||||
|
||||
module.exports = logger;
|
||||
79
api/db/connect.js
Normal file
79
api/db/connect.js
Normal file
@ -0,0 +1,79 @@
|
||||
require('dotenv').config();
|
||||
const { isEnabled } = require('@librechat/api');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
|
||||
const mongoose = require('mongoose');
|
||||
const MONGO_URI = process.env.MONGO_URI;
|
||||
|
||||
if (!MONGO_URI) {
|
||||
throw new Error('Please define the MONGO_URI environment variable');
|
||||
}
|
||||
/** The maximum number of connections in the connection pool. */
|
||||
const maxPoolSize = parseInt(process.env.MONGO_MAX_POOL_SIZE) || undefined;
|
||||
/** The minimum number of connections in the connection pool. */
|
||||
const minPoolSize = parseInt(process.env.MONGO_MIN_POOL_SIZE) || undefined;
|
||||
/** The maximum number of connections that may be in the process of being established concurrently by the connection pool. */
|
||||
const maxConnecting = parseInt(process.env.MONGO_MAX_CONNECTING) || undefined;
|
||||
/** The maximum number of milliseconds that a connection can remain idle in the pool before being removed and closed. */
|
||||
const maxIdleTimeMS = parseInt(process.env.MONGO_MAX_IDLE_TIME_MS) || undefined;
|
||||
/** The maximum time in milliseconds that a thread can wait for a connection to become available. */
|
||||
const waitQueueTimeoutMS = parseInt(process.env.MONGO_WAIT_QUEUE_TIMEOUT_MS) || undefined;
|
||||
/** Set to false to disable automatic index creation for all models associated with this connection. */
|
||||
const autoIndex =
|
||||
process.env.MONGO_AUTO_INDEX != undefined
|
||||
? isEnabled(process.env.MONGO_AUTO_INDEX) || false
|
||||
: undefined;
|
||||
|
||||
/** Set to `false` to disable Mongoose automatically calling `createCollection()` on every model created on this connection. */
|
||||
const autoCreate =
|
||||
process.env.MONGO_AUTO_CREATE != undefined
|
||||
? isEnabled(process.env.MONGO_AUTO_CREATE) || false
|
||||
: undefined;
|
||||
/**
|
||||
* Global is used here to maintain a cached connection across hot reloads
|
||||
* in development. This prevents connections growing exponentially
|
||||
* during API Route usage.
|
||||
*/
|
||||
let cached = global.mongoose;
|
||||
|
||||
if (!cached) {
|
||||
cached = global.mongoose = { conn: null, promise: null };
|
||||
}
|
||||
|
||||
async function connectDb() {
|
||||
if (cached.conn && cached.conn?._readyState === 1) {
|
||||
return cached.conn;
|
||||
}
|
||||
|
||||
const disconnected = cached.conn && cached.conn?._readyState !== 1;
|
||||
if (!cached.promise || disconnected) {
|
||||
const opts = {
|
||||
bufferCommands: false,
|
||||
...(maxPoolSize ? { maxPoolSize } : {}),
|
||||
...(minPoolSize ? { minPoolSize } : {}),
|
||||
...(maxConnecting ? { maxConnecting } : {}),
|
||||
...(maxIdleTimeMS ? { maxIdleTimeMS } : {}),
|
||||
...(waitQueueTimeoutMS ? { waitQueueTimeoutMS } : {}),
|
||||
...(autoIndex != undefined ? { autoIndex } : {}),
|
||||
...(autoCreate != undefined ? { autoCreate } : {}),
|
||||
// useNewUrlParser: true,
|
||||
// useUnifiedTopology: true,
|
||||
// bufferMaxEntries: 0,
|
||||
// useFindAndModify: true,
|
||||
// useCreateIndex: true
|
||||
};
|
||||
logger.info('Mongo Connection options');
|
||||
logger.info(JSON.stringify(opts, null, 2));
|
||||
mongoose.set('strictQuery', true);
|
||||
cached.promise = mongoose.connect(MONGO_URI, opts).then((mongoose) => {
|
||||
return mongoose;
|
||||
});
|
||||
}
|
||||
cached.conn = await cached.promise;
|
||||
|
||||
return cached.conn;
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
connectDb,
|
||||
};
|
||||
8
api/db/index.js
Normal file
8
api/db/index.js
Normal file
@ -0,0 +1,8 @@
|
||||
const mongoose = require('mongoose');
|
||||
const { createModels } = require('@librechat/data-schemas');
|
||||
const { connectDb } = require('./connect');
|
||||
const indexSync = require('./indexSync');
|
||||
|
||||
createModels(mongoose);
|
||||
|
||||
module.exports = { connectDb, indexSync };
|
||||
363
api/db/indexSync.js
Normal file
363
api/db/indexSync.js
Normal file
@ -0,0 +1,363 @@
|
||||
const mongoose = require('mongoose');
|
||||
const { MeiliSearch } = require('meilisearch');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { CacheKeys } = require('librechat-data-provider');
|
||||
const { isEnabled, FlowStateManager } = require('@librechat/api');
|
||||
const { getLogStores } = require('~/cache');
|
||||
const { batchResetMeiliFlags } = require('./utils');
|
||||
|
||||
const Conversation = mongoose.models.Conversation;
|
||||
const Message = mongoose.models.Message;
|
||||
|
||||
const searchEnabled = isEnabled(process.env.SEARCH);
|
||||
const indexingDisabled = isEnabled(process.env.MEILI_NO_SYNC);
|
||||
let currentTimeout = null;
|
||||
|
||||
const defaultSyncThreshold = 1000;
|
||||
const syncThreshold = process.env.MEILI_SYNC_THRESHOLD
|
||||
? parseInt(process.env.MEILI_SYNC_THRESHOLD, 10)
|
||||
: defaultSyncThreshold;
|
||||
|
||||
class MeiliSearchClient {
|
||||
static instance = null;
|
||||
|
||||
static getInstance() {
|
||||
if (!MeiliSearchClient.instance) {
|
||||
if (!process.env.MEILI_HOST || !process.env.MEILI_MASTER_KEY) {
|
||||
throw new Error('Meilisearch configuration is missing.');
|
||||
}
|
||||
MeiliSearchClient.instance = new MeiliSearch({
|
||||
host: process.env.MEILI_HOST,
|
||||
apiKey: process.env.MEILI_MASTER_KEY,
|
||||
});
|
||||
}
|
||||
return MeiliSearchClient.instance;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Deletes documents from MeiliSearch index that are missing the user field
|
||||
* @param {import('meilisearch').Index} index - MeiliSearch index instance
|
||||
* @param {string} indexName - Name of the index for logging
|
||||
* @returns {Promise<number>} - Number of documents deleted
|
||||
*/
|
||||
async function deleteDocumentsWithoutUserField(index, indexName) {
|
||||
let deletedCount = 0;
|
||||
let offset = 0;
|
||||
const batchSize = 1000;
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
const searchResult = await index.search('', {
|
||||
limit: batchSize,
|
||||
offset: offset,
|
||||
});
|
||||
|
||||
if (searchResult.hits.length === 0) {
|
||||
break;
|
||||
}
|
||||
|
||||
const idsToDelete = searchResult.hits.filter((hit) => !hit.user).map((hit) => hit.id);
|
||||
|
||||
if (idsToDelete.length > 0) {
|
||||
logger.info(
|
||||
`[indexSync] Deleting ${idsToDelete.length} documents without user field from ${indexName} index`,
|
||||
);
|
||||
await index.deleteDocuments(idsToDelete);
|
||||
deletedCount += idsToDelete.length;
|
||||
}
|
||||
|
||||
if (searchResult.hits.length < batchSize) {
|
||||
break;
|
||||
}
|
||||
|
||||
offset += batchSize;
|
||||
}
|
||||
|
||||
if (deletedCount > 0) {
|
||||
logger.info(`[indexSync] Deleted ${deletedCount} orphaned documents from ${indexName} index`);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error(`[indexSync] Error deleting documents from ${indexName}:`, error);
|
||||
}
|
||||
|
||||
return deletedCount;
|
||||
}
|
||||
|
||||
/**
|
||||
* Ensures indexes have proper filterable attributes configured and checks if documents have user field
|
||||
* @param {MeiliSearch} client - MeiliSearch client instance
|
||||
* @returns {Promise<{settingsUpdated: boolean, orphanedDocsFound: boolean}>} - Status of what was done
|
||||
*/
|
||||
async function ensureFilterableAttributes(client) {
|
||||
let settingsUpdated = false;
|
||||
let hasOrphanedDocs = false;
|
||||
|
||||
try {
|
||||
// Check and update messages index
|
||||
try {
|
||||
const messagesIndex = client.index('messages');
|
||||
const settings = await messagesIndex.getSettings();
|
||||
|
||||
if (!settings.filterableAttributes || !settings.filterableAttributes.includes('user')) {
|
||||
logger.info('[indexSync] Configuring messages index to filter by user...');
|
||||
await messagesIndex.updateSettings({
|
||||
filterableAttributes: ['user'],
|
||||
});
|
||||
logger.info('[indexSync] Messages index configured for user filtering');
|
||||
settingsUpdated = true;
|
||||
}
|
||||
|
||||
// Check if existing documents have user field indexed
|
||||
try {
|
||||
const searchResult = await messagesIndex.search('', { limit: 1 });
|
||||
if (searchResult.hits.length > 0 && !searchResult.hits[0].user) {
|
||||
logger.info(
|
||||
'[indexSync] Existing messages missing user field, will clean up orphaned documents...',
|
||||
);
|
||||
hasOrphanedDocs = true;
|
||||
}
|
||||
} catch (searchError) {
|
||||
logger.debug('[indexSync] Could not check message documents:', searchError.message);
|
||||
}
|
||||
} catch (error) {
|
||||
if (error.code !== 'index_not_found') {
|
||||
logger.warn('[indexSync] Could not check/update messages index settings:', error.message);
|
||||
}
|
||||
}
|
||||
|
||||
// Check and update conversations index
|
||||
try {
|
||||
const convosIndex = client.index('convos');
|
||||
const settings = await convosIndex.getSettings();
|
||||
|
||||
if (!settings.filterableAttributes || !settings.filterableAttributes.includes('user')) {
|
||||
logger.info('[indexSync] Configuring convos index to filter by user...');
|
||||
await convosIndex.updateSettings({
|
||||
filterableAttributes: ['user'],
|
||||
});
|
||||
logger.info('[indexSync] Convos index configured for user filtering');
|
||||
settingsUpdated = true;
|
||||
}
|
||||
|
||||
// Check if existing documents have user field indexed
|
||||
try {
|
||||
const searchResult = await convosIndex.search('', { limit: 1 });
|
||||
if (searchResult.hits.length > 0 && !searchResult.hits[0].user) {
|
||||
logger.info(
|
||||
'[indexSync] Existing conversations missing user field, will clean up orphaned documents...',
|
||||
);
|
||||
hasOrphanedDocs = true;
|
||||
}
|
||||
} catch (searchError) {
|
||||
logger.debug('[indexSync] Could not check conversation documents:', searchError.message);
|
||||
}
|
||||
} catch (error) {
|
||||
if (error.code !== 'index_not_found') {
|
||||
logger.warn('[indexSync] Could not check/update convos index settings:', error.message);
|
||||
}
|
||||
}
|
||||
|
||||
// If either index has orphaned documents, clean them up (but don't force resync)
|
||||
if (hasOrphanedDocs) {
|
||||
try {
|
||||
const messagesIndex = client.index('messages');
|
||||
await deleteDocumentsWithoutUserField(messagesIndex, 'messages');
|
||||
} catch (error) {
|
||||
logger.debug('[indexSync] Could not clean up messages:', error.message);
|
||||
}
|
||||
|
||||
try {
|
||||
const convosIndex = client.index('convos');
|
||||
await deleteDocumentsWithoutUserField(convosIndex, 'convos');
|
||||
} catch (error) {
|
||||
logger.debug('[indexSync] Could not clean up convos:', error.message);
|
||||
}
|
||||
|
||||
logger.info('[indexSync] Orphaned documents cleaned up without forcing resync.');
|
||||
}
|
||||
|
||||
if (settingsUpdated) {
|
||||
logger.info('[indexSync] Index settings updated. Full re-sync will be triggered.');
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error('[indexSync] Error ensuring filterable attributes:', error);
|
||||
}
|
||||
|
||||
return { settingsUpdated, orphanedDocsFound: hasOrphanedDocs };
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs the actual sync operations for messages and conversations
|
||||
* @param {FlowStateManager} flowManager - Flow state manager instance
|
||||
* @param {string} flowId - Flow identifier
|
||||
* @param {string} flowType - Flow type
|
||||
*/
|
||||
async function performSync(flowManager, flowId, flowType) {
|
||||
try {
|
||||
if (indexingDisabled === true) {
|
||||
logger.info('[indexSync] Indexing is disabled, skipping...');
|
||||
return { messagesSync: false, convosSync: false };
|
||||
}
|
||||
|
||||
const client = MeiliSearchClient.getInstance();
|
||||
|
||||
const { status } = await client.health();
|
||||
if (status !== 'available') {
|
||||
throw new Error('Meilisearch not available');
|
||||
}
|
||||
|
||||
/** Ensures indexes have proper filterable attributes configured */
|
||||
const { settingsUpdated, orphanedDocsFound: _orphanedDocsFound } =
|
||||
await ensureFilterableAttributes(client);
|
||||
|
||||
let messagesSync = false;
|
||||
let convosSync = false;
|
||||
|
||||
// Only reset flags if settings were actually updated (not just for orphaned doc cleanup)
|
||||
if (settingsUpdated) {
|
||||
logger.info(
|
||||
'[indexSync] Settings updated. Forcing full re-sync to reindex with new configuration...',
|
||||
);
|
||||
|
||||
// Reset sync flags to force full re-sync
|
||||
await batchResetMeiliFlags(Message.collection);
|
||||
await batchResetMeiliFlags(Conversation.collection);
|
||||
}
|
||||
|
||||
// Check if we need to sync messages
|
||||
logger.info('[indexSync] Requesting message sync progress...');
|
||||
const messageProgress = await Message.getSyncProgress();
|
||||
if (!messageProgress.isComplete || settingsUpdated) {
|
||||
logger.info(
|
||||
`[indexSync] Messages need syncing: ${messageProgress.totalProcessed}/${messageProgress.totalDocuments} indexed`,
|
||||
);
|
||||
|
||||
const messageCount = messageProgress.totalDocuments;
|
||||
const messagesIndexed = messageProgress.totalProcessed;
|
||||
const unindexedMessages = messageCount - messagesIndexed;
|
||||
|
||||
if (settingsUpdated || unindexedMessages > syncThreshold) {
|
||||
logger.info(`[indexSync] Starting message sync (${unindexedMessages} unindexed)`);
|
||||
await Message.syncWithMeili();
|
||||
messagesSync = true;
|
||||
} else if (unindexedMessages > 0) {
|
||||
logger.info(
|
||||
`[indexSync] ${unindexedMessages} messages unindexed (below threshold: ${syncThreshold}, skipping)`,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
logger.info(
|
||||
`[indexSync] Messages are fully synced: ${messageProgress.totalProcessed}/${messageProgress.totalDocuments}`,
|
||||
);
|
||||
}
|
||||
|
||||
// Check if we need to sync conversations
|
||||
const convoProgress = await Conversation.getSyncProgress();
|
||||
if (!convoProgress.isComplete || settingsUpdated) {
|
||||
logger.info(
|
||||
`[indexSync] Conversations need syncing: ${convoProgress.totalProcessed}/${convoProgress.totalDocuments} indexed`,
|
||||
);
|
||||
|
||||
const convoCount = convoProgress.totalDocuments;
|
||||
const convosIndexed = convoProgress.totalProcessed;
|
||||
|
||||
const unindexedConvos = convoCount - convosIndexed;
|
||||
if (settingsUpdated || unindexedConvos > syncThreshold) {
|
||||
logger.info(`[indexSync] Starting convos sync (${unindexedConvos} unindexed)`);
|
||||
await Conversation.syncWithMeili();
|
||||
convosSync = true;
|
||||
} else if (unindexedConvos > 0) {
|
||||
logger.info(
|
||||
`[indexSync] ${unindexedConvos} convos unindexed (below threshold: ${syncThreshold}, skipping)`,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
logger.info(
|
||||
`[indexSync] Conversations are fully synced: ${convoProgress.totalProcessed}/${convoProgress.totalDocuments}`,
|
||||
);
|
||||
}
|
||||
|
||||
return { messagesSync, convosSync };
|
||||
} finally {
|
||||
if (indexingDisabled === true) {
|
||||
logger.info('[indexSync] Indexing is disabled, skipping cleanup...');
|
||||
} else if (flowManager && flowId && flowType) {
|
||||
try {
|
||||
await flowManager.deleteFlow(flowId, flowType);
|
||||
logger.debug('[indexSync] Flow state cleaned up');
|
||||
} catch (cleanupErr) {
|
||||
logger.debug('[indexSync] Could not clean up flow state:', cleanupErr.message);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Main index sync function that uses FlowStateManager to prevent concurrent execution
|
||||
*/
|
||||
async function indexSync() {
|
||||
if (!searchEnabled) {
|
||||
return;
|
||||
}
|
||||
|
||||
logger.info('[indexSync] Starting index synchronization check...');
|
||||
|
||||
// Get or create FlowStateManager instance
|
||||
const flowsCache = getLogStores(CacheKeys.FLOWS);
|
||||
if (!flowsCache) {
|
||||
logger.warn('[indexSync] Flows cache not available, falling back to direct sync');
|
||||
return await performSync(null, null, null);
|
||||
}
|
||||
|
||||
const flowManager = new FlowStateManager(flowsCache, {
|
||||
ttl: 60000 * 10, // 10 minutes TTL for sync operations
|
||||
});
|
||||
|
||||
// Use a unique flow ID for the sync operation
|
||||
const flowId = 'meili-index-sync';
|
||||
const flowType = 'MEILI_SYNC';
|
||||
|
||||
try {
|
||||
// This will only execute the handler if no other instance is running the sync
|
||||
const result = await flowManager.createFlowWithHandler(flowId, flowType, () =>
|
||||
performSync(flowManager, flowId, flowType),
|
||||
);
|
||||
|
||||
if (result.messagesSync || result.convosSync) {
|
||||
logger.info('[indexSync] Sync completed successfully');
|
||||
} else {
|
||||
logger.debug('[indexSync] No sync was needed');
|
||||
}
|
||||
|
||||
return result;
|
||||
} catch (err) {
|
||||
if (err.message.includes('flow already exists')) {
|
||||
logger.info('[indexSync] Sync already running on another instance');
|
||||
return;
|
||||
}
|
||||
|
||||
if (err.message.includes('not found')) {
|
||||
logger.debug('[indexSync] Creating indices...');
|
||||
currentTimeout = setTimeout(async () => {
|
||||
try {
|
||||
await Message.syncWithMeili();
|
||||
await Conversation.syncWithMeili();
|
||||
} catch (err) {
|
||||
logger.error('[indexSync] Trouble creating indices, try restarting the server.', err);
|
||||
}
|
||||
}, 750);
|
||||
} else if (err.message.includes('Meilisearch not configured')) {
|
||||
logger.info('[indexSync] Meilisearch not configured, search will be disabled.');
|
||||
} else {
|
||||
logger.error('[indexSync] error', err);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
process.on('exit', () => {
|
||||
logger.debug('[indexSync] Clearing sync timeouts before exiting...');
|
||||
clearTimeout(currentTimeout);
|
||||
});
|
||||
|
||||
module.exports = indexSync;
|
||||
465
api/db/indexSync.spec.js
Normal file
465
api/db/indexSync.spec.js
Normal file
@ -0,0 +1,465 @@
|
||||
/**
|
||||
* Unit tests for performSync() function in indexSync.js
|
||||
*
|
||||
* Tests use real mongoose with mocked model methods, only mocking external calls.
|
||||
*/
|
||||
|
||||
const mongoose = require('mongoose');
|
||||
|
||||
// Mock only external dependencies (not internal classes/models)
|
||||
const mockLogger = {
|
||||
info: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
error: jest.fn(),
|
||||
debug: jest.fn(),
|
||||
};
|
||||
|
||||
const mockMeiliHealth = jest.fn();
|
||||
const mockMeiliIndex = jest.fn();
|
||||
const mockBatchResetMeiliFlags = jest.fn();
|
||||
const mockIsEnabled = jest.fn();
|
||||
const mockGetLogStores = jest.fn();
|
||||
|
||||
// Create mock models that will be reused
|
||||
const createMockModel = (collectionName) => ({
|
||||
collection: { name: collectionName },
|
||||
getSyncProgress: jest.fn(),
|
||||
syncWithMeili: jest.fn(),
|
||||
countDocuments: jest.fn(),
|
||||
});
|
||||
|
||||
const originalMessageModel = mongoose.models.Message;
|
||||
const originalConversationModel = mongoose.models.Conversation;
|
||||
|
||||
// Mock external modules
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: mockLogger,
|
||||
}));
|
||||
|
||||
jest.mock('meilisearch', () => ({
|
||||
MeiliSearch: jest.fn(() => ({
|
||||
health: mockMeiliHealth,
|
||||
index: mockMeiliIndex,
|
||||
})),
|
||||
}));
|
||||
|
||||
jest.mock('./utils', () => ({
|
||||
batchResetMeiliFlags: mockBatchResetMeiliFlags,
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
isEnabled: mockIsEnabled,
|
||||
FlowStateManager: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('~/cache', () => ({
|
||||
getLogStores: mockGetLogStores,
|
||||
}));
|
||||
|
||||
// Set environment before module load
|
||||
process.env.MEILI_HOST = 'http://localhost:7700';
|
||||
process.env.MEILI_MASTER_KEY = 'test-key';
|
||||
process.env.SEARCH = 'true';
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000'; // Set threshold before module loads
|
||||
|
||||
describe('performSync() - syncThreshold logic', () => {
|
||||
const ORIGINAL_ENV = process.env;
|
||||
let Message;
|
||||
let Conversation;
|
||||
|
||||
beforeAll(() => {
|
||||
Message = createMockModel('messages');
|
||||
Conversation = createMockModel('conversations');
|
||||
|
||||
mongoose.models.Message = Message;
|
||||
mongoose.models.Conversation = Conversation;
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
// Reset all mocks
|
||||
jest.clearAllMocks();
|
||||
// Reset modules to ensure fresh load of indexSync.js and its top-level consts (like syncThreshold)
|
||||
jest.resetModules();
|
||||
|
||||
// Set up environment
|
||||
process.env = { ...ORIGINAL_ENV };
|
||||
process.env.MEILI_HOST = 'http://localhost:7700';
|
||||
process.env.MEILI_MASTER_KEY = 'test-key';
|
||||
process.env.SEARCH = 'true';
|
||||
delete process.env.MEILI_NO_SYNC;
|
||||
|
||||
// Re-ensure models are available in mongoose after resetModules
|
||||
// We must require mongoose again to get the fresh instance that indexSync will use
|
||||
const mongoose = require('mongoose');
|
||||
mongoose.models.Message = Message;
|
||||
mongoose.models.Conversation = Conversation;
|
||||
|
||||
// Mock isEnabled
|
||||
mockIsEnabled.mockImplementation((val) => val === 'true' || val === true);
|
||||
|
||||
// Mock MeiliSearch client responses
|
||||
mockMeiliHealth.mockResolvedValue({ status: 'available' });
|
||||
mockMeiliIndex.mockReturnValue({
|
||||
getSettings: jest.fn().mockResolvedValue({ filterableAttributes: ['user'] }),
|
||||
updateSettings: jest.fn().mockResolvedValue({}),
|
||||
search: jest.fn().mockResolvedValue({ hits: [] }),
|
||||
});
|
||||
|
||||
mockBatchResetMeiliFlags.mockResolvedValue(undefined);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
process.env = ORIGINAL_ENV;
|
||||
});
|
||||
|
||||
afterAll(() => {
|
||||
mongoose.models.Message = originalMessageModel;
|
||||
mongoose.models.Conversation = originalConversationModel;
|
||||
});
|
||||
|
||||
test('triggers sync when unindexed messages exceed syncThreshold', async () => {
|
||||
// Arrange: Set threshold before module load
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Arrange: 1050 unindexed messages > 1000 threshold
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 1150, // 1050 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 50,
|
||||
totalDocuments: 50,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
Message.syncWithMeili.mockResolvedValue(undefined);
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: No countDocuments calls
|
||||
expect(Message.countDocuments).not.toHaveBeenCalled();
|
||||
expect(Conversation.countDocuments).not.toHaveBeenCalled();
|
||||
|
||||
// Assert: Message sync triggered because 1050 > 1000
|
||||
expect(Message.syncWithMeili).toHaveBeenCalledTimes(1);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Messages need syncing: 100/1150 indexed',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Starting message sync (1050 unindexed)',
|
||||
);
|
||||
|
||||
// Assert: Conversation sync NOT triggered (already complete)
|
||||
expect(Conversation.syncWithMeili).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test('skips sync when unindexed messages are below syncThreshold', async () => {
|
||||
// Arrange: 50 unindexed messages < 1000 threshold
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 150, // 50 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 50,
|
||||
totalDocuments: 50,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: No countDocuments calls
|
||||
expect(Message.countDocuments).not.toHaveBeenCalled();
|
||||
expect(Conversation.countDocuments).not.toHaveBeenCalled();
|
||||
|
||||
// Assert: Message sync NOT triggered because 50 < 1000
|
||||
expect(Message.syncWithMeili).not.toHaveBeenCalled();
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Messages need syncing: 100/150 indexed',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] 50 messages unindexed (below threshold: 1000, skipping)',
|
||||
);
|
||||
|
||||
// Assert: Conversation sync NOT triggered (already complete)
|
||||
expect(Conversation.syncWithMeili).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test('respects syncThreshold at boundary (exactly at threshold)', async () => {
|
||||
// Arrange: 1000 unindexed messages = 1000 threshold (NOT greater than)
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 1100, // 1000 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 0,
|
||||
totalDocuments: 0,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: No countDocuments calls
|
||||
expect(Message.countDocuments).not.toHaveBeenCalled();
|
||||
|
||||
// Assert: Message sync NOT triggered because 1000 is NOT > 1000
|
||||
expect(Message.syncWithMeili).not.toHaveBeenCalled();
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Messages need syncing: 100/1100 indexed',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] 1000 messages unindexed (below threshold: 1000, skipping)',
|
||||
);
|
||||
});
|
||||
|
||||
test('triggers sync when unindexed is threshold + 1', async () => {
|
||||
// Arrange: 1001 unindexed messages > 1000 threshold
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 1101, // 1001 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 0,
|
||||
totalDocuments: 0,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
Message.syncWithMeili.mockResolvedValue(undefined);
|
||||
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: No countDocuments calls
|
||||
expect(Message.countDocuments).not.toHaveBeenCalled();
|
||||
|
||||
// Assert: Message sync triggered because 1001 > 1000
|
||||
expect(Message.syncWithMeili).toHaveBeenCalledTimes(1);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Messages need syncing: 100/1101 indexed',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Starting message sync (1001 unindexed)',
|
||||
);
|
||||
});
|
||||
|
||||
test('uses totalDocuments from convoProgress for conversation sync decisions', async () => {
|
||||
// Arrange: Messages complete, conversations need sync
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 100,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 50,
|
||||
totalDocuments: 1100, // 1050 unindexed > 1000 threshold
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.syncWithMeili.mockResolvedValue(undefined);
|
||||
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: No countDocuments calls (the optimization)
|
||||
expect(Message.countDocuments).not.toHaveBeenCalled();
|
||||
expect(Conversation.countDocuments).not.toHaveBeenCalled();
|
||||
|
||||
// Assert: Only conversation sync triggered
|
||||
expect(Message.syncWithMeili).not.toHaveBeenCalled();
|
||||
expect(Conversation.syncWithMeili).toHaveBeenCalledTimes(1);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Conversations need syncing: 50/1100 indexed',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Starting convos sync (1050 unindexed)',
|
||||
);
|
||||
});
|
||||
|
||||
test('skips sync when collections are fully synced', async () => {
|
||||
// Arrange: Everything already synced
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 100,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 50,
|
||||
totalDocuments: 50,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: No countDocuments calls
|
||||
expect(Message.countDocuments).not.toHaveBeenCalled();
|
||||
expect(Conversation.countDocuments).not.toHaveBeenCalled();
|
||||
|
||||
// Assert: No sync triggered
|
||||
expect(Message.syncWithMeili).not.toHaveBeenCalled();
|
||||
expect(Conversation.syncWithMeili).not.toHaveBeenCalled();
|
||||
|
||||
// Assert: Correct logs
|
||||
expect(mockLogger.info).toHaveBeenCalledWith('[indexSync] Messages are fully synced: 100/100');
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Conversations are fully synced: 50/50',
|
||||
);
|
||||
});
|
||||
|
||||
test('triggers message sync when settingsUpdated even if below syncThreshold', async () => {
|
||||
// Arrange: Only 50 unindexed messages (< 1000 threshold), but settings were updated
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 150, // 50 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 50,
|
||||
totalDocuments: 50,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
Message.syncWithMeili.mockResolvedValue(undefined);
|
||||
|
||||
// Mock settings update scenario
|
||||
mockMeiliIndex.mockReturnValue({
|
||||
getSettings: jest.fn().mockResolvedValue({ filterableAttributes: [] }), // No user field
|
||||
updateSettings: jest.fn().mockResolvedValue({}),
|
||||
search: jest.fn().mockResolvedValue({ hits: [] }),
|
||||
});
|
||||
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: Flags were reset due to settings update
|
||||
expect(mockBatchResetMeiliFlags).toHaveBeenCalledWith(Message.collection);
|
||||
expect(mockBatchResetMeiliFlags).toHaveBeenCalledWith(Conversation.collection);
|
||||
|
||||
// Assert: Message sync triggered despite being below threshold (50 < 1000)
|
||||
expect(Message.syncWithMeili).toHaveBeenCalledTimes(1);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Settings updated. Forcing full re-sync to reindex with new configuration...',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Starting message sync (50 unindexed)',
|
||||
);
|
||||
});
|
||||
|
||||
test('triggers conversation sync when settingsUpdated even if below syncThreshold', async () => {
|
||||
// Arrange: Messages complete, conversations have 50 unindexed (< 1000 threshold), but settings were updated
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 100,
|
||||
isComplete: true,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 50,
|
||||
totalDocuments: 100, // 50 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.syncWithMeili.mockResolvedValue(undefined);
|
||||
|
||||
// Mock settings update scenario
|
||||
mockMeiliIndex.mockReturnValue({
|
||||
getSettings: jest.fn().mockResolvedValue({ filterableAttributes: [] }), // No user field
|
||||
updateSettings: jest.fn().mockResolvedValue({}),
|
||||
search: jest.fn().mockResolvedValue({ hits: [] }),
|
||||
});
|
||||
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: Flags were reset due to settings update
|
||||
expect(mockBatchResetMeiliFlags).toHaveBeenCalledWith(Message.collection);
|
||||
expect(mockBatchResetMeiliFlags).toHaveBeenCalledWith(Conversation.collection);
|
||||
|
||||
// Assert: Conversation sync triggered despite being below threshold (50 < 1000)
|
||||
expect(Conversation.syncWithMeili).toHaveBeenCalledTimes(1);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Settings updated. Forcing full re-sync to reindex with new configuration...',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith('[indexSync] Starting convos sync (50 unindexed)');
|
||||
});
|
||||
|
||||
test('triggers both message and conversation sync when settingsUpdated even if both below syncThreshold', async () => {
|
||||
// Arrange: Set threshold before module load
|
||||
process.env.MEILI_SYNC_THRESHOLD = '1000';
|
||||
|
||||
// Arrange: Both have documents below threshold (50 each), but settings were updated
|
||||
Message.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 100,
|
||||
totalDocuments: 150, // 50 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Conversation.getSyncProgress.mockResolvedValue({
|
||||
totalProcessed: 50,
|
||||
totalDocuments: 100, // 50 unindexed
|
||||
isComplete: false,
|
||||
});
|
||||
|
||||
Message.syncWithMeili.mockResolvedValue(undefined);
|
||||
Conversation.syncWithMeili.mockResolvedValue(undefined);
|
||||
|
||||
// Mock settings update scenario
|
||||
mockMeiliIndex.mockReturnValue({
|
||||
getSettings: jest.fn().mockResolvedValue({ filterableAttributes: [] }), // No user field
|
||||
updateSettings: jest.fn().mockResolvedValue({}),
|
||||
search: jest.fn().mockResolvedValue({ hits: [] }),
|
||||
});
|
||||
|
||||
// Act
|
||||
const indexSync = require('./indexSync');
|
||||
await indexSync();
|
||||
|
||||
// Assert: Flags were reset due to settings update
|
||||
expect(mockBatchResetMeiliFlags).toHaveBeenCalledWith(Message.collection);
|
||||
expect(mockBatchResetMeiliFlags).toHaveBeenCalledWith(Conversation.collection);
|
||||
|
||||
// Assert: Both syncs triggered despite both being below threshold
|
||||
expect(Message.syncWithMeili).toHaveBeenCalledTimes(1);
|
||||
expect(Conversation.syncWithMeili).toHaveBeenCalledTimes(1);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Settings updated. Forcing full re-sync to reindex with new configuration...',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
'[indexSync] Starting message sync (50 unindexed)',
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith('[indexSync] Starting convos sync (50 unindexed)');
|
||||
});
|
||||
});
|
||||
5
api/db/models.js
Normal file
5
api/db/models.js
Normal file
@ -0,0 +1,5 @@
|
||||
const mongoose = require('mongoose');
|
||||
const { createModels } = require('@librechat/data-schemas');
|
||||
const models = createModels(mongoose);
|
||||
|
||||
module.exports = { ...models };
|
||||
90
api/db/utils.js
Normal file
90
api/db/utils.js
Normal file
@ -0,0 +1,90 @@
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
|
||||
const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
|
||||
|
||||
/**
|
||||
* Batch update documents in chunks to avoid timeouts on weak instances
|
||||
* @param {mongoose.Collection} collection - MongoDB collection
|
||||
* @returns {Promise<number>} - Total modified count
|
||||
* @throws {Error} - Throws if database operations fail (e.g., network issues, connection loss, permission problems)
|
||||
*/
|
||||
async function batchResetMeiliFlags(collection) {
|
||||
const DEFAULT_BATCH_SIZE = 1000;
|
||||
|
||||
let BATCH_SIZE = parseEnvInt('MEILI_SYNC_BATCH_SIZE', DEFAULT_BATCH_SIZE);
|
||||
if (BATCH_SIZE === 0) {
|
||||
logger.warn(
|
||||
`[batchResetMeiliFlags] MEILI_SYNC_BATCH_SIZE cannot be 0. Using default: ${DEFAULT_BATCH_SIZE}`,
|
||||
);
|
||||
BATCH_SIZE = DEFAULT_BATCH_SIZE;
|
||||
}
|
||||
|
||||
const BATCH_DELAY_MS = parseEnvInt('MEILI_SYNC_DELAY_MS', 100);
|
||||
let totalModified = 0;
|
||||
let hasMore = true;
|
||||
|
||||
try {
|
||||
while (hasMore) {
|
||||
const docs = await collection
|
||||
.find({ expiredAt: null, _meiliIndex: true }, { projection: { _id: 1 } })
|
||||
.limit(BATCH_SIZE)
|
||||
.toArray();
|
||||
|
||||
if (docs.length === 0) {
|
||||
break;
|
||||
}
|
||||
|
||||
const ids = docs.map((doc) => doc._id);
|
||||
const result = await collection.updateMany(
|
||||
{ _id: { $in: ids } },
|
||||
{ $set: { _meiliIndex: false } },
|
||||
);
|
||||
|
||||
totalModified += result.modifiedCount;
|
||||
process.stdout.write(
|
||||
`\r Updating ${collection.collectionName}: ${totalModified} documents...`,
|
||||
);
|
||||
|
||||
if (docs.length < BATCH_SIZE) {
|
||||
hasMore = false;
|
||||
}
|
||||
|
||||
if (hasMore && BATCH_DELAY_MS > 0) {
|
||||
await sleep(BATCH_DELAY_MS);
|
||||
}
|
||||
}
|
||||
|
||||
return totalModified;
|
||||
} catch (error) {
|
||||
throw new Error(
|
||||
`Failed to batch reset Meili flags for collection '${collection.collectionName}' after processing ${totalModified} documents: ${error.message}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse and validate an environment variable as a positive integer
|
||||
* @param {string} varName - Environment variable name
|
||||
* @param {number} defaultValue - Default value to use if invalid or missing
|
||||
* @returns {number} - Parsed value or default
|
||||
*/
|
||||
function parseEnvInt(varName, defaultValue) {
|
||||
const value = process.env[varName];
|
||||
if (!value) {
|
||||
return defaultValue;
|
||||
}
|
||||
|
||||
const parsed = parseInt(value, 10);
|
||||
if (isNaN(parsed) || parsed < 0) {
|
||||
logger.warn(
|
||||
`[batchResetMeiliFlags] Invalid value for ${varName}="${value}". Expected a positive integer. Using default: ${defaultValue}`,
|
||||
);
|
||||
return defaultValue;
|
||||
}
|
||||
|
||||
return parsed;
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
batchResetMeiliFlags,
|
||||
};
|
||||
521
api/db/utils.spec.js
Normal file
521
api/db/utils.spec.js
Normal file
@ -0,0 +1,521 @@
|
||||
const mongoose = require('mongoose');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
const { batchResetMeiliFlags } = require('./utils');
|
||||
|
||||
describe('batchResetMeiliFlags', () => {
|
||||
let mongoServer;
|
||||
let testCollection;
|
||||
const ORIGINAL_BATCH_SIZE = process.env.MEILI_SYNC_BATCH_SIZE;
|
||||
const ORIGINAL_BATCH_DELAY = process.env.MEILI_SYNC_DELAY_MS;
|
||||
|
||||
beforeAll(async () => {
|
||||
mongoServer = await MongoMemoryServer.create();
|
||||
const mongoUri = mongoServer.getUri();
|
||||
await mongoose.connect(mongoUri);
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await mongoose.disconnect();
|
||||
await mongoServer.stop();
|
||||
|
||||
// Restore original env variables
|
||||
if (ORIGINAL_BATCH_SIZE !== undefined) {
|
||||
process.env.MEILI_SYNC_BATCH_SIZE = ORIGINAL_BATCH_SIZE;
|
||||
} else {
|
||||
delete process.env.MEILI_SYNC_BATCH_SIZE;
|
||||
}
|
||||
|
||||
if (ORIGINAL_BATCH_DELAY !== undefined) {
|
||||
process.env.MEILI_SYNC_DELAY_MS = ORIGINAL_BATCH_DELAY;
|
||||
} else {
|
||||
delete process.env.MEILI_SYNC_DELAY_MS;
|
||||
}
|
||||
});
|
||||
|
||||
beforeEach(async () => {
|
||||
// Create a fresh collection for each test
|
||||
testCollection = mongoose.connection.db.collection('test_meili_batch');
|
||||
await testCollection.deleteMany({});
|
||||
|
||||
// Reset env variables to defaults
|
||||
delete process.env.MEILI_SYNC_BATCH_SIZE;
|
||||
delete process.env.MEILI_SYNC_DELAY_MS;
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
if (testCollection) {
|
||||
await testCollection.deleteMany({});
|
||||
}
|
||||
});
|
||||
|
||||
describe('basic functionality', () => {
|
||||
it('should reset _meiliIndex flag for documents with expiredAt: null and _meiliIndex: true', async () => {
|
||||
// Insert test documents
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true, name: 'doc1' },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true, name: 'doc2' },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true, name: 'doc3' },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(3);
|
||||
|
||||
const updatedDocs = await testCollection.find({ _meiliIndex: false }).toArray();
|
||||
expect(updatedDocs).toHaveLength(3);
|
||||
|
||||
const notUpdatedDocs = await testCollection.find({ _meiliIndex: true }).toArray();
|
||||
expect(notUpdatedDocs).toHaveLength(0);
|
||||
});
|
||||
|
||||
it('should not modify documents with expiredAt set', async () => {
|
||||
const expiredDate = new Date();
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: expiredDate, _meiliIndex: true },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(1);
|
||||
|
||||
const expiredDoc = await testCollection.findOne({ expiredAt: expiredDate });
|
||||
expect(expiredDoc._meiliIndex).toBe(true);
|
||||
});
|
||||
|
||||
it('should not modify documents with _meiliIndex: false', async () => {
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: false },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(1);
|
||||
});
|
||||
|
||||
it('should return 0 when no documents match the criteria', async () => {
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: new Date(), _meiliIndex: true },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: false },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(0);
|
||||
});
|
||||
|
||||
it('should return 0 when collection is empty', async () => {
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('batch processing', () => {
|
||||
it('should process documents in batches according to MEILI_SYNC_BATCH_SIZE', async () => {
|
||||
process.env.MEILI_SYNC_BATCH_SIZE = '2';
|
||||
|
||||
const docs = [];
|
||||
for (let i = 0; i < 5; i++) {
|
||||
docs.push({
|
||||
_id: new mongoose.Types.ObjectId(),
|
||||
expiredAt: null,
|
||||
_meiliIndex: true,
|
||||
name: `doc${i}`,
|
||||
});
|
||||
}
|
||||
await testCollection.insertMany(docs);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(5);
|
||||
|
||||
const updatedDocs = await testCollection.find({ _meiliIndex: false }).toArray();
|
||||
expect(updatedDocs).toHaveLength(5);
|
||||
});
|
||||
|
||||
it('should handle large datasets with small batch sizes', async () => {
|
||||
process.env.MEILI_SYNC_BATCH_SIZE = '10';
|
||||
|
||||
const docs = [];
|
||||
for (let i = 0; i < 25; i++) {
|
||||
docs.push({
|
||||
_id: new mongoose.Types.ObjectId(),
|
||||
expiredAt: null,
|
||||
_meiliIndex: true,
|
||||
});
|
||||
}
|
||||
await testCollection.insertMany(docs);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(25);
|
||||
});
|
||||
|
||||
it('should use default batch size of 1000 when env variable is not set', async () => {
|
||||
// Create exactly 1000 documents to verify default batch behavior
|
||||
const docs = [];
|
||||
for (let i = 0; i < 1000; i++) {
|
||||
docs.push({
|
||||
_id: new mongoose.Types.ObjectId(),
|
||||
expiredAt: null,
|
||||
_meiliIndex: true,
|
||||
});
|
||||
}
|
||||
await testCollection.insertMany(docs);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(1000);
|
||||
});
|
||||
});
|
||||
|
||||
describe('return value', () => {
|
||||
it('should return correct modified count', async () => {
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
]);
|
||||
|
||||
await expect(batchResetMeiliFlags(testCollection)).resolves.toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe('batch delay', () => {
|
||||
it('should respect MEILI_SYNC_DELAY_MS between batches', async () => {
|
||||
process.env.MEILI_SYNC_BATCH_SIZE = '2';
|
||||
process.env.MEILI_SYNC_DELAY_MS = '50';
|
||||
|
||||
const docs = [];
|
||||
for (let i = 0; i < 5; i++) {
|
||||
docs.push({
|
||||
_id: new mongoose.Types.ObjectId(),
|
||||
expiredAt: null,
|
||||
_meiliIndex: true,
|
||||
});
|
||||
}
|
||||
await testCollection.insertMany(docs);
|
||||
|
||||
const startTime = Date.now();
|
||||
await batchResetMeiliFlags(testCollection);
|
||||
const endTime = Date.now();
|
||||
|
||||
// With 5 documents and batch size 2, we need 3 batches
|
||||
// That means 2 delays between batches (not after the last one)
|
||||
// So minimum time should be around 100ms (2 * 50ms)
|
||||
// Using a slightly lower threshold to account for timing variations
|
||||
const elapsed = endTime - startTime;
|
||||
expect(elapsed).toBeGreaterThanOrEqual(80);
|
||||
});
|
||||
|
||||
it('should not delay when MEILI_SYNC_DELAY_MS is 0', async () => {
|
||||
process.env.MEILI_SYNC_BATCH_SIZE = '2';
|
||||
process.env.MEILI_SYNC_DELAY_MS = '0';
|
||||
|
||||
const docs = [];
|
||||
for (let i = 0; i < 5; i++) {
|
||||
docs.push({
|
||||
_id: new mongoose.Types.ObjectId(),
|
||||
expiredAt: null,
|
||||
_meiliIndex: true,
|
||||
});
|
||||
}
|
||||
await testCollection.insertMany(docs);
|
||||
|
||||
const startTime = Date.now();
|
||||
await batchResetMeiliFlags(testCollection);
|
||||
const endTime = Date.now();
|
||||
|
||||
const elapsed = endTime - startTime;
|
||||
// Should complete without intentional delays, but database operations still take time
|
||||
// Just verify it completes and returns the correct count
|
||||
expect(elapsed).toBeLessThan(1000); // More reasonable upper bound
|
||||
|
||||
const result = await testCollection.countDocuments({ _meiliIndex: false });
|
||||
expect(result).toBe(5);
|
||||
});
|
||||
|
||||
it('should not delay after the last batch', async () => {
|
||||
process.env.MEILI_SYNC_BATCH_SIZE = '3';
|
||||
process.env.MEILI_SYNC_DELAY_MS = '100';
|
||||
|
||||
// Exactly 3 documents - should fit in one batch, no delay
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
// Verify all 3 documents were processed in a single batch
|
||||
expect(result).toBe(3);
|
||||
|
||||
const updatedDocs = await testCollection.countDocuments({ _meiliIndex: false });
|
||||
expect(updatedDocs).toBe(3);
|
||||
});
|
||||
});
|
||||
|
||||
describe('edge cases', () => {
|
||||
it('should handle documents without _meiliIndex field', async () => {
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
// Only one document has _meiliIndex: true
|
||||
expect(result).toBe(1);
|
||||
});
|
||||
|
||||
it('should handle mixed document states correctly', async () => {
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: false },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: new Date(), _meiliIndex: true },
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(2);
|
||||
|
||||
const flaggedDocs = await testCollection
|
||||
.find({ expiredAt: null, _meiliIndex: false })
|
||||
.toArray();
|
||||
expect(flaggedDocs).toHaveLength(3); // 2 were updated, 1 was already false
|
||||
});
|
||||
});
|
||||
|
||||
describe('error handling', () => {
|
||||
it('should throw error with context when find operation fails', async () => {
|
||||
const mockCollection = {
|
||||
collectionName: 'test_meili_batch',
|
||||
find: jest.fn().mockReturnValue({
|
||||
limit: jest.fn().mockReturnValue({
|
||||
toArray: jest.fn().mockRejectedValue(new Error('Network error')),
|
||||
}),
|
||||
}),
|
||||
};
|
||||
|
||||
await expect(batchResetMeiliFlags(mockCollection)).rejects.toThrow(
|
||||
"Failed to batch reset Meili flags for collection 'test_meili_batch' after processing 0 documents: Network error",
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw error with context when updateMany operation fails', async () => {
|
||||
const mockCollection = {
|
||||
collectionName: 'test_meili_batch',
|
||||
find: jest.fn().mockReturnValue({
|
||||
limit: jest.fn().mockReturnValue({
|
||||
toArray: jest
|
||||
.fn()
|
||||
.mockResolvedValue([
|
||||
{ _id: new mongoose.Types.ObjectId() },
|
||||
{ _id: new mongoose.Types.ObjectId() },
|
||||
]),
|
||||
}),
|
||||
}),
|
||||
updateMany: jest.fn().mockRejectedValue(new Error('Connection lost')),
|
||||
};
|
||||
|
||||
await expect(batchResetMeiliFlags(mockCollection)).rejects.toThrow(
|
||||
"Failed to batch reset Meili flags for collection 'test_meili_batch' after processing 0 documents: Connection lost",
|
||||
);
|
||||
});
|
||||
|
||||
it('should include documents processed count in error when failure occurs mid-batch', async () => {
|
||||
// Set batch size to 2 to force multiple batches
|
||||
process.env.MEILI_SYNC_BATCH_SIZE = '2';
|
||||
process.env.MEILI_SYNC_DELAY_MS = '0'; // No delay for faster test
|
||||
|
||||
let findCallCount = 0;
|
||||
let updateCallCount = 0;
|
||||
|
||||
const mockCollection = {
|
||||
collectionName: 'test_meili_batch',
|
||||
find: jest.fn().mockReturnValue({
|
||||
limit: jest.fn().mockReturnValue({
|
||||
toArray: jest.fn().mockImplementation(() => {
|
||||
findCallCount++;
|
||||
// Return 2 documents for first two calls (to keep loop going)
|
||||
// Return 2 documents for third call (to trigger third update which will fail)
|
||||
if (findCallCount <= 3) {
|
||||
return Promise.resolve([
|
||||
{ _id: new mongoose.Types.ObjectId() },
|
||||
{ _id: new mongoose.Types.ObjectId() },
|
||||
]);
|
||||
}
|
||||
// Should not reach here due to error
|
||||
return Promise.resolve([]);
|
||||
}),
|
||||
}),
|
||||
}),
|
||||
updateMany: jest.fn().mockImplementation(() => {
|
||||
updateCallCount++;
|
||||
if (updateCallCount === 1) {
|
||||
return Promise.resolve({ modifiedCount: 2 });
|
||||
} else if (updateCallCount === 2) {
|
||||
return Promise.resolve({ modifiedCount: 2 });
|
||||
} else {
|
||||
return Promise.reject(new Error('Database timeout'));
|
||||
}
|
||||
}),
|
||||
};
|
||||
|
||||
await expect(batchResetMeiliFlags(mockCollection)).rejects.toThrow(
|
||||
"Failed to batch reset Meili flags for collection 'test_meili_batch' after processing 4 documents: Database timeout",
|
||||
);
|
||||
});
|
||||
|
||||
it('should use collection.collectionName in error messages', async () => {
|
||||
const mockCollection = {
|
||||
collectionName: 'messages',
|
||||
find: jest.fn().mockReturnValue({
|
||||
limit: jest.fn().mockReturnValue({
|
||||
toArray: jest.fn().mockRejectedValue(new Error('Permission denied')),
|
||||
}),
|
||||
}),
|
||||
};
|
||||
|
||||
await expect(batchResetMeiliFlags(mockCollection)).rejects.toThrow(
|
||||
"Failed to batch reset Meili flags for collection 'messages' after processing 0 documents: Permission denied",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('environment variable validation', () => {
|
||||
let warnSpy;
|
||||
|
||||
beforeEach(() => {
|
||||
// Mock logger.warn to track warning calls
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
warnSpy = jest.spyOn(logger, 'warn').mockImplementation(() => {});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
if (warnSpy) {
|
||||
warnSpy.mockRestore();
|
||||
}
|
||||
});
|
||||
|
||||
it('should log warning and use default when MEILI_SYNC_BATCH_SIZE is not a number', async () => {
|
||||
process.env.MEILI_SYNC_BATCH_SIZE = 'abc';
|
||||
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(1);
|
||||
expect(warnSpy).toHaveBeenCalledWith(
|
||||
expect.stringContaining('Invalid value for MEILI_SYNC_BATCH_SIZE="abc"'),
|
||||
);
|
||||
expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('Using default: 1000'));
|
||||
});
|
||||
|
||||
it('should log warning and use default when MEILI_SYNC_DELAY_MS is not a number', async () => {
|
||||
process.env.MEILI_SYNC_DELAY_MS = 'xyz';
|
||||
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(1);
|
||||
expect(warnSpy).toHaveBeenCalledWith(
|
||||
expect.stringContaining('Invalid value for MEILI_SYNC_DELAY_MS="xyz"'),
|
||||
);
|
||||
expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('Using default: 100'));
|
||||
});
|
||||
|
||||
it('should log warning and use default when MEILI_SYNC_BATCH_SIZE is negative', async () => {
|
||||
process.env.MEILI_SYNC_BATCH_SIZE = '-50';
|
||||
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(1);
|
||||
expect(warnSpy).toHaveBeenCalledWith(
|
||||
expect.stringContaining('Invalid value for MEILI_SYNC_BATCH_SIZE="-50"'),
|
||||
);
|
||||
});
|
||||
|
||||
it('should log warning and use default when MEILI_SYNC_DELAY_MS is negative', async () => {
|
||||
process.env.MEILI_SYNC_DELAY_MS = '-100';
|
||||
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(1);
|
||||
expect(warnSpy).toHaveBeenCalledWith(
|
||||
expect.stringContaining('Invalid value for MEILI_SYNC_DELAY_MS="-100"'),
|
||||
);
|
||||
});
|
||||
|
||||
it('should accept valid positive integer values without warnings', async () => {
|
||||
process.env.MEILI_SYNC_BATCH_SIZE = '500';
|
||||
process.env.MEILI_SYNC_DELAY_MS = '50';
|
||||
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(1);
|
||||
expect(warnSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should log warning and use default when MEILI_SYNC_BATCH_SIZE is zero', async () => {
|
||||
process.env.MEILI_SYNC_BATCH_SIZE = '0';
|
||||
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(1);
|
||||
expect(warnSpy).toHaveBeenCalledWith(
|
||||
expect.stringContaining('MEILI_SYNC_BATCH_SIZE cannot be 0. Using default: 1000'),
|
||||
);
|
||||
});
|
||||
|
||||
it('should accept zero as a valid value for MEILI_SYNC_DELAY_MS without warnings', async () => {
|
||||
process.env.MEILI_SYNC_DELAY_MS = '0';
|
||||
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(1);
|
||||
expect(warnSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should not log warnings when environment variables are not set', async () => {
|
||||
delete process.env.MEILI_SYNC_BATCH_SIZE;
|
||||
delete process.env.MEILI_SYNC_DELAY_MS;
|
||||
|
||||
await testCollection.insertMany([
|
||||
{ _id: new mongoose.Types.ObjectId(), expiredAt: null, _meiliIndex: true },
|
||||
]);
|
||||
|
||||
const result = await batchResetMeiliFlags(testCollection);
|
||||
|
||||
expect(result).toBe(1);
|
||||
expect(warnSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
15
api/jest.config.js
Normal file
15
api/jest.config.js
Normal file
@ -0,0 +1,15 @@
|
||||
module.exports = {
|
||||
testEnvironment: 'node',
|
||||
clearMocks: true,
|
||||
roots: ['<rootDir>'],
|
||||
coverageDirectory: 'coverage',
|
||||
testTimeout: 30000, // 30 seconds timeout for all tests
|
||||
setupFiles: ['./test/jestSetup.js', './test/__mocks__/logger.js'],
|
||||
moduleNameMapper: {
|
||||
'~/(.*)': '<rootDir>/$1',
|
||||
'~/data/auth.json': '<rootDir>/__mocks__/auth.mock.json',
|
||||
'^openid-client/passport$': '<rootDir>/test/__mocks__/openid-client-passport.js', // Mock for the passport strategy part
|
||||
'^openid-client$': '<rootDir>/test/__mocks__/openid-client.js',
|
||||
},
|
||||
transformIgnorePatterns: ['/node_modules/(?!(openid-client|oauth4webapi|jose)/).*/'],
|
||||
};
|
||||
13
api/jsconfig.json
Normal file
13
api/jsconfig.json
Normal file
@ -0,0 +1,13 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"target": "ES6",
|
||||
"module": "CommonJS",
|
||||
// "checkJs": true, // Report errors in JavaScript files
|
||||
"baseUrl": "./",
|
||||
"paths": {
|
||||
"*": ["*", "node_modules/*"],
|
||||
"~/*": ["./*"]
|
||||
}
|
||||
},
|
||||
"exclude": ["node_modules"]
|
||||
}
|
||||
77
api/models/Action.js
Normal file
77
api/models/Action.js
Normal file
@ -0,0 +1,77 @@
|
||||
const { Action } = require('~/db/models');
|
||||
|
||||
/**
|
||||
* Update an action with new data without overwriting existing properties,
|
||||
* or create a new action if it doesn't exist.
|
||||
*
|
||||
* @param {Object} searchParams - The search parameters to find the action to update.
|
||||
* @param {string} searchParams.action_id - The ID of the action to update.
|
||||
* @param {string} searchParams.user - The user ID of the action's author.
|
||||
* @param {Object} updateData - An object containing the properties to update.
|
||||
* @returns {Promise<Action>} The updated or newly created action document as a plain object.
|
||||
*/
|
||||
const updateAction = async (searchParams, updateData) => {
|
||||
const options = { new: true, upsert: true };
|
||||
return await Action.findOneAndUpdate(searchParams, updateData, options).lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Retrieves all actions that match the given search parameters.
|
||||
*
|
||||
* @param {Object} searchParams - The search parameters to find matching actions.
|
||||
* @param {boolean} includeSensitive - Flag to include sensitive data in the metadata.
|
||||
* @returns {Promise<Array<Action>>} A promise that resolves to an array of action documents as plain objects.
|
||||
*/
|
||||
const getActions = async (searchParams, includeSensitive = false) => {
|
||||
const actions = await Action.find(searchParams).lean();
|
||||
|
||||
if (!includeSensitive) {
|
||||
for (let i = 0; i < actions.length; i++) {
|
||||
const metadata = actions[i].metadata;
|
||||
if (!metadata) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const sensitiveFields = ['api_key', 'oauth_client_id', 'oauth_client_secret'];
|
||||
for (let field of sensitiveFields) {
|
||||
if (metadata[field]) {
|
||||
delete metadata[field];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return actions;
|
||||
};
|
||||
|
||||
/**
|
||||
* Deletes an action by params.
|
||||
*
|
||||
* @param {Object} searchParams - The search parameters to find the action to delete.
|
||||
* @param {string} searchParams.action_id - The ID of the action to delete.
|
||||
* @param {string} searchParams.user - The user ID of the action's author.
|
||||
* @returns {Promise<Action>} A promise that resolves to the deleted action document as a plain object, or null if no document was found.
|
||||
*/
|
||||
const deleteAction = async (searchParams) => {
|
||||
return await Action.findOneAndDelete(searchParams).lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Deletes actions by params.
|
||||
*
|
||||
* @param {Object} searchParams - The search parameters to find the actions to delete.
|
||||
* @param {string} searchParams.action_id - The ID of the action(s) to delete.
|
||||
* @param {string} searchParams.user - The user ID of the action's author.
|
||||
* @returns {Promise<Number>} A promise that resolves to the number of deleted action documents.
|
||||
*/
|
||||
const deleteActions = async (searchParams) => {
|
||||
const result = await Action.deleteMany(searchParams);
|
||||
return result.deletedCount;
|
||||
};
|
||||
|
||||
module.exports = {
|
||||
getActions,
|
||||
updateAction,
|
||||
deleteAction,
|
||||
deleteActions,
|
||||
};
|
||||
925
api/models/Agent.js
Normal file
925
api/models/Agent.js
Normal file
@ -0,0 +1,925 @@
|
||||
const mongoose = require('mongoose');
|
||||
const crypto = require('node:crypto');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { getCustomEndpointConfig } = require('@librechat/api');
|
||||
const {
|
||||
Tools,
|
||||
SystemRoles,
|
||||
ResourceType,
|
||||
actionDelimiter,
|
||||
isAgentsEndpoint,
|
||||
isEphemeralAgentId,
|
||||
encodeEphemeralAgentId,
|
||||
} = require('librechat-data-provider');
|
||||
const { mcp_all, mcp_delimiter } = require('librechat-data-provider').Constants;
|
||||
const {
|
||||
removeAgentFromAllProjects,
|
||||
removeAgentIdsFromProject,
|
||||
addAgentIdsToProject,
|
||||
} = require('./Project');
|
||||
const { removeAllPermissions } = require('~/server/services/PermissionService');
|
||||
const { getMCPServerTools } = require('~/server/services/Config');
|
||||
const { Agent, AclEntry, User } = require('~/db/models');
|
||||
const { getActions } = require('./Action');
|
||||
|
||||
/**
|
||||
* Extracts unique MCP server names from tools array
|
||||
* Tools format: "toolName_mcp_serverName" or "sys__server__sys_mcp_serverName"
|
||||
* @param {string[]} tools - Array of tool identifiers
|
||||
* @returns {string[]} Array of unique MCP server names
|
||||
*/
|
||||
const extractMCPServerNames = (tools) => {
|
||||
if (!tools || !Array.isArray(tools)) {
|
||||
return [];
|
||||
}
|
||||
const serverNames = new Set();
|
||||
for (const tool of tools) {
|
||||
if (!tool || !tool.includes(mcp_delimiter)) {
|
||||
continue;
|
||||
}
|
||||
const parts = tool.split(mcp_delimiter);
|
||||
if (parts.length >= 2) {
|
||||
serverNames.add(parts[parts.length - 1]);
|
||||
}
|
||||
}
|
||||
return Array.from(serverNames);
|
||||
};
|
||||
|
||||
/**
|
||||
* Create an agent with the provided data.
|
||||
* @param {Object} agentData - The agent data to create.
|
||||
* @returns {Promise<Agent>} The created agent document as a plain object.
|
||||
* @throws {Error} If the agent creation fails.
|
||||
*/
|
||||
const createAgent = async (agentData) => {
|
||||
const { author: _author, ...versionData } = agentData;
|
||||
const timestamp = new Date();
|
||||
const initialAgentData = {
|
||||
...agentData,
|
||||
versions: [
|
||||
{
|
||||
...versionData,
|
||||
createdAt: timestamp,
|
||||
updatedAt: timestamp,
|
||||
},
|
||||
],
|
||||
category: agentData.category || 'general',
|
||||
mcpServerNames: extractMCPServerNames(agentData.tools),
|
||||
};
|
||||
|
||||
return (await Agent.create(initialAgentData)).toObject();
|
||||
};
|
||||
|
||||
/**
|
||||
* Get an agent document based on the provided ID.
|
||||
*
|
||||
* @param {Object} searchParameter - The search parameters to find the agent to update.
|
||||
* @param {string} searchParameter.id - The ID of the agent to update.
|
||||
* @param {string} searchParameter.author - The user ID of the agent's author.
|
||||
* @returns {Promise<Agent|null>} The agent document as a plain object, or null if not found.
|
||||
*/
|
||||
const getAgent = async (searchParameter) => await Agent.findOne(searchParameter).lean();
|
||||
|
||||
/**
|
||||
* Get multiple agent documents based on the provided search parameters.
|
||||
*
|
||||
* @param {Object} searchParameter - The search parameters to find agents.
|
||||
* @returns {Promise<Agent[]>} Array of agent documents as plain objects.
|
||||
*/
|
||||
const getAgents = async (searchParameter) => await Agent.find(searchParameter).lean();
|
||||
|
||||
/**
|
||||
* Load an agent based on the provided ID
|
||||
*
|
||||
* @param {Object} params
|
||||
* @param {ServerRequest} params.req
|
||||
* @param {string} params.spec
|
||||
* @param {string} params.agent_id
|
||||
* @param {string} params.endpoint
|
||||
* @param {import('@librechat/agents').ClientOptions} [params.model_parameters]
|
||||
* @returns {Promise<Agent|null>} The agent document as a plain object, or null if not found.
|
||||
*/
|
||||
const loadEphemeralAgent = async ({ req, spec, endpoint, model_parameters: _m }) => {
|
||||
const { model, ...model_parameters } = _m;
|
||||
const modelSpecs = req.config?.modelSpecs?.list;
|
||||
/** @type {TModelSpec | null} */
|
||||
let modelSpec = null;
|
||||
if (spec != null && spec !== '') {
|
||||
modelSpec = modelSpecs?.find((s) => s.name === spec) || null;
|
||||
}
|
||||
/** @type {TEphemeralAgent | null} */
|
||||
const ephemeralAgent = req.body.ephemeralAgent;
|
||||
const mcpServers = new Set(ephemeralAgent?.mcp);
|
||||
const userId = req.user?.id; // note: userId cannot be undefined at runtime
|
||||
if (modelSpec?.mcpServers) {
|
||||
for (const mcpServer of modelSpec.mcpServers) {
|
||||
mcpServers.add(mcpServer);
|
||||
}
|
||||
}
|
||||
/** @type {string[]} */
|
||||
const tools = [];
|
||||
if (ephemeralAgent?.execute_code === true || modelSpec?.executeCode === true) {
|
||||
tools.push(Tools.execute_code);
|
||||
}
|
||||
if (ephemeralAgent?.file_search === true || modelSpec?.fileSearch === true) {
|
||||
tools.push(Tools.file_search);
|
||||
}
|
||||
if (ephemeralAgent?.web_search === true || modelSpec?.webSearch === true) {
|
||||
tools.push(Tools.web_search);
|
||||
}
|
||||
|
||||
const addedServers = new Set();
|
||||
if (mcpServers.size > 0) {
|
||||
for (const mcpServer of mcpServers) {
|
||||
if (addedServers.has(mcpServer)) {
|
||||
continue;
|
||||
}
|
||||
const serverTools = await getMCPServerTools(userId, mcpServer);
|
||||
if (!serverTools) {
|
||||
tools.push(`${mcp_all}${mcp_delimiter}${mcpServer}`);
|
||||
addedServers.add(mcpServer);
|
||||
continue;
|
||||
}
|
||||
tools.push(...Object.keys(serverTools));
|
||||
addedServers.add(mcpServer);
|
||||
}
|
||||
}
|
||||
|
||||
const instructions = req.body.promptPrefix;
|
||||
|
||||
// Get endpoint config for modelDisplayLabel fallback
|
||||
const appConfig = req.config;
|
||||
let endpointConfig = appConfig?.endpoints?.[endpoint];
|
||||
if (!isAgentsEndpoint(endpoint) && !endpointConfig) {
|
||||
try {
|
||||
endpointConfig = getCustomEndpointConfig({ endpoint, appConfig });
|
||||
} catch (err) {
|
||||
logger.error('[loadEphemeralAgent] Error getting custom endpoint config', err);
|
||||
}
|
||||
}
|
||||
|
||||
// For ephemeral agents, use modelLabel if provided, then model spec's label,
|
||||
// then modelDisplayLabel from endpoint config, otherwise empty string to show model name
|
||||
const sender =
|
||||
model_parameters?.modelLabel ?? modelSpec?.label ?? endpointConfig?.modelDisplayLabel ?? '';
|
||||
|
||||
// Encode ephemeral agent ID with endpoint, model, and computed sender for display
|
||||
const ephemeralId = encodeEphemeralAgentId({ endpoint, model, sender });
|
||||
|
||||
const result = {
|
||||
id: ephemeralId,
|
||||
instructions,
|
||||
provider: endpoint,
|
||||
model_parameters,
|
||||
model,
|
||||
tools,
|
||||
};
|
||||
|
||||
if (ephemeralAgent?.artifacts != null && ephemeralAgent.artifacts) {
|
||||
result.artifacts = ephemeralAgent.artifacts;
|
||||
}
|
||||
return result;
|
||||
};
|
||||
|
||||
/**
|
||||
* Load an agent based on the provided ID
|
||||
*
|
||||
* @param {Object} params
|
||||
* @param {ServerRequest} params.req
|
||||
* @param {string} params.spec
|
||||
* @param {string} params.agent_id
|
||||
* @param {string} params.endpoint
|
||||
* @param {import('@librechat/agents').ClientOptions} [params.model_parameters]
|
||||
* @returns {Promise<Agent|null>} The agent document as a plain object, or null if not found.
|
||||
*/
|
||||
const loadAgent = async ({ req, spec, agent_id, endpoint, model_parameters }) => {
|
||||
if (!agent_id) {
|
||||
return null;
|
||||
}
|
||||
if (isEphemeralAgentId(agent_id)) {
|
||||
return await loadEphemeralAgent({ req, spec, endpoint, model_parameters });
|
||||
}
|
||||
const agent = await getAgent({
|
||||
id: agent_id,
|
||||
});
|
||||
|
||||
if (!agent) {
|
||||
return null;
|
||||
}
|
||||
|
||||
agent.version = agent.versions ? agent.versions.length : 0;
|
||||
return agent;
|
||||
};
|
||||
|
||||
/**
|
||||
* Check if a version already exists in the versions array, excluding timestamp and author fields
|
||||
* @param {Object} updateData - The update data to compare
|
||||
* @param {Object} currentData - The current agent data
|
||||
* @param {Array} versions - The existing versions array
|
||||
* @param {string} [actionsHash] - Hash of current action metadata
|
||||
* @returns {Object|null} - The matching version if found, null otherwise
|
||||
*/
|
||||
const isDuplicateVersion = (updateData, currentData, versions, actionsHash = null) => {
|
||||
if (!versions || versions.length === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const excludeFields = [
|
||||
'_id',
|
||||
'id',
|
||||
'createdAt',
|
||||
'updatedAt',
|
||||
'author',
|
||||
'updatedBy',
|
||||
'created_at',
|
||||
'updated_at',
|
||||
'__v',
|
||||
'versions',
|
||||
'actionsHash', // Exclude actionsHash from direct comparison
|
||||
];
|
||||
|
||||
const { $push: _$push, $pull: _$pull, $addToSet: _$addToSet, ...directUpdates } = updateData;
|
||||
|
||||
if (Object.keys(directUpdates).length === 0 && !actionsHash) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const wouldBeVersion = { ...currentData, ...directUpdates };
|
||||
const lastVersion = versions[versions.length - 1];
|
||||
|
||||
if (actionsHash && lastVersion.actionsHash !== actionsHash) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const allFields = new Set([...Object.keys(wouldBeVersion), ...Object.keys(lastVersion)]);
|
||||
|
||||
const importantFields = Array.from(allFields).filter((field) => !excludeFields.includes(field));
|
||||
|
||||
let isMatch = true;
|
||||
for (const field of importantFields) {
|
||||
const wouldBeValue = wouldBeVersion[field];
|
||||
const lastVersionValue = lastVersion[field];
|
||||
|
||||
// Skip if both are undefined/null
|
||||
if (!wouldBeValue && !lastVersionValue) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Handle arrays
|
||||
if (Array.isArray(wouldBeValue) || Array.isArray(lastVersionValue)) {
|
||||
// Normalize: treat undefined/null as empty array for comparison
|
||||
let wouldBeArr;
|
||||
if (Array.isArray(wouldBeValue)) {
|
||||
wouldBeArr = wouldBeValue;
|
||||
} else if (wouldBeValue == null) {
|
||||
wouldBeArr = [];
|
||||
} else {
|
||||
wouldBeArr = [wouldBeValue];
|
||||
}
|
||||
|
||||
let lastVersionArr;
|
||||
if (Array.isArray(lastVersionValue)) {
|
||||
lastVersionArr = lastVersionValue;
|
||||
} else if (lastVersionValue == null) {
|
||||
lastVersionArr = [];
|
||||
} else {
|
||||
lastVersionArr = [lastVersionValue];
|
||||
}
|
||||
|
||||
if (wouldBeArr.length !== lastVersionArr.length) {
|
||||
isMatch = false;
|
||||
break;
|
||||
}
|
||||
|
||||
// Special handling for projectIds (MongoDB ObjectIds)
|
||||
if (field === 'projectIds') {
|
||||
const wouldBeIds = wouldBeArr.map((id) => id.toString()).sort();
|
||||
const versionIds = lastVersionArr.map((id) => id.toString()).sort();
|
||||
|
||||
if (!wouldBeIds.every((id, i) => id === versionIds[i])) {
|
||||
isMatch = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
// Handle arrays of objects
|
||||
else if (
|
||||
wouldBeArr.length > 0 &&
|
||||
typeof wouldBeArr[0] === 'object' &&
|
||||
wouldBeArr[0] !== null
|
||||
) {
|
||||
const sortedWouldBe = [...wouldBeArr].map((item) => JSON.stringify(item)).sort();
|
||||
const sortedVersion = [...lastVersionArr].map((item) => JSON.stringify(item)).sort();
|
||||
|
||||
if (!sortedWouldBe.every((item, i) => item === sortedVersion[i])) {
|
||||
isMatch = false;
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
const sortedWouldBe = [...wouldBeArr].sort();
|
||||
const sortedVersion = [...lastVersionArr].sort();
|
||||
|
||||
if (!sortedWouldBe.every((item, i) => item === sortedVersion[i])) {
|
||||
isMatch = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
// Handle objects
|
||||
else if (typeof wouldBeValue === 'object' && wouldBeValue !== null) {
|
||||
const lastVersionObj =
|
||||
typeof lastVersionValue === 'object' && lastVersionValue !== null ? lastVersionValue : {};
|
||||
|
||||
// For empty objects, normalize the comparison
|
||||
const wouldBeKeys = Object.keys(wouldBeValue);
|
||||
const lastVersionKeys = Object.keys(lastVersionObj);
|
||||
|
||||
// If both are empty objects, they're equal
|
||||
if (wouldBeKeys.length === 0 && lastVersionKeys.length === 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Otherwise do a deep comparison
|
||||
if (JSON.stringify(wouldBeValue) !== JSON.stringify(lastVersionObj)) {
|
||||
isMatch = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
// Handle primitive values
|
||||
else {
|
||||
// For primitives, handle the case where one is undefined and the other is a default value
|
||||
if (wouldBeValue !== lastVersionValue) {
|
||||
// Special handling for boolean false vs undefined
|
||||
if (
|
||||
typeof wouldBeValue === 'boolean' &&
|
||||
wouldBeValue === false &&
|
||||
lastVersionValue === undefined
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
// Special handling for empty string vs undefined
|
||||
if (
|
||||
typeof wouldBeValue === 'string' &&
|
||||
wouldBeValue === '' &&
|
||||
lastVersionValue === undefined
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
isMatch = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return isMatch ? lastVersion : null;
|
||||
};
|
||||
|
||||
/**
|
||||
* Update an agent with new data without overwriting existing
|
||||
* properties, or create a new agent if it doesn't exist.
|
||||
* When an agent is updated, a copy of the current state will be saved to the versions array.
|
||||
*
|
||||
* @param {Object} searchParameter - The search parameters to find the agent to update.
|
||||
* @param {string} searchParameter.id - The ID of the agent to update.
|
||||
* @param {string} [searchParameter.author] - The user ID of the agent's author.
|
||||
* @param {Object} updateData - An object containing the properties to update.
|
||||
* @param {Object} [options] - Optional configuration object.
|
||||
* @param {string} [options.updatingUserId] - The ID of the user performing the update (used for tracking non-author updates).
|
||||
* @param {boolean} [options.forceVersion] - Force creation of a new version even if no fields changed.
|
||||
* @param {boolean} [options.skipVersioning] - Skip version creation entirely (useful for isolated operations like sharing).
|
||||
* @returns {Promise<Agent>} The updated or newly created agent document as a plain object.
|
||||
* @throws {Error} If the update would create a duplicate version
|
||||
*/
|
||||
const updateAgent = async (searchParameter, updateData, options = {}) => {
|
||||
const { updatingUserId = null, forceVersion = false, skipVersioning = false } = options;
|
||||
const mongoOptions = { new: true, upsert: false };
|
||||
|
||||
const currentAgent = await Agent.findOne(searchParameter);
|
||||
if (currentAgent) {
|
||||
const {
|
||||
__v,
|
||||
_id,
|
||||
id: __id,
|
||||
versions,
|
||||
author: _author,
|
||||
...versionData
|
||||
} = currentAgent.toObject();
|
||||
const { $push, $pull, $addToSet, ...directUpdates } = updateData;
|
||||
|
||||
// Sync mcpServerNames when tools are updated
|
||||
if (directUpdates.tools !== undefined) {
|
||||
const mcpServerNames = extractMCPServerNames(directUpdates.tools);
|
||||
directUpdates.mcpServerNames = mcpServerNames;
|
||||
updateData.mcpServerNames = mcpServerNames; // Also update the original updateData
|
||||
}
|
||||
|
||||
let actionsHash = null;
|
||||
|
||||
// Generate actions hash if agent has actions
|
||||
if (currentAgent.actions && currentAgent.actions.length > 0) {
|
||||
// Extract action IDs from the format "domain_action_id"
|
||||
const actionIds = currentAgent.actions
|
||||
.map((action) => {
|
||||
const parts = action.split(actionDelimiter);
|
||||
return parts[1]; // Get just the action ID part
|
||||
})
|
||||
.filter(Boolean);
|
||||
|
||||
if (actionIds.length > 0) {
|
||||
try {
|
||||
const actions = await getActions(
|
||||
{
|
||||
action_id: { $in: actionIds },
|
||||
},
|
||||
true,
|
||||
); // Include sensitive data for hash
|
||||
|
||||
actionsHash = await generateActionMetadataHash(currentAgent.actions, actions);
|
||||
} catch (error) {
|
||||
logger.error('Error fetching actions for hash generation:', error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const shouldCreateVersion =
|
||||
!skipVersioning &&
|
||||
(forceVersion || Object.keys(directUpdates).length > 0 || $push || $pull || $addToSet);
|
||||
|
||||
if (shouldCreateVersion) {
|
||||
const duplicateVersion = isDuplicateVersion(updateData, versionData, versions, actionsHash);
|
||||
if (duplicateVersion && !forceVersion) {
|
||||
// No changes detected, return the current agent without creating a new version
|
||||
const agentObj = currentAgent.toObject();
|
||||
agentObj.version = versions.length;
|
||||
return agentObj;
|
||||
}
|
||||
}
|
||||
|
||||
const versionEntry = {
|
||||
...versionData,
|
||||
...directUpdates,
|
||||
updatedAt: new Date(),
|
||||
};
|
||||
|
||||
// Include actions hash in version if available
|
||||
if (actionsHash) {
|
||||
versionEntry.actionsHash = actionsHash;
|
||||
}
|
||||
|
||||
// Always store updatedBy field to track who made the change
|
||||
if (updatingUserId) {
|
||||
versionEntry.updatedBy = new mongoose.Types.ObjectId(updatingUserId);
|
||||
}
|
||||
|
||||
if (shouldCreateVersion) {
|
||||
updateData.$push = {
|
||||
...($push || {}),
|
||||
versions: versionEntry,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
return Agent.findOneAndUpdate(searchParameter, updateData, mongoOptions).lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Modifies an agent with the resource file id.
|
||||
* @param {object} params
|
||||
* @param {ServerRequest} params.req
|
||||
* @param {string} params.agent_id
|
||||
* @param {string} params.tool_resource
|
||||
* @param {string} params.file_id
|
||||
* @returns {Promise<Agent>} The updated agent.
|
||||
*/
|
||||
const addAgentResourceFile = async ({ req, agent_id, tool_resource, file_id }) => {
|
||||
const searchParameter = { id: agent_id };
|
||||
let agent = await getAgent(searchParameter);
|
||||
if (!agent) {
|
||||
throw new Error('Agent not found for adding resource file');
|
||||
}
|
||||
const fileIdsPath = `tool_resources.${tool_resource}.file_ids`;
|
||||
await Agent.updateOne(
|
||||
{
|
||||
id: agent_id,
|
||||
[`${fileIdsPath}`]: { $exists: false },
|
||||
},
|
||||
{
|
||||
$set: {
|
||||
[`${fileIdsPath}`]: [],
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
const updateData = {
|
||||
$addToSet: {
|
||||
tools: tool_resource,
|
||||
[fileIdsPath]: file_id,
|
||||
},
|
||||
};
|
||||
|
||||
const updatedAgent = await updateAgent(searchParameter, updateData, {
|
||||
updatingUserId: req?.user?.id,
|
||||
});
|
||||
if (updatedAgent) {
|
||||
return updatedAgent;
|
||||
} else {
|
||||
throw new Error('Agent not found for adding resource file');
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Removes multiple resource files from an agent using atomic operations.
|
||||
* @param {object} params
|
||||
* @param {string} params.agent_id
|
||||
* @param {Array<{tool_resource: string, file_id: string}>} params.files
|
||||
* @returns {Promise<Agent>} The updated agent.
|
||||
* @throws {Error} If the agent is not found or update fails.
|
||||
*/
|
||||
const removeAgentResourceFiles = async ({ agent_id, files }) => {
|
||||
const searchParameter = { id: agent_id };
|
||||
|
||||
// Group files to remove by resource
|
||||
const filesByResource = files.reduce((acc, { tool_resource, file_id }) => {
|
||||
if (!acc[tool_resource]) {
|
||||
acc[tool_resource] = [];
|
||||
}
|
||||
acc[tool_resource].push(file_id);
|
||||
return acc;
|
||||
}, {});
|
||||
|
||||
// Step 1: Atomically remove file IDs using $pull
|
||||
const pullOps = {};
|
||||
const resourcesToCheck = new Set();
|
||||
for (const [resource, fileIds] of Object.entries(filesByResource)) {
|
||||
const fileIdsPath = `tool_resources.${resource}.file_ids`;
|
||||
pullOps[fileIdsPath] = { $in: fileIds };
|
||||
resourcesToCheck.add(resource);
|
||||
}
|
||||
|
||||
const updatePullData = { $pull: pullOps };
|
||||
const agentAfterPull = await Agent.findOneAndUpdate(searchParameter, updatePullData, {
|
||||
new: true,
|
||||
}).lean();
|
||||
|
||||
if (!agentAfterPull) {
|
||||
// Agent might have been deleted concurrently, or never existed.
|
||||
// Check if it existed before trying to throw.
|
||||
const agentExists = await getAgent(searchParameter);
|
||||
if (!agentExists) {
|
||||
throw new Error('Agent not found for removing resource files');
|
||||
}
|
||||
// If it existed but findOneAndUpdate returned null, something else went wrong.
|
||||
throw new Error('Failed to update agent during file removal (pull step)');
|
||||
}
|
||||
|
||||
// Return the agent state directly after the $pull operation.
|
||||
// Skipping the $unset step for now to simplify and test core $pull atomicity.
|
||||
// Empty arrays might remain, but the removal itself should be correct.
|
||||
return agentAfterPull;
|
||||
};
|
||||
|
||||
/**
|
||||
* Deletes an agent based on the provided ID.
|
||||
*
|
||||
* @param {Object} searchParameter - The search parameters to find the agent to delete.
|
||||
* @param {string} searchParameter.id - The ID of the agent to delete.
|
||||
* @param {string} [searchParameter.author] - The user ID of the agent's author.
|
||||
* @returns {Promise<void>} Resolves when the agent has been successfully deleted.
|
||||
*/
|
||||
const deleteAgent = async (searchParameter) => {
|
||||
const agent = await Agent.findOneAndDelete(searchParameter);
|
||||
if (agent) {
|
||||
await removeAgentFromAllProjects(agent.id);
|
||||
await removeAllPermissions({
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceId: agent._id,
|
||||
});
|
||||
try {
|
||||
await Agent.updateMany({ 'edges.to': agent.id }, { $pull: { edges: { to: agent.id } } });
|
||||
} catch (error) {
|
||||
logger.error('[deleteAgent] Error removing agent from handoff edges', error);
|
||||
}
|
||||
try {
|
||||
await User.updateMany(
|
||||
{ 'favorites.agentId': agent.id },
|
||||
{ $pull: { favorites: { agentId: agent.id } } },
|
||||
);
|
||||
} catch (error) {
|
||||
logger.error('[deleteAgent] Error removing agent from user favorites', error);
|
||||
}
|
||||
}
|
||||
return agent;
|
||||
};
|
||||
|
||||
/**
|
||||
* Deletes all agents created by a specific user.
|
||||
* @param {string} userId - The ID of the user whose agents should be deleted.
|
||||
* @returns {Promise<void>} A promise that resolves when all user agents have been deleted.
|
||||
*/
|
||||
const deleteUserAgents = async (userId) => {
|
||||
try {
|
||||
const userAgents = await getAgents({ author: userId });
|
||||
|
||||
if (userAgents.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const agentIds = userAgents.map((agent) => agent.id);
|
||||
const agentObjectIds = userAgents.map((agent) => agent._id);
|
||||
|
||||
for (const agentId of agentIds) {
|
||||
await removeAgentFromAllProjects(agentId);
|
||||
}
|
||||
|
||||
await AclEntry.deleteMany({
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceId: { $in: agentObjectIds },
|
||||
});
|
||||
|
||||
try {
|
||||
await User.updateMany(
|
||||
{ 'favorites.agentId': { $in: agentIds } },
|
||||
{ $pull: { favorites: { agentId: { $in: agentIds } } } },
|
||||
);
|
||||
} catch (error) {
|
||||
logger.error('[deleteUserAgents] Error removing agents from user favorites', error);
|
||||
}
|
||||
|
||||
await Agent.deleteMany({ author: userId });
|
||||
} catch (error) {
|
||||
logger.error('[deleteUserAgents] General error:', error);
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Get agents by accessible IDs with optional cursor-based pagination.
|
||||
* @param {Object} params - The parameters for getting accessible agents.
|
||||
* @param {Array} [params.accessibleIds] - Array of agent ObjectIds the user has ACL access to.
|
||||
* @param {Object} [params.otherParams] - Additional query parameters (including author filter).
|
||||
* @param {number} [params.limit] - Number of agents to return (max 100). If not provided, returns all agents.
|
||||
* @param {string} [params.after] - Cursor for pagination - get agents after this cursor. // base64 encoded JSON string with updatedAt and _id.
|
||||
* @returns {Promise<Object>} A promise that resolves to an object containing the agents data and pagination info.
|
||||
*/
|
||||
const getListAgentsByAccess = async ({
|
||||
accessibleIds = [],
|
||||
otherParams = {},
|
||||
limit = null,
|
||||
after = null,
|
||||
}) => {
|
||||
const isPaginated = limit !== null && limit !== undefined;
|
||||
const normalizedLimit = isPaginated ? Math.min(Math.max(1, parseInt(limit) || 20), 100) : null;
|
||||
|
||||
// Build base query combining ACL accessible agents with other filters
|
||||
const baseQuery = { ...otherParams, _id: { $in: accessibleIds } };
|
||||
|
||||
// Add cursor condition
|
||||
if (after) {
|
||||
try {
|
||||
const cursor = JSON.parse(Buffer.from(after, 'base64').toString('utf8'));
|
||||
const { updatedAt, _id } = cursor;
|
||||
|
||||
const cursorCondition = {
|
||||
$or: [
|
||||
{ updatedAt: { $lt: new Date(updatedAt) } },
|
||||
{ updatedAt: new Date(updatedAt), _id: { $gt: new mongoose.Types.ObjectId(_id) } },
|
||||
],
|
||||
};
|
||||
|
||||
// Merge cursor condition with base query
|
||||
if (Object.keys(baseQuery).length > 0) {
|
||||
baseQuery.$and = [{ ...baseQuery }, cursorCondition];
|
||||
// Remove the original conditions from baseQuery to avoid duplication
|
||||
Object.keys(baseQuery).forEach((key) => {
|
||||
if (key !== '$and') delete baseQuery[key];
|
||||
});
|
||||
} else {
|
||||
Object.assign(baseQuery, cursorCondition);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.warn('Invalid cursor:', error.message);
|
||||
}
|
||||
}
|
||||
|
||||
let query = Agent.find(baseQuery, {
|
||||
id: 1,
|
||||
_id: 1,
|
||||
name: 1,
|
||||
avatar: 1,
|
||||
author: 1,
|
||||
projectIds: 1,
|
||||
description: 1,
|
||||
updatedAt: 1,
|
||||
category: 1,
|
||||
support_contact: 1,
|
||||
is_promoted: 1,
|
||||
}).sort({ updatedAt: -1, _id: 1 });
|
||||
|
||||
// Only apply limit if pagination is requested
|
||||
if (isPaginated) {
|
||||
query = query.limit(normalizedLimit + 1);
|
||||
}
|
||||
|
||||
const agents = await query.lean();
|
||||
|
||||
const hasMore = isPaginated ? agents.length > normalizedLimit : false;
|
||||
const data = (isPaginated ? agents.slice(0, normalizedLimit) : agents).map((agent) => {
|
||||
if (agent.author) {
|
||||
agent.author = agent.author.toString();
|
||||
}
|
||||
return agent;
|
||||
});
|
||||
|
||||
// Generate next cursor only if paginated
|
||||
let nextCursor = null;
|
||||
if (isPaginated && hasMore && data.length > 0) {
|
||||
const lastAgent = agents[normalizedLimit - 1];
|
||||
nextCursor = Buffer.from(
|
||||
JSON.stringify({
|
||||
updatedAt: lastAgent.updatedAt.toISOString(),
|
||||
_id: lastAgent._id.toString(),
|
||||
}),
|
||||
).toString('base64');
|
||||
}
|
||||
|
||||
return {
|
||||
object: 'list',
|
||||
data,
|
||||
first_id: data.length > 0 ? data[0].id : null,
|
||||
last_id: data.length > 0 ? data[data.length - 1].id : null,
|
||||
has_more: hasMore,
|
||||
after: nextCursor,
|
||||
};
|
||||
};
|
||||
|
||||
/**
|
||||
* Updates the projects associated with an agent, adding and removing project IDs as specified.
|
||||
* This function also updates the corresponding projects to include or exclude the agent ID.
|
||||
*
|
||||
* @param {Object} params - Parameters for updating the agent's projects.
|
||||
* @param {IUser} params.user - Parameters for updating the agent's projects.
|
||||
* @param {string} params.agentId - The ID of the agent to update.
|
||||
* @param {string[]} [params.projectIds] - Array of project IDs to add to the agent.
|
||||
* @param {string[]} [params.removeProjectIds] - Array of project IDs to remove from the agent.
|
||||
* @returns {Promise<MongoAgent>} The updated agent document.
|
||||
* @throws {Error} If there's an error updating the agent or projects.
|
||||
*/
|
||||
const updateAgentProjects = async ({ user, agentId, projectIds, removeProjectIds }) => {
|
||||
const updateOps = {};
|
||||
|
||||
if (removeProjectIds && removeProjectIds.length > 0) {
|
||||
for (const projectId of removeProjectIds) {
|
||||
await removeAgentIdsFromProject(projectId, [agentId]);
|
||||
}
|
||||
updateOps.$pull = { projectIds: { $in: removeProjectIds } };
|
||||
}
|
||||
|
||||
if (projectIds && projectIds.length > 0) {
|
||||
for (const projectId of projectIds) {
|
||||
await addAgentIdsToProject(projectId, [agentId]);
|
||||
}
|
||||
updateOps.$addToSet = { projectIds: { $each: projectIds } };
|
||||
}
|
||||
|
||||
if (Object.keys(updateOps).length === 0) {
|
||||
return await getAgent({ id: agentId });
|
||||
}
|
||||
|
||||
const updateQuery = { id: agentId, author: user.id };
|
||||
if (user.role === SystemRoles.ADMIN) {
|
||||
delete updateQuery.author;
|
||||
}
|
||||
|
||||
const updatedAgent = await updateAgent(updateQuery, updateOps, {
|
||||
updatingUserId: user.id,
|
||||
skipVersioning: true,
|
||||
});
|
||||
if (updatedAgent) {
|
||||
return updatedAgent;
|
||||
}
|
||||
if (updateOps.$addToSet) {
|
||||
for (const projectId of projectIds) {
|
||||
await removeAgentIdsFromProject(projectId, [agentId]);
|
||||
}
|
||||
} else if (updateOps.$pull) {
|
||||
for (const projectId of removeProjectIds) {
|
||||
await addAgentIdsToProject(projectId, [agentId]);
|
||||
}
|
||||
}
|
||||
|
||||
return await getAgent({ id: agentId });
|
||||
};
|
||||
|
||||
/**
|
||||
* Reverts an agent to a specific version in its version history.
|
||||
* @param {Object} searchParameter - The search parameters to find the agent to revert.
|
||||
* @param {string} searchParameter.id - The ID of the agent to revert.
|
||||
* @param {string} [searchParameter.author] - The user ID of the agent's author.
|
||||
* @param {number} versionIndex - The index of the version to revert to in the versions array.
|
||||
* @returns {Promise<MongoAgent>} The updated agent document after reverting.
|
||||
* @throws {Error} If the agent is not found or the specified version does not exist.
|
||||
*/
|
||||
const revertAgentVersion = async (searchParameter, versionIndex) => {
|
||||
const agent = await Agent.findOne(searchParameter);
|
||||
if (!agent) {
|
||||
throw new Error('Agent not found');
|
||||
}
|
||||
|
||||
if (!agent.versions || !agent.versions[versionIndex]) {
|
||||
throw new Error(`Version ${versionIndex} not found`);
|
||||
}
|
||||
|
||||
const revertToVersion = agent.versions[versionIndex];
|
||||
|
||||
const updateData = {
|
||||
...revertToVersion,
|
||||
};
|
||||
|
||||
delete updateData._id;
|
||||
delete updateData.id;
|
||||
delete updateData.versions;
|
||||
delete updateData.author;
|
||||
delete updateData.updatedBy;
|
||||
|
||||
return Agent.findOneAndUpdate(searchParameter, updateData, { new: true }).lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Generates a hash of action metadata for version comparison
|
||||
* @param {string[]} actionIds - Array of action IDs in format "domain_action_id"
|
||||
* @param {Action[]} actions - Array of action documents
|
||||
* @returns {Promise<string>} - SHA256 hash of the action metadata
|
||||
*/
|
||||
const generateActionMetadataHash = async (actionIds, actions) => {
|
||||
if (!actionIds || actionIds.length === 0) {
|
||||
return '';
|
||||
}
|
||||
|
||||
// Create a map of action_id to metadata for quick lookup
|
||||
const actionMap = new Map();
|
||||
actions.forEach((action) => {
|
||||
actionMap.set(action.action_id, action.metadata);
|
||||
});
|
||||
|
||||
// Sort action IDs for consistent hashing
|
||||
const sortedActionIds = [...actionIds].sort();
|
||||
|
||||
// Build a deterministic string representation of all action metadata
|
||||
const metadataString = sortedActionIds
|
||||
.map((actionFullId) => {
|
||||
// Extract just the action_id part (after the delimiter)
|
||||
const parts = actionFullId.split(actionDelimiter);
|
||||
const actionId = parts[1];
|
||||
|
||||
const metadata = actionMap.get(actionId);
|
||||
if (!metadata) {
|
||||
return `${actionId}:null`;
|
||||
}
|
||||
|
||||
// Sort metadata keys for deterministic output
|
||||
const sortedKeys = Object.keys(metadata).sort();
|
||||
const metadataStr = sortedKeys
|
||||
.map((key) => `${key}:${JSON.stringify(metadata[key])}`)
|
||||
.join(',');
|
||||
return `${actionId}:{${metadataStr}}`;
|
||||
})
|
||||
.join(';');
|
||||
|
||||
// Use Web Crypto API to generate hash
|
||||
const encoder = new TextEncoder();
|
||||
const data = encoder.encode(metadataString);
|
||||
const hashBuffer = await crypto.webcrypto.subtle.digest('SHA-256', data);
|
||||
const hashArray = Array.from(new Uint8Array(hashBuffer));
|
||||
const hashHex = hashArray.map((b) => b.toString(16).padStart(2, '0')).join('');
|
||||
|
||||
return hashHex;
|
||||
};
|
||||
/**
|
||||
* Counts the number of promoted agents.
|
||||
* @returns {Promise<number>} - The count of promoted agents
|
||||
*/
|
||||
const countPromotedAgents = async () => {
|
||||
const count = await Agent.countDocuments({ is_promoted: true });
|
||||
return count;
|
||||
};
|
||||
|
||||
/**
|
||||
* Load a default agent based on the endpoint
|
||||
* @param {string} endpoint
|
||||
* @returns {Agent | null}
|
||||
*/
|
||||
|
||||
module.exports = {
|
||||
getAgent,
|
||||
getAgents,
|
||||
loadAgent,
|
||||
createAgent,
|
||||
updateAgent,
|
||||
deleteAgent,
|
||||
deleteUserAgents,
|
||||
revertAgentVersion,
|
||||
updateAgentProjects,
|
||||
countPromotedAgents,
|
||||
addAgentResourceFile,
|
||||
getListAgentsByAccess,
|
||||
removeAgentResourceFiles,
|
||||
generateActionMetadataHash,
|
||||
};
|
||||
3829
api/models/Agent.spec.js
Normal file
3829
api/models/Agent.spec.js
Normal file
File diff suppressed because it is too large
Load Diff
62
api/models/Assistant.js
Normal file
62
api/models/Assistant.js
Normal file
@ -0,0 +1,62 @@
|
||||
const { Assistant } = require('~/db/models');
|
||||
|
||||
/**
|
||||
* Update an assistant with new data without overwriting existing properties,
|
||||
* or create a new assistant if it doesn't exist.
|
||||
*
|
||||
* @param {Object} searchParams - The search parameters to find the assistant to update.
|
||||
* @param {string} searchParams.assistant_id - The ID of the assistant to update.
|
||||
* @param {string} searchParams.user - The user ID of the assistant's author.
|
||||
* @param {Object} updateData - An object containing the properties to update.
|
||||
* @returns {Promise<AssistantDocument>} The updated or newly created assistant document as a plain object.
|
||||
*/
|
||||
const updateAssistantDoc = async (searchParams, updateData) => {
|
||||
const options = { new: true, upsert: true };
|
||||
return await Assistant.findOneAndUpdate(searchParams, updateData, options).lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Retrieves an assistant document based on the provided ID.
|
||||
*
|
||||
* @param {Object} searchParams - The search parameters to find the assistant to update.
|
||||
* @param {string} searchParams.assistant_id - The ID of the assistant to update.
|
||||
* @param {string} searchParams.user - The user ID of the assistant's author.
|
||||
* @returns {Promise<AssistantDocument|null>} The assistant document as a plain object, or null if not found.
|
||||
*/
|
||||
const getAssistant = async (searchParams) => await Assistant.findOne(searchParams).lean();
|
||||
|
||||
/**
|
||||
* Retrieves all assistants that match the given search parameters.
|
||||
*
|
||||
* @param {Object} searchParams - The search parameters to find matching assistants.
|
||||
* @param {Object} [select] - Optional. Specifies which document fields to include or exclude.
|
||||
* @returns {Promise<Array<AssistantDocument>>} A promise that resolves to an array of assistant documents as plain objects.
|
||||
*/
|
||||
const getAssistants = async (searchParams, select = null) => {
|
||||
let query = Assistant.find(searchParams);
|
||||
|
||||
if (select) {
|
||||
query = query.select(select);
|
||||
}
|
||||
|
||||
return await query.lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Deletes an assistant based on the provided ID.
|
||||
*
|
||||
* @param {Object} searchParams - The search parameters to find the assistant to delete.
|
||||
* @param {string} searchParams.assistant_id - The ID of the assistant to delete.
|
||||
* @param {string} searchParams.user - The user ID of the assistant's author.
|
||||
* @returns {Promise<void>} Resolves when the assistant has been successfully deleted.
|
||||
*/
|
||||
const deleteAssistant = async (searchParams) => {
|
||||
return await Assistant.findOneAndDelete(searchParams);
|
||||
};
|
||||
|
||||
module.exports = {
|
||||
updateAssistantDoc,
|
||||
deleteAssistant,
|
||||
getAssistants,
|
||||
getAssistant,
|
||||
};
|
||||
28
api/models/Banner.js
Normal file
28
api/models/Banner.js
Normal file
@ -0,0 +1,28 @@
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { Banner } = require('~/db/models');
|
||||
|
||||
/**
|
||||
* Retrieves the current active banner.
|
||||
* @returns {Promise<Object|null>} The active banner object or null if no active banner is found.
|
||||
*/
|
||||
const getBanner = async (user) => {
|
||||
try {
|
||||
const now = new Date();
|
||||
const banner = await Banner.findOne({
|
||||
displayFrom: { $lte: now },
|
||||
$or: [{ displayTo: { $gte: now } }, { displayTo: null }],
|
||||
type: 'banner',
|
||||
}).lean();
|
||||
|
||||
if (!banner || banner.isPublic || user) {
|
||||
return banner;
|
||||
}
|
||||
|
||||
return null;
|
||||
} catch (error) {
|
||||
logger.error('[getBanners] Error getting banners', error);
|
||||
throw new Error('Error getting banners');
|
||||
}
|
||||
};
|
||||
|
||||
module.exports = { getBanner };
|
||||
57
api/models/Categories.js
Normal file
57
api/models/Categories.js
Normal file
@ -0,0 +1,57 @@
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
|
||||
const options = [
|
||||
{
|
||||
label: 'com_ui_idea',
|
||||
value: 'idea',
|
||||
},
|
||||
{
|
||||
label: 'com_ui_travel',
|
||||
value: 'travel',
|
||||
},
|
||||
{
|
||||
label: 'com_ui_teach_or_explain',
|
||||
value: 'teach_or_explain',
|
||||
},
|
||||
{
|
||||
label: 'com_ui_write',
|
||||
value: 'write',
|
||||
},
|
||||
{
|
||||
label: 'com_ui_shop',
|
||||
value: 'shop',
|
||||
},
|
||||
{
|
||||
label: 'com_ui_code',
|
||||
value: 'code',
|
||||
},
|
||||
{
|
||||
label: 'com_ui_misc',
|
||||
value: 'misc',
|
||||
},
|
||||
{
|
||||
label: 'com_ui_roleplay',
|
||||
value: 'roleplay',
|
||||
},
|
||||
{
|
||||
label: 'com_ui_finance',
|
||||
value: 'finance',
|
||||
},
|
||||
];
|
||||
|
||||
module.exports = {
|
||||
/**
|
||||
* Retrieves the categories asynchronously.
|
||||
* @returns {Promise<TGetCategoriesResponse>} An array of category objects.
|
||||
* @throws {Error} If there is an error retrieving the categories.
|
||||
*/
|
||||
getCategories: async () => {
|
||||
try {
|
||||
// const categories = await Categories.find();
|
||||
return options;
|
||||
} catch (error) {
|
||||
logger.error('Error getting categories', error);
|
||||
return [];
|
||||
}
|
||||
},
|
||||
};
|
||||
367
api/models/Conversation.js
Normal file
367
api/models/Conversation.js
Normal file
@ -0,0 +1,367 @@
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { createTempChatExpirationDate } = require('@librechat/api');
|
||||
const { getMessages, deleteMessages } = require('./Message');
|
||||
const { Conversation } = require('~/db/models');
|
||||
|
||||
/**
|
||||
* Searches for a conversation by conversationId and returns a lean document with only conversationId and user.
|
||||
* @param {string} conversationId - The conversation's ID.
|
||||
* @returns {Promise<{conversationId: string, user: string} | null>} The conversation object with selected fields or null if not found.
|
||||
*/
|
||||
const searchConversation = async (conversationId) => {
|
||||
try {
|
||||
return await Conversation.findOne({ conversationId }, 'conversationId user').lean();
|
||||
} catch (error) {
|
||||
logger.error('[searchConversation] Error searching conversation', error);
|
||||
throw new Error('Error searching conversation');
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Retrieves a single conversation for a given user and conversation ID.
|
||||
* @param {string} user - The user's ID.
|
||||
* @param {string} conversationId - The conversation's ID.
|
||||
* @returns {Promise<TConversation>} The conversation object.
|
||||
*/
|
||||
const getConvo = async (user, conversationId) => {
|
||||
try {
|
||||
return await Conversation.findOne({ user, conversationId }).lean();
|
||||
} catch (error) {
|
||||
logger.error('[getConvo] Error getting single conversation', error);
|
||||
throw new Error('Error getting single conversation');
|
||||
}
|
||||
};
|
||||
|
||||
const deleteNullOrEmptyConversations = async () => {
|
||||
try {
|
||||
const filter = {
|
||||
$or: [
|
||||
{ conversationId: null },
|
||||
{ conversationId: '' },
|
||||
{ conversationId: { $exists: false } },
|
||||
],
|
||||
};
|
||||
|
||||
const result = await Conversation.deleteMany(filter);
|
||||
|
||||
// Delete associated messages
|
||||
const messageDeleteResult = await deleteMessages(filter);
|
||||
|
||||
logger.info(
|
||||
`[deleteNullOrEmptyConversations] Deleted ${result.deletedCount} conversations and ${messageDeleteResult.deletedCount} messages`,
|
||||
);
|
||||
|
||||
return {
|
||||
conversations: result,
|
||||
messages: messageDeleteResult,
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error('[deleteNullOrEmptyConversations] Error deleting conversations', error);
|
||||
throw new Error('Error deleting conversations with null or empty conversationId');
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Searches for a conversation by conversationId and returns associated file ids.
|
||||
* @param {string} conversationId - The conversation's ID.
|
||||
* @returns {Promise<string[] | null>}
|
||||
*/
|
||||
const getConvoFiles = async (conversationId) => {
|
||||
try {
|
||||
return (await Conversation.findOne({ conversationId }, 'files').lean())?.files ?? [];
|
||||
} catch (error) {
|
||||
logger.error('[getConvoFiles] Error getting conversation files', error);
|
||||
throw new Error('Error getting conversation files');
|
||||
}
|
||||
};
|
||||
|
||||
module.exports = {
|
||||
getConvoFiles,
|
||||
searchConversation,
|
||||
deleteNullOrEmptyConversations,
|
||||
/**
|
||||
* Saves a conversation to the database.
|
||||
* @param {Object} req - The request object.
|
||||
* @param {string} conversationId - The conversation's ID.
|
||||
* @param {Object} metadata - Additional metadata to log for operation.
|
||||
* @returns {Promise<TConversation>} The conversation object.
|
||||
*/
|
||||
saveConvo: async (req, { conversationId, newConversationId, ...convo }, metadata) => {
|
||||
try {
|
||||
if (metadata?.context) {
|
||||
logger.debug(`[saveConvo] ${metadata.context}`);
|
||||
}
|
||||
|
||||
const messages = await getMessages({ conversationId }, '_id');
|
||||
const update = { ...convo, messages, user: req.user.id };
|
||||
|
||||
if (newConversationId) {
|
||||
update.conversationId = newConversationId;
|
||||
}
|
||||
|
||||
if (req?.body?.isTemporary) {
|
||||
try {
|
||||
const appConfig = req.config;
|
||||
update.expiredAt = createTempChatExpirationDate(appConfig?.interfaceConfig);
|
||||
} catch (err) {
|
||||
logger.error('Error creating temporary chat expiration date:', err);
|
||||
logger.info(`---\`saveConvo\` context: ${metadata?.context}`);
|
||||
update.expiredAt = null;
|
||||
}
|
||||
} else {
|
||||
update.expiredAt = null;
|
||||
}
|
||||
|
||||
/** @type {{ $set: Partial<TConversation>; $unset?: Record<keyof TConversation, number> }} */
|
||||
const updateOperation = { $set: update };
|
||||
if (metadata && metadata.unsetFields && Object.keys(metadata.unsetFields).length > 0) {
|
||||
updateOperation.$unset = metadata.unsetFields;
|
||||
}
|
||||
|
||||
/** Note: the resulting Model object is necessary for Meilisearch operations */
|
||||
const conversation = await Conversation.findOneAndUpdate(
|
||||
{ conversationId, user: req.user.id },
|
||||
updateOperation,
|
||||
{
|
||||
new: true,
|
||||
upsert: true,
|
||||
},
|
||||
);
|
||||
|
||||
return conversation.toObject();
|
||||
} catch (error) {
|
||||
logger.error('[saveConvo] Error saving conversation', error);
|
||||
if (metadata && metadata?.context) {
|
||||
logger.info(`[saveConvo] ${metadata.context}`);
|
||||
}
|
||||
return { message: 'Error saving conversation' };
|
||||
}
|
||||
},
|
||||
bulkSaveConvos: async (conversations) => {
|
||||
try {
|
||||
const bulkOps = conversations.map((convo) => ({
|
||||
updateOne: {
|
||||
filter: { conversationId: convo.conversationId, user: convo.user },
|
||||
update: convo,
|
||||
upsert: true,
|
||||
timestamps: false,
|
||||
},
|
||||
}));
|
||||
|
||||
const result = await Conversation.bulkWrite(bulkOps);
|
||||
return result;
|
||||
} catch (error) {
|
||||
logger.error('[bulkSaveConvos] Error saving conversations in bulk', error);
|
||||
throw new Error('Failed to save conversations in bulk.');
|
||||
}
|
||||
},
|
||||
getConvosByCursor: async (
|
||||
user,
|
||||
{
|
||||
cursor,
|
||||
limit = 25,
|
||||
isArchived = false,
|
||||
tags,
|
||||
search,
|
||||
sortBy = 'updatedAt',
|
||||
sortDirection = 'desc',
|
||||
} = {},
|
||||
) => {
|
||||
const filters = [{ user }];
|
||||
if (isArchived) {
|
||||
filters.push({ isArchived: true });
|
||||
} else {
|
||||
filters.push({ $or: [{ isArchived: false }, { isArchived: { $exists: false } }] });
|
||||
}
|
||||
|
||||
if (Array.isArray(tags) && tags.length > 0) {
|
||||
filters.push({ tags: { $in: tags } });
|
||||
}
|
||||
|
||||
filters.push({ $or: [{ expiredAt: null }, { expiredAt: { $exists: false } }] });
|
||||
|
||||
if (search) {
|
||||
try {
|
||||
const meiliResults = await Conversation.meiliSearch(search, { filter: `user = "${user}"` });
|
||||
const matchingIds = Array.isArray(meiliResults.hits)
|
||||
? meiliResults.hits.map((result) => result.conversationId)
|
||||
: [];
|
||||
if (!matchingIds.length) {
|
||||
return { conversations: [], nextCursor: null };
|
||||
}
|
||||
filters.push({ conversationId: { $in: matchingIds } });
|
||||
} catch (error) {
|
||||
logger.error('[getConvosByCursor] Error during meiliSearch', error);
|
||||
throw new Error('Error during meiliSearch');
|
||||
}
|
||||
}
|
||||
|
||||
const validSortFields = ['title', 'createdAt', 'updatedAt'];
|
||||
if (!validSortFields.includes(sortBy)) {
|
||||
throw new Error(
|
||||
`Invalid sortBy field: ${sortBy}. Must be one of ${validSortFields.join(', ')}`,
|
||||
);
|
||||
}
|
||||
const finalSortBy = sortBy;
|
||||
const finalSortDirection = sortDirection === 'asc' ? 'asc' : 'desc';
|
||||
|
||||
let cursorFilter = null;
|
||||
if (cursor) {
|
||||
try {
|
||||
const decoded = JSON.parse(Buffer.from(cursor, 'base64').toString());
|
||||
const { primary, secondary } = decoded;
|
||||
const primaryValue = finalSortBy === 'title' ? primary : new Date(primary);
|
||||
const secondaryValue = new Date(secondary);
|
||||
const op = finalSortDirection === 'asc' ? '$gt' : '$lt';
|
||||
|
||||
cursorFilter = {
|
||||
$or: [
|
||||
{ [finalSortBy]: { [op]: primaryValue } },
|
||||
{
|
||||
[finalSortBy]: primaryValue,
|
||||
updatedAt: { [op]: secondaryValue },
|
||||
},
|
||||
],
|
||||
};
|
||||
} catch (err) {
|
||||
logger.warn('[getConvosByCursor] Invalid cursor format, starting from beginning');
|
||||
}
|
||||
if (cursorFilter) {
|
||||
filters.push(cursorFilter);
|
||||
}
|
||||
}
|
||||
|
||||
const query = filters.length === 1 ? filters[0] : { $and: filters };
|
||||
|
||||
try {
|
||||
const sortOrder = finalSortDirection === 'asc' ? 1 : -1;
|
||||
const sortObj = { [finalSortBy]: sortOrder };
|
||||
|
||||
if (finalSortBy !== 'updatedAt') {
|
||||
sortObj.updatedAt = sortOrder;
|
||||
}
|
||||
|
||||
const convos = await Conversation.find(query)
|
||||
.select(
|
||||
'conversationId endpoint title createdAt updatedAt user model agent_id assistant_id spec iconURL',
|
||||
)
|
||||
.sort(sortObj)
|
||||
.limit(limit + 1)
|
||||
.lean();
|
||||
|
||||
let nextCursor = null;
|
||||
if (convos.length > limit) {
|
||||
convos.pop(); // Remove extra item used to detect next page
|
||||
// Create cursor from the last RETURNED item (not the popped one)
|
||||
const lastReturned = convos[convos.length - 1];
|
||||
const primaryValue = lastReturned[finalSortBy];
|
||||
const primaryStr = finalSortBy === 'title' ? primaryValue : primaryValue.toISOString();
|
||||
const secondaryStr = lastReturned.updatedAt.toISOString();
|
||||
const composite = { primary: primaryStr, secondary: secondaryStr };
|
||||
nextCursor = Buffer.from(JSON.stringify(composite)).toString('base64');
|
||||
}
|
||||
|
||||
return { conversations: convos, nextCursor };
|
||||
} catch (error) {
|
||||
logger.error('[getConvosByCursor] Error getting conversations', error);
|
||||
throw new Error('Error getting conversations');
|
||||
}
|
||||
},
|
||||
getConvosQueried: async (user, convoIds, cursor = null, limit = 25) => {
|
||||
try {
|
||||
if (!convoIds?.length) {
|
||||
return { conversations: [], nextCursor: null, convoMap: {} };
|
||||
}
|
||||
|
||||
const conversationIds = convoIds.map((convo) => convo.conversationId);
|
||||
|
||||
const results = await Conversation.find({
|
||||
user,
|
||||
conversationId: { $in: conversationIds },
|
||||
$or: [{ expiredAt: { $exists: false } }, { expiredAt: null }],
|
||||
}).lean();
|
||||
|
||||
results.sort((a, b) => new Date(b.updatedAt) - new Date(a.updatedAt));
|
||||
|
||||
let filtered = results;
|
||||
if (cursor && cursor !== 'start') {
|
||||
const cursorDate = new Date(cursor);
|
||||
filtered = results.filter((convo) => new Date(convo.updatedAt) < cursorDate);
|
||||
}
|
||||
|
||||
const limited = filtered.slice(0, limit + 1);
|
||||
let nextCursor = null;
|
||||
if (limited.length > limit) {
|
||||
limited.pop(); // Remove extra item used to detect next page
|
||||
// Create cursor from the last RETURNED item (not the popped one)
|
||||
nextCursor = limited[limited.length - 1].updatedAt.toISOString();
|
||||
}
|
||||
|
||||
const convoMap = {};
|
||||
limited.forEach((convo) => {
|
||||
convoMap[convo.conversationId] = convo;
|
||||
});
|
||||
|
||||
return { conversations: limited, nextCursor, convoMap };
|
||||
} catch (error) {
|
||||
logger.error('[getConvosQueried] Error getting conversations', error);
|
||||
throw new Error('Error fetching conversations');
|
||||
}
|
||||
},
|
||||
getConvo,
|
||||
/* chore: this method is not properly error handled */
|
||||
getConvoTitle: async (user, conversationId) => {
|
||||
try {
|
||||
const convo = await getConvo(user, conversationId);
|
||||
/* ChatGPT Browser was triggering error here due to convo being saved later */
|
||||
if (convo && !convo.title) {
|
||||
return null;
|
||||
} else {
|
||||
// TypeError: Cannot read properties of null (reading 'title')
|
||||
return convo?.title || 'New Chat';
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error('[getConvoTitle] Error getting conversation title', error);
|
||||
throw new Error('Error getting conversation title');
|
||||
}
|
||||
},
|
||||
/**
|
||||
* Asynchronously deletes conversations and associated messages for a given user and filter.
|
||||
*
|
||||
* @async
|
||||
* @function
|
||||
* @param {string|ObjectId} user - The user's ID.
|
||||
* @param {Object} filter - Additional filter criteria for the conversations to be deleted.
|
||||
* @returns {Promise<{ n: number, ok: number, deletedCount: number, messages: { n: number, ok: number, deletedCount: number } }>}
|
||||
* An object containing the count of deleted conversations and associated messages.
|
||||
* @throws {Error} Throws an error if there's an issue with the database operations.
|
||||
*
|
||||
* @example
|
||||
* const user = 'someUserId';
|
||||
* const filter = { someField: 'someValue' };
|
||||
* const result = await deleteConvos(user, filter);
|
||||
* logger.error(result); // { n: 5, ok: 1, deletedCount: 5, messages: { n: 10, ok: 1, deletedCount: 10 } }
|
||||
*/
|
||||
deleteConvos: async (user, filter) => {
|
||||
try {
|
||||
const userFilter = { ...filter, user };
|
||||
const conversations = await Conversation.find(userFilter).select('conversationId');
|
||||
const conversationIds = conversations.map((c) => c.conversationId);
|
||||
|
||||
if (!conversationIds.length) {
|
||||
throw new Error('Conversation not found or already deleted.');
|
||||
}
|
||||
|
||||
const deleteConvoResult = await Conversation.deleteMany(userFilter);
|
||||
|
||||
const deleteMessagesResult = await deleteMessages({
|
||||
conversationId: { $in: conversationIds },
|
||||
});
|
||||
|
||||
return { ...deleteConvoResult, messages: deleteMessagesResult };
|
||||
} catch (error) {
|
||||
logger.error('[deleteConvos] Error deleting conversations and messages', error);
|
||||
throw error;
|
||||
}
|
||||
},
|
||||
};
|
||||
833
api/models/Conversation.spec.js
Normal file
833
api/models/Conversation.spec.js
Normal file
@ -0,0 +1,833 @@
|
||||
const mongoose = require('mongoose');
|
||||
const { v4: uuidv4 } = require('uuid');
|
||||
const { EModelEndpoint } = require('librechat-data-provider');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
const {
|
||||
deleteNullOrEmptyConversations,
|
||||
searchConversation,
|
||||
getConvosByCursor,
|
||||
getConvosQueried,
|
||||
getConvoFiles,
|
||||
getConvoTitle,
|
||||
deleteConvos,
|
||||
saveConvo,
|
||||
getConvo,
|
||||
} = require('./Conversation');
|
||||
jest.mock('~/server/services/Config/app');
|
||||
jest.mock('./Message');
|
||||
const { getMessages, deleteMessages } = require('./Message');
|
||||
|
||||
const { Conversation } = require('~/db/models');
|
||||
|
||||
describe('Conversation Operations', () => {
|
||||
let mongoServer;
|
||||
let mockReq;
|
||||
let mockConversationData;
|
||||
|
||||
beforeAll(async () => {
|
||||
mongoServer = await MongoMemoryServer.create();
|
||||
const mongoUri = mongoServer.getUri();
|
||||
await mongoose.connect(mongoUri);
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await mongoose.disconnect();
|
||||
await mongoServer.stop();
|
||||
});
|
||||
|
||||
beforeEach(async () => {
|
||||
// Clear database
|
||||
await Conversation.deleteMany({});
|
||||
|
||||
// Reset mocks
|
||||
jest.clearAllMocks();
|
||||
|
||||
// Default mock implementations
|
||||
getMessages.mockResolvedValue([]);
|
||||
deleteMessages.mockResolvedValue({ deletedCount: 0 });
|
||||
|
||||
mockReq = {
|
||||
user: { id: 'user123' },
|
||||
body: {},
|
||||
config: {
|
||||
interfaceConfig: {
|
||||
temporaryChatRetention: 24, // Default 24 hours
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
mockConversationData = {
|
||||
conversationId: uuidv4(),
|
||||
title: 'Test Conversation',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
};
|
||||
});
|
||||
|
||||
describe('saveConvo', () => {
|
||||
it('should save a conversation for an authenticated user', async () => {
|
||||
const result = await saveConvo(mockReq, mockConversationData);
|
||||
|
||||
expect(result.conversationId).toBe(mockConversationData.conversationId);
|
||||
expect(result.user).toBe('user123');
|
||||
expect(result.title).toBe('Test Conversation');
|
||||
expect(result.endpoint).toBe(EModelEndpoint.openAI);
|
||||
|
||||
// Verify the conversation was actually saved to the database
|
||||
const savedConvo = await Conversation.findOne({
|
||||
conversationId: mockConversationData.conversationId,
|
||||
user: 'user123',
|
||||
});
|
||||
expect(savedConvo).toBeTruthy();
|
||||
expect(savedConvo.title).toBe('Test Conversation');
|
||||
});
|
||||
|
||||
it('should query messages when saving a conversation', async () => {
|
||||
// Mock messages as ObjectIds
|
||||
const mongoose = require('mongoose');
|
||||
const mockMessages = [new mongoose.Types.ObjectId(), new mongoose.Types.ObjectId()];
|
||||
getMessages.mockResolvedValue(mockMessages);
|
||||
|
||||
await saveConvo(mockReq, mockConversationData);
|
||||
|
||||
// Verify that getMessages was called with correct parameters
|
||||
expect(getMessages).toHaveBeenCalledWith(
|
||||
{ conversationId: mockConversationData.conversationId },
|
||||
'_id',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle newConversationId when provided', async () => {
|
||||
const newConversationId = uuidv4();
|
||||
const result = await saveConvo(mockReq, {
|
||||
...mockConversationData,
|
||||
newConversationId,
|
||||
});
|
||||
|
||||
expect(result.conversationId).toBe(newConversationId);
|
||||
});
|
||||
|
||||
it('should handle unsetFields metadata', async () => {
|
||||
const metadata = {
|
||||
unsetFields: { someField: 1 },
|
||||
};
|
||||
|
||||
await saveConvo(mockReq, mockConversationData, metadata);
|
||||
|
||||
const savedConvo = await Conversation.findOne({
|
||||
conversationId: mockConversationData.conversationId,
|
||||
});
|
||||
expect(savedConvo.someField).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('isTemporary conversation handling', () => {
|
||||
it('should save a conversation with expiredAt when isTemporary is true', async () => {
|
||||
// Mock app config with 24 hour retention
|
||||
mockReq.config.interfaceConfig.temporaryChatRetention = 24;
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
|
||||
const beforeSave = new Date();
|
||||
const result = await saveConvo(mockReq, mockConversationData);
|
||||
const afterSave = new Date();
|
||||
|
||||
expect(result.conversationId).toBe(mockConversationData.conversationId);
|
||||
expect(result.expiredAt).toBeDefined();
|
||||
expect(result.expiredAt).toBeInstanceOf(Date);
|
||||
|
||||
// Verify expiredAt is approximately 24 hours in the future
|
||||
const expectedExpirationTime = new Date(beforeSave.getTime() + 24 * 60 * 60 * 1000);
|
||||
const actualExpirationTime = new Date(result.expiredAt);
|
||||
|
||||
expect(actualExpirationTime.getTime()).toBeGreaterThanOrEqual(
|
||||
expectedExpirationTime.getTime() - 1000,
|
||||
);
|
||||
expect(actualExpirationTime.getTime()).toBeLessThanOrEqual(
|
||||
new Date(afterSave.getTime() + 24 * 60 * 60 * 1000 + 1000).getTime(),
|
||||
);
|
||||
});
|
||||
|
||||
it('should save a conversation without expiredAt when isTemporary is false', async () => {
|
||||
mockReq.body = { isTemporary: false };
|
||||
|
||||
const result = await saveConvo(mockReq, mockConversationData);
|
||||
|
||||
expect(result.conversationId).toBe(mockConversationData.conversationId);
|
||||
expect(result.expiredAt).toBeNull();
|
||||
});
|
||||
|
||||
it('should save a conversation without expiredAt when isTemporary is not provided', async () => {
|
||||
// No isTemporary in body
|
||||
mockReq.body = {};
|
||||
|
||||
const result = await saveConvo(mockReq, mockConversationData);
|
||||
|
||||
expect(result.conversationId).toBe(mockConversationData.conversationId);
|
||||
expect(result.expiredAt).toBeNull();
|
||||
});
|
||||
|
||||
it('should use custom retention period from config', async () => {
|
||||
// Mock app config with 48 hour retention
|
||||
mockReq.config.interfaceConfig.temporaryChatRetention = 48;
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
|
||||
const beforeSave = new Date();
|
||||
const result = await saveConvo(mockReq, mockConversationData);
|
||||
|
||||
expect(result.expiredAt).toBeDefined();
|
||||
|
||||
// Verify expiredAt is approximately 48 hours in the future
|
||||
const expectedExpirationTime = new Date(beforeSave.getTime() + 48 * 60 * 60 * 1000);
|
||||
const actualExpirationTime = new Date(result.expiredAt);
|
||||
|
||||
expect(actualExpirationTime.getTime()).toBeGreaterThanOrEqual(
|
||||
expectedExpirationTime.getTime() - 1000,
|
||||
);
|
||||
expect(actualExpirationTime.getTime()).toBeLessThanOrEqual(
|
||||
expectedExpirationTime.getTime() + 1000,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle minimum retention period (1 hour)', async () => {
|
||||
// Mock app config with less than minimum retention
|
||||
mockReq.config.interfaceConfig.temporaryChatRetention = 0.5; // Half hour - should be clamped to 1 hour
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
|
||||
const beforeSave = new Date();
|
||||
const result = await saveConvo(mockReq, mockConversationData);
|
||||
|
||||
expect(result.expiredAt).toBeDefined();
|
||||
|
||||
// Verify expiredAt is approximately 1 hour in the future (minimum)
|
||||
const expectedExpirationTime = new Date(beforeSave.getTime() + 1 * 60 * 60 * 1000);
|
||||
const actualExpirationTime = new Date(result.expiredAt);
|
||||
|
||||
expect(actualExpirationTime.getTime()).toBeGreaterThanOrEqual(
|
||||
expectedExpirationTime.getTime() - 1000,
|
||||
);
|
||||
expect(actualExpirationTime.getTime()).toBeLessThanOrEqual(
|
||||
expectedExpirationTime.getTime() + 1000,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle maximum retention period (8760 hours)', async () => {
|
||||
// Mock app config with more than maximum retention
|
||||
mockReq.config.interfaceConfig.temporaryChatRetention = 10000; // Should be clamped to 8760 hours
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
|
||||
const beforeSave = new Date();
|
||||
const result = await saveConvo(mockReq, mockConversationData);
|
||||
|
||||
expect(result.expiredAt).toBeDefined();
|
||||
|
||||
// Verify expiredAt is approximately 8760 hours (1 year) in the future
|
||||
const expectedExpirationTime = new Date(beforeSave.getTime() + 8760 * 60 * 60 * 1000);
|
||||
const actualExpirationTime = new Date(result.expiredAt);
|
||||
|
||||
expect(actualExpirationTime.getTime()).toBeGreaterThanOrEqual(
|
||||
expectedExpirationTime.getTime() - 1000,
|
||||
);
|
||||
expect(actualExpirationTime.getTime()).toBeLessThanOrEqual(
|
||||
expectedExpirationTime.getTime() + 1000,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle missing config gracefully', async () => {
|
||||
// Simulate missing config - should use default retention period
|
||||
delete mockReq.config;
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
|
||||
const beforeSave = new Date();
|
||||
const result = await saveConvo(mockReq, mockConversationData);
|
||||
const afterSave = new Date();
|
||||
|
||||
// Should still save the conversation with default retention period (30 days)
|
||||
expect(result.conversationId).toBe(mockConversationData.conversationId);
|
||||
expect(result.expiredAt).toBeDefined();
|
||||
expect(result.expiredAt).toBeInstanceOf(Date);
|
||||
|
||||
// Verify expiredAt is approximately 30 days in the future (720 hours)
|
||||
const expectedExpirationTime = new Date(beforeSave.getTime() + 720 * 60 * 60 * 1000);
|
||||
const actualExpirationTime = new Date(result.expiredAt);
|
||||
|
||||
expect(actualExpirationTime.getTime()).toBeGreaterThanOrEqual(
|
||||
expectedExpirationTime.getTime() - 1000,
|
||||
);
|
||||
expect(actualExpirationTime.getTime()).toBeLessThanOrEqual(
|
||||
new Date(afterSave.getTime() + 720 * 60 * 60 * 1000 + 1000).getTime(),
|
||||
);
|
||||
});
|
||||
|
||||
it('should use default retention when config is not provided', async () => {
|
||||
// Mock getAppConfig to return empty config
|
||||
mockReq.config = {}; // Empty config
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
|
||||
const beforeSave = new Date();
|
||||
const result = await saveConvo(mockReq, mockConversationData);
|
||||
|
||||
expect(result.expiredAt).toBeDefined();
|
||||
|
||||
// Default retention is 30 days (720 hours)
|
||||
const expectedExpirationTime = new Date(beforeSave.getTime() + 30 * 24 * 60 * 60 * 1000);
|
||||
const actualExpirationTime = new Date(result.expiredAt);
|
||||
|
||||
expect(actualExpirationTime.getTime()).toBeGreaterThanOrEqual(
|
||||
expectedExpirationTime.getTime() - 1000,
|
||||
);
|
||||
expect(actualExpirationTime.getTime()).toBeLessThanOrEqual(
|
||||
expectedExpirationTime.getTime() + 1000,
|
||||
);
|
||||
});
|
||||
|
||||
it('should update expiredAt when saving existing temporary conversation', async () => {
|
||||
// First save a temporary conversation
|
||||
mockReq.config.interfaceConfig.temporaryChatRetention = 24;
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
const firstSave = await saveConvo(mockReq, mockConversationData);
|
||||
const originalExpiredAt = firstSave.expiredAt;
|
||||
|
||||
// Wait a bit to ensure time difference
|
||||
await new Promise((resolve) => setTimeout(resolve, 100));
|
||||
|
||||
// Save again with same conversationId but different title
|
||||
const updatedData = { ...mockConversationData, title: 'Updated Title' };
|
||||
const secondSave = await saveConvo(mockReq, updatedData);
|
||||
|
||||
// Should update title and create new expiredAt
|
||||
expect(secondSave.title).toBe('Updated Title');
|
||||
expect(secondSave.expiredAt).toBeDefined();
|
||||
expect(new Date(secondSave.expiredAt).getTime()).toBeGreaterThan(
|
||||
new Date(originalExpiredAt).getTime(),
|
||||
);
|
||||
});
|
||||
|
||||
it('should not set expiredAt when updating non-temporary conversation', async () => {
|
||||
// First save a non-temporary conversation
|
||||
mockReq.body = { isTemporary: false };
|
||||
const firstSave = await saveConvo(mockReq, mockConversationData);
|
||||
expect(firstSave.expiredAt).toBeNull();
|
||||
|
||||
// Update without isTemporary flag
|
||||
mockReq.body = {};
|
||||
const updatedData = { ...mockConversationData, title: 'Updated Title' };
|
||||
const secondSave = await saveConvo(mockReq, updatedData);
|
||||
|
||||
expect(secondSave.title).toBe('Updated Title');
|
||||
expect(secondSave.expiredAt).toBeNull();
|
||||
});
|
||||
|
||||
it('should filter out expired conversations in getConvosByCursor', async () => {
|
||||
// Create some test conversations
|
||||
const nonExpiredConvo = await Conversation.create({
|
||||
conversationId: uuidv4(),
|
||||
user: 'user123',
|
||||
title: 'Non-expired',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
expiredAt: null,
|
||||
updatedAt: new Date(),
|
||||
});
|
||||
|
||||
await Conversation.create({
|
||||
conversationId: uuidv4(),
|
||||
user: 'user123',
|
||||
title: 'Future expired',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
expiredAt: new Date(Date.now() + 24 * 60 * 60 * 1000), // 24 hours from now
|
||||
updatedAt: new Date(),
|
||||
});
|
||||
|
||||
// Mock Meili search
|
||||
Conversation.meiliSearch = jest.fn().mockResolvedValue({ hits: [] });
|
||||
|
||||
const result = await getConvosByCursor('user123');
|
||||
|
||||
// Should only return conversations with null or non-existent expiredAt
|
||||
expect(result.conversations).toHaveLength(1);
|
||||
expect(result.conversations[0].conversationId).toBe(nonExpiredConvo.conversationId);
|
||||
});
|
||||
|
||||
it('should filter out expired conversations in getConvosQueried', async () => {
|
||||
// Create test conversations
|
||||
const nonExpiredConvo = await Conversation.create({
|
||||
conversationId: uuidv4(),
|
||||
user: 'user123',
|
||||
title: 'Non-expired',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
expiredAt: null,
|
||||
});
|
||||
|
||||
const expiredConvo = await Conversation.create({
|
||||
conversationId: uuidv4(),
|
||||
user: 'user123',
|
||||
title: 'Expired',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
expiredAt: new Date(Date.now() + 24 * 60 * 60 * 1000),
|
||||
});
|
||||
|
||||
const convoIds = [
|
||||
{ conversationId: nonExpiredConvo.conversationId },
|
||||
{ conversationId: expiredConvo.conversationId },
|
||||
];
|
||||
|
||||
const result = await getConvosQueried('user123', convoIds);
|
||||
|
||||
// Should only return the non-expired conversation
|
||||
expect(result.conversations).toHaveLength(1);
|
||||
expect(result.conversations[0].conversationId).toBe(nonExpiredConvo.conversationId);
|
||||
expect(result.convoMap[nonExpiredConvo.conversationId]).toBeDefined();
|
||||
expect(result.convoMap[expiredConvo.conversationId]).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('searchConversation', () => {
|
||||
it('should find a conversation by conversationId', async () => {
|
||||
await Conversation.create({
|
||||
conversationId: mockConversationData.conversationId,
|
||||
user: 'user123',
|
||||
title: 'Test',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
});
|
||||
|
||||
const result = await searchConversation(mockConversationData.conversationId);
|
||||
|
||||
expect(result).toBeTruthy();
|
||||
expect(result.conversationId).toBe(mockConversationData.conversationId);
|
||||
expect(result.user).toBe('user123');
|
||||
expect(result.title).toBeUndefined(); // Only returns conversationId and user
|
||||
});
|
||||
|
||||
it('should return null if conversation not found', async () => {
|
||||
const result = await searchConversation('non-existent-id');
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe('getConvo', () => {
|
||||
it('should retrieve a conversation for a user', async () => {
|
||||
await Conversation.create({
|
||||
conversationId: mockConversationData.conversationId,
|
||||
user: 'user123',
|
||||
title: 'Test Conversation',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
});
|
||||
|
||||
const result = await getConvo('user123', mockConversationData.conversationId);
|
||||
|
||||
expect(result.conversationId).toBe(mockConversationData.conversationId);
|
||||
expect(result.user).toBe('user123');
|
||||
expect(result.title).toBe('Test Conversation');
|
||||
});
|
||||
|
||||
it('should return null if conversation not found', async () => {
|
||||
const result = await getConvo('user123', 'non-existent-id');
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe('getConvoTitle', () => {
|
||||
it('should return the conversation title', async () => {
|
||||
await Conversation.create({
|
||||
conversationId: mockConversationData.conversationId,
|
||||
user: 'user123',
|
||||
title: 'Test Title',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
});
|
||||
|
||||
const result = await getConvoTitle('user123', mockConversationData.conversationId);
|
||||
expect(result).toBe('Test Title');
|
||||
});
|
||||
|
||||
it('should return null if conversation has no title', async () => {
|
||||
await Conversation.create({
|
||||
conversationId: mockConversationData.conversationId,
|
||||
user: 'user123',
|
||||
title: null,
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
});
|
||||
|
||||
const result = await getConvoTitle('user123', mockConversationData.conversationId);
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
it('should return "New Chat" if conversation not found', async () => {
|
||||
const result = await getConvoTitle('user123', 'non-existent-id');
|
||||
expect(result).toBe('New Chat');
|
||||
});
|
||||
});
|
||||
|
||||
describe('getConvoFiles', () => {
|
||||
it('should return conversation files', async () => {
|
||||
const files = ['file1', 'file2'];
|
||||
await Conversation.create({
|
||||
conversationId: mockConversationData.conversationId,
|
||||
user: 'user123',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
files,
|
||||
});
|
||||
|
||||
const result = await getConvoFiles(mockConversationData.conversationId);
|
||||
expect(result).toEqual(files);
|
||||
});
|
||||
|
||||
it('should return empty array if no files', async () => {
|
||||
await Conversation.create({
|
||||
conversationId: mockConversationData.conversationId,
|
||||
user: 'user123',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
});
|
||||
|
||||
const result = await getConvoFiles(mockConversationData.conversationId);
|
||||
expect(result).toEqual([]);
|
||||
});
|
||||
|
||||
it('should return empty array if conversation not found', async () => {
|
||||
const result = await getConvoFiles('non-existent-id');
|
||||
expect(result).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('deleteConvos', () => {
|
||||
it('should delete conversations and associated messages', async () => {
|
||||
await Conversation.create({
|
||||
conversationId: mockConversationData.conversationId,
|
||||
user: 'user123',
|
||||
title: 'To Delete',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
});
|
||||
|
||||
deleteMessages.mockResolvedValue({ deletedCount: 5 });
|
||||
|
||||
const result = await deleteConvos('user123', {
|
||||
conversationId: mockConversationData.conversationId,
|
||||
});
|
||||
|
||||
expect(result.deletedCount).toBe(1);
|
||||
expect(result.messages.deletedCount).toBe(5);
|
||||
expect(deleteMessages).toHaveBeenCalledWith({
|
||||
conversationId: { $in: [mockConversationData.conversationId] },
|
||||
});
|
||||
|
||||
// Verify conversation was deleted
|
||||
const deletedConvo = await Conversation.findOne({
|
||||
conversationId: mockConversationData.conversationId,
|
||||
});
|
||||
expect(deletedConvo).toBeNull();
|
||||
});
|
||||
|
||||
it('should throw error if no conversations found', async () => {
|
||||
await expect(deleteConvos('user123', { conversationId: 'non-existent' })).rejects.toThrow(
|
||||
'Conversation not found or already deleted.',
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('deleteNullOrEmptyConversations', () => {
|
||||
it('should delete conversations with null, empty, or missing conversationIds', async () => {
|
||||
// Since conversationId is required by the schema, we can't create documents with null/missing IDs
|
||||
// This test should verify the function works when such documents exist (e.g., from data corruption)
|
||||
|
||||
// For this test, let's create a valid conversation and verify the function doesn't delete it
|
||||
await Conversation.create({
|
||||
conversationId: mockConversationData.conversationId,
|
||||
user: 'user4',
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
});
|
||||
|
||||
deleteMessages.mockResolvedValue({ deletedCount: 0 });
|
||||
|
||||
const result = await deleteNullOrEmptyConversations();
|
||||
|
||||
expect(result.conversations.deletedCount).toBe(0); // No invalid conversations to delete
|
||||
expect(result.messages.deletedCount).toBe(0);
|
||||
|
||||
// Verify valid conversation remains
|
||||
const remainingConvos = await Conversation.find({});
|
||||
expect(remainingConvos).toHaveLength(1);
|
||||
expect(remainingConvos[0].conversationId).toBe(mockConversationData.conversationId);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Error Handling', () => {
|
||||
it('should handle database errors in saveConvo', async () => {
|
||||
// Force a database error by disconnecting
|
||||
await mongoose.disconnect();
|
||||
|
||||
const result = await saveConvo(mockReq, mockConversationData);
|
||||
|
||||
expect(result).toEqual({ message: 'Error saving conversation' });
|
||||
|
||||
// Reconnect for other tests
|
||||
await mongoose.connect(mongoServer.getUri());
|
||||
});
|
||||
});
|
||||
|
||||
describe('getConvosByCursor pagination', () => {
|
||||
/**
|
||||
* Helper to create conversations with specific timestamps
|
||||
* Uses collection.insertOne to bypass Mongoose timestamps entirely
|
||||
*/
|
||||
const createConvoWithTimestamps = async (index, createdAt, updatedAt) => {
|
||||
const conversationId = uuidv4();
|
||||
// Use collection-level insert to bypass Mongoose timestamps
|
||||
await Conversation.collection.insertOne({
|
||||
conversationId,
|
||||
user: 'user123',
|
||||
title: `Conversation ${index}`,
|
||||
endpoint: EModelEndpoint.openAI,
|
||||
expiredAt: null,
|
||||
isArchived: false,
|
||||
createdAt,
|
||||
updatedAt,
|
||||
});
|
||||
return Conversation.findOne({ conversationId }).lean();
|
||||
};
|
||||
|
||||
it('should not skip conversations at page boundaries', async () => {
|
||||
// Create 30 conversations to ensure pagination (limit is 25)
|
||||
const baseTime = new Date('2026-01-01T00:00:00.000Z');
|
||||
const convos = [];
|
||||
|
||||
for (let i = 0; i < 30; i++) {
|
||||
const updatedAt = new Date(baseTime.getTime() - i * 60000); // Each 1 minute apart
|
||||
const convo = await createConvoWithTimestamps(i, updatedAt, updatedAt);
|
||||
convos.push(convo);
|
||||
}
|
||||
|
||||
// Fetch first page
|
||||
const page1 = await getConvosByCursor('user123', { limit: 25 });
|
||||
|
||||
expect(page1.conversations).toHaveLength(25);
|
||||
expect(page1.nextCursor).toBeTruthy();
|
||||
|
||||
// Fetch second page using cursor
|
||||
const page2 = await getConvosByCursor('user123', {
|
||||
limit: 25,
|
||||
cursor: page1.nextCursor,
|
||||
});
|
||||
|
||||
// Should get remaining 5 conversations
|
||||
expect(page2.conversations).toHaveLength(5);
|
||||
expect(page2.nextCursor).toBeNull();
|
||||
|
||||
// Verify no duplicates and no gaps
|
||||
const allIds = [
|
||||
...page1.conversations.map((c) => c.conversationId),
|
||||
...page2.conversations.map((c) => c.conversationId),
|
||||
];
|
||||
const uniqueIds = new Set(allIds);
|
||||
|
||||
expect(uniqueIds.size).toBe(30); // All 30 conversations accounted for
|
||||
expect(allIds.length).toBe(30); // No duplicates
|
||||
});
|
||||
|
||||
it('should include conversation at exact page boundary (item 26 bug fix)', async () => {
|
||||
// This test specifically verifies the fix for the bug where item 26
|
||||
// (the first item that should appear on page 2) was being skipped
|
||||
|
||||
const baseTime = new Date('2026-01-01T12:00:00.000Z');
|
||||
|
||||
// Create exactly 26 conversations
|
||||
const convos = [];
|
||||
for (let i = 0; i < 26; i++) {
|
||||
const updatedAt = new Date(baseTime.getTime() - i * 60000);
|
||||
const convo = await createConvoWithTimestamps(i, updatedAt, updatedAt);
|
||||
convos.push(convo);
|
||||
}
|
||||
|
||||
// The 26th conversation (index 25) should be on page 2
|
||||
const item26 = convos[25];
|
||||
|
||||
// Fetch first page with limit 25
|
||||
const page1 = await getConvosByCursor('user123', { limit: 25 });
|
||||
|
||||
expect(page1.conversations).toHaveLength(25);
|
||||
expect(page1.nextCursor).toBeTruthy();
|
||||
|
||||
// Item 26 should NOT be in page 1
|
||||
const page1Ids = page1.conversations.map((c) => c.conversationId);
|
||||
expect(page1Ids).not.toContain(item26.conversationId);
|
||||
|
||||
// Fetch second page
|
||||
const page2 = await getConvosByCursor('user123', {
|
||||
limit: 25,
|
||||
cursor: page1.nextCursor,
|
||||
});
|
||||
|
||||
// Item 26 MUST be in page 2 (this was the bug - it was being skipped)
|
||||
expect(page2.conversations).toHaveLength(1);
|
||||
expect(page2.conversations[0].conversationId).toBe(item26.conversationId);
|
||||
});
|
||||
|
||||
it('should sort by updatedAt DESC by default', async () => {
|
||||
// Create conversations with different updatedAt times
|
||||
// Note: createdAt is older but updatedAt varies
|
||||
const convo1 = await createConvoWithTimestamps(
|
||||
1,
|
||||
new Date('2026-01-01T00:00:00.000Z'), // oldest created
|
||||
new Date('2026-01-03T00:00:00.000Z'), // most recently updated
|
||||
);
|
||||
|
||||
const convo2 = await createConvoWithTimestamps(
|
||||
2,
|
||||
new Date('2026-01-02T00:00:00.000Z'), // middle created
|
||||
new Date('2026-01-02T00:00:00.000Z'), // middle updated
|
||||
);
|
||||
|
||||
const convo3 = await createConvoWithTimestamps(
|
||||
3,
|
||||
new Date('2026-01-03T00:00:00.000Z'), // newest created
|
||||
new Date('2026-01-01T00:00:00.000Z'), // oldest updated
|
||||
);
|
||||
|
||||
const result = await getConvosByCursor('user123');
|
||||
|
||||
// Should be sorted by updatedAt DESC (most recent first)
|
||||
expect(result.conversations).toHaveLength(3);
|
||||
expect(result.conversations[0].conversationId).toBe(convo1.conversationId); // Jan 3 updatedAt
|
||||
expect(result.conversations[1].conversationId).toBe(convo2.conversationId); // Jan 2 updatedAt
|
||||
expect(result.conversations[2].conversationId).toBe(convo3.conversationId); // Jan 1 updatedAt
|
||||
});
|
||||
|
||||
it('should handle conversations with same updatedAt (tie-breaker)', async () => {
|
||||
const sameTime = new Date('2026-01-01T12:00:00.000Z');
|
||||
|
||||
// Create 3 conversations with exact same updatedAt
|
||||
const convo1 = await createConvoWithTimestamps(1, sameTime, sameTime);
|
||||
const convo2 = await createConvoWithTimestamps(2, sameTime, sameTime);
|
||||
const convo3 = await createConvoWithTimestamps(3, sameTime, sameTime);
|
||||
|
||||
const result = await getConvosByCursor('user123');
|
||||
|
||||
// All 3 should be returned (no skipping due to same timestamps)
|
||||
expect(result.conversations).toHaveLength(3);
|
||||
|
||||
const returnedIds = result.conversations.map((c) => c.conversationId);
|
||||
expect(returnedIds).toContain(convo1.conversationId);
|
||||
expect(returnedIds).toContain(convo2.conversationId);
|
||||
expect(returnedIds).toContain(convo3.conversationId);
|
||||
});
|
||||
|
||||
it('should handle cursor pagination with conversations updated during pagination', async () => {
|
||||
// Simulate the scenario where a conversation is updated between page fetches
|
||||
const baseTime = new Date('2026-01-01T00:00:00.000Z');
|
||||
|
||||
// Create 30 conversations
|
||||
for (let i = 0; i < 30; i++) {
|
||||
const updatedAt = new Date(baseTime.getTime() - i * 60000);
|
||||
await createConvoWithTimestamps(i, updatedAt, updatedAt);
|
||||
}
|
||||
|
||||
// Fetch first page
|
||||
const page1 = await getConvosByCursor('user123', { limit: 25 });
|
||||
expect(page1.conversations).toHaveLength(25);
|
||||
|
||||
// Now update one of the conversations that should be on page 2
|
||||
// to have a newer updatedAt (simulating user activity during pagination)
|
||||
const convosOnPage2 = await Conversation.find({ user: 'user123' })
|
||||
.sort({ updatedAt: -1 })
|
||||
.skip(25)
|
||||
.limit(5);
|
||||
|
||||
if (convosOnPage2.length > 0) {
|
||||
const updatedConvo = convosOnPage2[0];
|
||||
await Conversation.updateOne(
|
||||
{ _id: updatedConvo._id },
|
||||
{ updatedAt: new Date('2026-01-02T00:00:00.000Z') }, // Much newer
|
||||
);
|
||||
}
|
||||
|
||||
// Fetch second page with original cursor
|
||||
const page2 = await getConvosByCursor('user123', {
|
||||
limit: 25,
|
||||
cursor: page1.nextCursor,
|
||||
});
|
||||
|
||||
// The updated conversation might not be in page 2 anymore
|
||||
// (it moved to the front), but we should still get remaining items
|
||||
// without errors and without infinite loops
|
||||
expect(page2.conversations.length).toBeGreaterThanOrEqual(0);
|
||||
});
|
||||
|
||||
it('should correctly decode and use cursor for pagination', async () => {
|
||||
const baseTime = new Date('2026-01-01T00:00:00.000Z');
|
||||
|
||||
// Create 30 conversations
|
||||
for (let i = 0; i < 30; i++) {
|
||||
const updatedAt = new Date(baseTime.getTime() - i * 60000);
|
||||
await createConvoWithTimestamps(i, updatedAt, updatedAt);
|
||||
}
|
||||
|
||||
// Fetch first page
|
||||
const page1 = await getConvosByCursor('user123', { limit: 25 });
|
||||
|
||||
// Decode the cursor to verify it's based on the last RETURNED item
|
||||
const decodedCursor = JSON.parse(Buffer.from(page1.nextCursor, 'base64').toString());
|
||||
|
||||
// The cursor should match the last item in page1 (item at index 24)
|
||||
const lastReturnedItem = page1.conversations[24];
|
||||
|
||||
expect(new Date(decodedCursor.primary).getTime()).toBe(
|
||||
new Date(lastReturnedItem.updatedAt).getTime(),
|
||||
);
|
||||
});
|
||||
|
||||
it('should support sortBy createdAt when explicitly requested', async () => {
|
||||
// Create conversations with different timestamps
|
||||
const convo1 = await createConvoWithTimestamps(
|
||||
1,
|
||||
new Date('2026-01-03T00:00:00.000Z'), // newest created
|
||||
new Date('2026-01-01T00:00:00.000Z'), // oldest updated
|
||||
);
|
||||
|
||||
const convo2 = await createConvoWithTimestamps(
|
||||
2,
|
||||
new Date('2026-01-01T00:00:00.000Z'), // oldest created
|
||||
new Date('2026-01-03T00:00:00.000Z'), // newest updated
|
||||
);
|
||||
|
||||
// Verify timestamps were set correctly
|
||||
expect(new Date(convo1.createdAt).getTime()).toBe(
|
||||
new Date('2026-01-03T00:00:00.000Z').getTime(),
|
||||
);
|
||||
expect(new Date(convo2.createdAt).getTime()).toBe(
|
||||
new Date('2026-01-01T00:00:00.000Z').getTime(),
|
||||
);
|
||||
|
||||
const result = await getConvosByCursor('user123', { sortBy: 'createdAt' });
|
||||
|
||||
// Should be sorted by createdAt DESC
|
||||
expect(result.conversations).toHaveLength(2);
|
||||
expect(result.conversations[0].conversationId).toBe(convo1.conversationId); // Jan 3 createdAt
|
||||
expect(result.conversations[1].conversationId).toBe(convo2.conversationId); // Jan 1 createdAt
|
||||
});
|
||||
|
||||
it('should handle empty result set gracefully', async () => {
|
||||
const result = await getConvosByCursor('user123');
|
||||
|
||||
expect(result.conversations).toHaveLength(0);
|
||||
expect(result.nextCursor).toBeNull();
|
||||
});
|
||||
|
||||
it('should handle exactly limit number of conversations (no next page)', async () => {
|
||||
const baseTime = new Date('2026-01-01T00:00:00.000Z');
|
||||
|
||||
// Create exactly 25 conversations (equal to default limit)
|
||||
for (let i = 0; i < 25; i++) {
|
||||
const updatedAt = new Date(baseTime.getTime() - i * 60000);
|
||||
await createConvoWithTimestamps(i, updatedAt, updatedAt);
|
||||
}
|
||||
|
||||
const result = await getConvosByCursor('user123', { limit: 25 });
|
||||
|
||||
expect(result.conversations).toHaveLength(25);
|
||||
expect(result.nextCursor).toBeNull(); // No next page
|
||||
});
|
||||
});
|
||||
});
|
||||
284
api/models/ConversationTag.js
Normal file
284
api/models/ConversationTag.js
Normal file
@ -0,0 +1,284 @@
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { ConversationTag, Conversation } = require('~/db/models');
|
||||
|
||||
/**
|
||||
* Retrieves all conversation tags for a user.
|
||||
* @param {string} user - The user ID.
|
||||
* @returns {Promise<Array>} An array of conversation tags.
|
||||
*/
|
||||
const getConversationTags = async (user) => {
|
||||
try {
|
||||
return await ConversationTag.find({ user }).sort({ position: 1 }).lean();
|
||||
} catch (error) {
|
||||
logger.error('[getConversationTags] Error getting conversation tags', error);
|
||||
throw new Error('Error getting conversation tags');
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Creates a new conversation tag.
|
||||
* @param {string} user - The user ID.
|
||||
* @param {Object} data - The tag data.
|
||||
* @param {string} data.tag - The tag name.
|
||||
* @param {string} [data.description] - The tag description.
|
||||
* @param {boolean} [data.addToConversation] - Whether to add the tag to a conversation.
|
||||
* @param {string} [data.conversationId] - The conversation ID to add the tag to.
|
||||
* @returns {Promise<Object>} The created tag.
|
||||
*/
|
||||
const createConversationTag = async (user, data) => {
|
||||
try {
|
||||
const { tag, description, addToConversation, conversationId } = data;
|
||||
|
||||
const existingTag = await ConversationTag.findOne({ user, tag }).lean();
|
||||
if (existingTag) {
|
||||
return existingTag;
|
||||
}
|
||||
|
||||
const maxPosition = await ConversationTag.findOne({ user }).sort('-position').lean();
|
||||
const position = (maxPosition?.position || 0) + 1;
|
||||
|
||||
const newTag = await ConversationTag.findOneAndUpdate(
|
||||
{ tag, user },
|
||||
{
|
||||
tag,
|
||||
user,
|
||||
count: addToConversation ? 1 : 0,
|
||||
position,
|
||||
description,
|
||||
$setOnInsert: { createdAt: new Date() },
|
||||
},
|
||||
{
|
||||
new: true,
|
||||
upsert: true,
|
||||
lean: true,
|
||||
},
|
||||
);
|
||||
|
||||
if (addToConversation && conversationId) {
|
||||
await Conversation.findOneAndUpdate(
|
||||
{ user, conversationId },
|
||||
{ $addToSet: { tags: tag } },
|
||||
{ new: true },
|
||||
);
|
||||
}
|
||||
|
||||
return newTag;
|
||||
} catch (error) {
|
||||
logger.error('[createConversationTag] Error creating conversation tag', error);
|
||||
throw new Error('Error creating conversation tag');
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Updates an existing conversation tag.
|
||||
* @param {string} user - The user ID.
|
||||
* @param {string} oldTag - The current tag name.
|
||||
* @param {Object} data - The updated tag data.
|
||||
* @param {string} [data.tag] - The new tag name.
|
||||
* @param {string} [data.description] - The updated description.
|
||||
* @param {number} [data.position] - The new position.
|
||||
* @returns {Promise<Object>} The updated tag.
|
||||
*/
|
||||
const updateConversationTag = async (user, oldTag, data) => {
|
||||
try {
|
||||
const { tag: newTag, description, position } = data;
|
||||
|
||||
const existingTag = await ConversationTag.findOne({ user, tag: oldTag }).lean();
|
||||
if (!existingTag) {
|
||||
return null;
|
||||
}
|
||||
|
||||
if (newTag && newTag !== oldTag) {
|
||||
const tagAlreadyExists = await ConversationTag.findOne({ user, tag: newTag }).lean();
|
||||
if (tagAlreadyExists) {
|
||||
throw new Error('Tag already exists');
|
||||
}
|
||||
|
||||
await Conversation.updateMany({ user, tags: oldTag }, { $set: { 'tags.$': newTag } });
|
||||
}
|
||||
|
||||
const updateData = {};
|
||||
if (newTag) {
|
||||
updateData.tag = newTag;
|
||||
}
|
||||
if (description !== undefined) {
|
||||
updateData.description = description;
|
||||
}
|
||||
if (position !== undefined) {
|
||||
await adjustPositions(user, existingTag.position, position);
|
||||
updateData.position = position;
|
||||
}
|
||||
|
||||
return await ConversationTag.findOneAndUpdate({ user, tag: oldTag }, updateData, {
|
||||
new: true,
|
||||
lean: true,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error('[updateConversationTag] Error updating conversation tag', error);
|
||||
throw new Error('Error updating conversation tag');
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Adjusts positions of tags when a tag's position is changed.
|
||||
* @param {string} user - The user ID.
|
||||
* @param {number} oldPosition - The old position of the tag.
|
||||
* @param {number} newPosition - The new position of the tag.
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
const adjustPositions = async (user, oldPosition, newPosition) => {
|
||||
if (oldPosition === newPosition) {
|
||||
return;
|
||||
}
|
||||
|
||||
const update = oldPosition < newPosition ? { $inc: { position: -1 } } : { $inc: { position: 1 } };
|
||||
const position =
|
||||
oldPosition < newPosition
|
||||
? {
|
||||
$gt: Math.min(oldPosition, newPosition),
|
||||
$lte: Math.max(oldPosition, newPosition),
|
||||
}
|
||||
: {
|
||||
$gte: Math.min(oldPosition, newPosition),
|
||||
$lt: Math.max(oldPosition, newPosition),
|
||||
};
|
||||
|
||||
await ConversationTag.updateMany(
|
||||
{
|
||||
user,
|
||||
position,
|
||||
},
|
||||
update,
|
||||
);
|
||||
};
|
||||
|
||||
/**
|
||||
* Deletes a conversation tag.
|
||||
* @param {string} user - The user ID.
|
||||
* @param {string} tag - The tag to delete.
|
||||
* @returns {Promise<Object>} The deleted tag.
|
||||
*/
|
||||
const deleteConversationTag = async (user, tag) => {
|
||||
try {
|
||||
const deletedTag = await ConversationTag.findOneAndDelete({ user, tag }).lean();
|
||||
if (!deletedTag) {
|
||||
return null;
|
||||
}
|
||||
|
||||
await Conversation.updateMany({ user, tags: tag }, { $pull: { tags: tag } });
|
||||
|
||||
await ConversationTag.updateMany(
|
||||
{ user, position: { $gt: deletedTag.position } },
|
||||
{ $inc: { position: -1 } },
|
||||
);
|
||||
|
||||
return deletedTag;
|
||||
} catch (error) {
|
||||
logger.error('[deleteConversationTag] Error deleting conversation tag', error);
|
||||
throw new Error('Error deleting conversation tag');
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Updates tags for a specific conversation.
|
||||
* @param {string} user - The user ID.
|
||||
* @param {string} conversationId - The conversation ID.
|
||||
* @param {string[]} tags - The new set of tags for the conversation.
|
||||
* @returns {Promise<string[]>} The updated list of tags for the conversation.
|
||||
*/
|
||||
const updateTagsForConversation = async (user, conversationId, tags) => {
|
||||
try {
|
||||
const conversation = await Conversation.findOne({ user, conversationId }).lean();
|
||||
if (!conversation) {
|
||||
throw new Error('Conversation not found');
|
||||
}
|
||||
|
||||
const oldTags = new Set(conversation.tags);
|
||||
const newTags = new Set(tags);
|
||||
|
||||
const addedTags = [...newTags].filter((tag) => !oldTags.has(tag));
|
||||
const removedTags = [...oldTags].filter((tag) => !newTags.has(tag));
|
||||
|
||||
const bulkOps = [];
|
||||
|
||||
for (const tag of addedTags) {
|
||||
bulkOps.push({
|
||||
updateOne: {
|
||||
filter: { user, tag },
|
||||
update: { $inc: { count: 1 } },
|
||||
upsert: true,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
for (const tag of removedTags) {
|
||||
bulkOps.push({
|
||||
updateOne: {
|
||||
filter: { user, tag },
|
||||
update: { $inc: { count: -1 } },
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
if (bulkOps.length > 0) {
|
||||
await ConversationTag.bulkWrite(bulkOps);
|
||||
}
|
||||
|
||||
const updatedConversation = (
|
||||
await Conversation.findOneAndUpdate(
|
||||
{ user, conversationId },
|
||||
{ $set: { tags: [...newTags] } },
|
||||
{ new: true },
|
||||
)
|
||||
).toObject();
|
||||
|
||||
return updatedConversation.tags;
|
||||
} catch (error) {
|
||||
logger.error('[updateTagsForConversation] Error updating tags', error);
|
||||
throw new Error('Error updating tags for conversation');
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Increments tag counts for existing tags only.
|
||||
* @param {string} user - The user ID.
|
||||
* @param {string[]} tags - Array of tag names to increment
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
const bulkIncrementTagCounts = async (user, tags) => {
|
||||
if (!tags || tags.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const uniqueTags = [...new Set(tags.filter(Boolean))];
|
||||
if (uniqueTags.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const bulkOps = uniqueTags.map((tag) => ({
|
||||
updateOne: {
|
||||
filter: { user, tag },
|
||||
update: { $inc: { count: 1 } },
|
||||
},
|
||||
}));
|
||||
|
||||
const result = await ConversationTag.bulkWrite(bulkOps);
|
||||
if (result && result.modifiedCount > 0) {
|
||||
logger.debug(
|
||||
`user: ${user} | Incremented tag counts - modified ${result.modifiedCount} tags`,
|
||||
);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error('[bulkIncrementTagCounts] Error incrementing tag counts', error);
|
||||
}
|
||||
};
|
||||
|
||||
module.exports = {
|
||||
getConversationTags,
|
||||
createConversationTag,
|
||||
updateConversationTag,
|
||||
deleteConversationTag,
|
||||
bulkIncrementTagCounts,
|
||||
updateTagsForConversation,
|
||||
};
|
||||
179
api/models/File.js
Normal file
179
api/models/File.js
Normal file
@ -0,0 +1,179 @@
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { EToolResources, FileContext } = require('librechat-data-provider');
|
||||
const { File } = require('~/db/models');
|
||||
|
||||
/**
|
||||
* Finds a file by its file_id with additional query options.
|
||||
* @param {string} file_id - The unique identifier of the file.
|
||||
* @param {object} options - Query options for filtering, projection, etc.
|
||||
* @returns {Promise<MongoFile>} A promise that resolves to the file document or null.
|
||||
*/
|
||||
const findFileById = async (file_id, options = {}) => {
|
||||
return await File.findOne({ file_id, ...options }).lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Retrieves files matching a given filter, sorted by the most recently updated.
|
||||
* @param {Object} filter - The filter criteria to apply.
|
||||
* @param {Object} [_sortOptions] - Optional sort parameters.
|
||||
* @param {Object|String} [selectFields={ text: 0 }] - Fields to include/exclude in the query results.
|
||||
* Default excludes the 'text' field.
|
||||
* @returns {Promise<Array<MongoFile>>} A promise that resolves to an array of file documents.
|
||||
*/
|
||||
const getFiles = async (filter, _sortOptions, selectFields = { text: 0 }) => {
|
||||
const sortOptions = { updatedAt: -1, ..._sortOptions };
|
||||
return await File.find(filter).select(selectFields).sort(sortOptions).lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Retrieves tool files (files that are embedded or have a fileIdentifier) from an array of file IDs
|
||||
* @param {string[]} fileIds - Array of file_id strings to search for
|
||||
* @param {Set<EToolResources>} toolResourceSet - Optional filter for tool resources
|
||||
* @returns {Promise<Array<MongoFile>>} Files that match the criteria
|
||||
*/
|
||||
const getToolFilesByIds = async (fileIds, toolResourceSet) => {
|
||||
if (!fileIds || !fileIds.length || !toolResourceSet?.size) {
|
||||
return [];
|
||||
}
|
||||
|
||||
try {
|
||||
const filter = {
|
||||
file_id: { $in: fileIds },
|
||||
$or: [],
|
||||
};
|
||||
|
||||
if (toolResourceSet.has(EToolResources.context)) {
|
||||
filter.$or.push({ text: { $exists: true, $ne: null }, context: FileContext.agents });
|
||||
}
|
||||
if (toolResourceSet.has(EToolResources.file_search)) {
|
||||
filter.$or.push({ embedded: true });
|
||||
}
|
||||
if (toolResourceSet.has(EToolResources.execute_code)) {
|
||||
filter.$or.push({ 'metadata.fileIdentifier': { $exists: true } });
|
||||
}
|
||||
|
||||
const selectFields = { text: 0 };
|
||||
const sortOptions = { updatedAt: -1 };
|
||||
|
||||
return await getFiles(filter, sortOptions, selectFields);
|
||||
} catch (error) {
|
||||
logger.error('[getToolFilesByIds] Error retrieving tool files:', error);
|
||||
throw new Error('Error retrieving tool files');
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Creates a new file with a TTL of 1 hour.
|
||||
* @param {MongoFile} data - The file data to be created, must contain file_id.
|
||||
* @param {boolean} disableTTL - Whether to disable the TTL.
|
||||
* @returns {Promise<MongoFile>} A promise that resolves to the created file document.
|
||||
*/
|
||||
const createFile = async (data, disableTTL) => {
|
||||
const fileData = {
|
||||
...data,
|
||||
expiresAt: new Date(Date.now() + 3600 * 1000),
|
||||
};
|
||||
|
||||
if (disableTTL) {
|
||||
delete fileData.expiresAt;
|
||||
}
|
||||
|
||||
return await File.findOneAndUpdate({ file_id: data.file_id }, fileData, {
|
||||
new: true,
|
||||
upsert: true,
|
||||
}).lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Updates a file identified by file_id with new data and removes the TTL.
|
||||
* @param {MongoFile} data - The data to update, must contain file_id.
|
||||
* @returns {Promise<MongoFile>} A promise that resolves to the updated file document.
|
||||
*/
|
||||
const updateFile = async (data) => {
|
||||
const { file_id, ...update } = data;
|
||||
const updateOperation = {
|
||||
$set: update,
|
||||
$unset: { expiresAt: '' }, // Remove the expiresAt field to prevent TTL
|
||||
};
|
||||
return await File.findOneAndUpdate({ file_id }, updateOperation, { new: true }).lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Increments the usage of a file identified by file_id.
|
||||
* @param {MongoFile} data - The data to update, must contain file_id and the increment value for usage.
|
||||
* @returns {Promise<MongoFile>} A promise that resolves to the updated file document.
|
||||
*/
|
||||
const updateFileUsage = async (data) => {
|
||||
const { file_id, inc = 1 } = data;
|
||||
const updateOperation = {
|
||||
$inc: { usage: inc },
|
||||
$unset: { expiresAt: '', temp_file_id: '' },
|
||||
};
|
||||
return await File.findOneAndUpdate({ file_id }, updateOperation, { new: true }).lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Deletes a file identified by file_id.
|
||||
* @param {string} file_id - The unique identifier of the file to delete.
|
||||
* @returns {Promise<MongoFile>} A promise that resolves to the deleted file document or null.
|
||||
*/
|
||||
const deleteFile = async (file_id) => {
|
||||
return await File.findOneAndDelete({ file_id }).lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Deletes a file identified by a filter.
|
||||
* @param {object} filter - The filter criteria to apply.
|
||||
* @returns {Promise<MongoFile>} A promise that resolves to the deleted file document or null.
|
||||
*/
|
||||
const deleteFileByFilter = async (filter) => {
|
||||
return await File.findOneAndDelete(filter).lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Deletes multiple files identified by an array of file_ids.
|
||||
* @param {Array<string>} file_ids - The unique identifiers of the files to delete.
|
||||
* @returns {Promise<Object>} A promise that resolves to the result of the deletion operation.
|
||||
*/
|
||||
const deleteFiles = async (file_ids, user) => {
|
||||
let deleteQuery = { file_id: { $in: file_ids } };
|
||||
if (user) {
|
||||
deleteQuery = { user: user };
|
||||
}
|
||||
return await File.deleteMany(deleteQuery);
|
||||
};
|
||||
|
||||
/**
|
||||
* Batch updates files with new signed URLs in MongoDB
|
||||
*
|
||||
* @param {MongoFile[]} updates - Array of updates in the format { file_id, filepath }
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
async function batchUpdateFiles(updates) {
|
||||
if (!updates || updates.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const bulkOperations = updates.map((update) => ({
|
||||
updateOne: {
|
||||
filter: { file_id: update.file_id },
|
||||
update: { $set: { filepath: update.filepath } },
|
||||
},
|
||||
}));
|
||||
|
||||
const result = await File.bulkWrite(bulkOperations);
|
||||
logger.info(`Updated ${result.modifiedCount} files with new S3 URLs`);
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
findFileById,
|
||||
getFiles,
|
||||
getToolFilesByIds,
|
||||
createFile,
|
||||
updateFile,
|
||||
updateFileUsage,
|
||||
deleteFile,
|
||||
deleteFiles,
|
||||
deleteFileByFilter,
|
||||
batchUpdateFiles,
|
||||
};
|
||||
629
api/models/File.spec.js
Normal file
629
api/models/File.spec.js
Normal file
@ -0,0 +1,629 @@
|
||||
const mongoose = require('mongoose');
|
||||
const { v4: uuidv4 } = require('uuid');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
const { createModels, createMethods } = require('@librechat/data-schemas');
|
||||
const {
|
||||
SystemRoles,
|
||||
ResourceType,
|
||||
AccessRoleIds,
|
||||
PrincipalType,
|
||||
} = require('librechat-data-provider');
|
||||
const { grantPermission } = require('~/server/services/PermissionService');
|
||||
const { createAgent } = require('./Agent');
|
||||
|
||||
let File;
|
||||
let Agent;
|
||||
let AclEntry;
|
||||
let User;
|
||||
let modelsToCleanup = [];
|
||||
let methods;
|
||||
let getFiles;
|
||||
let createFile;
|
||||
let seedDefaultRoles;
|
||||
|
||||
describe('File Access Control', () => {
|
||||
let mongoServer;
|
||||
|
||||
beforeAll(async () => {
|
||||
mongoServer = await MongoMemoryServer.create();
|
||||
const mongoUri = mongoServer.getUri();
|
||||
await mongoose.connect(mongoUri);
|
||||
|
||||
// Initialize all models
|
||||
const models = createModels(mongoose);
|
||||
|
||||
// Track which models we're adding
|
||||
modelsToCleanup = Object.keys(models);
|
||||
|
||||
// Register models on mongoose.models so methods can access them
|
||||
const dbModels = require('~/db/models');
|
||||
Object.assign(mongoose.models, dbModels);
|
||||
|
||||
File = dbModels.File;
|
||||
Agent = dbModels.Agent;
|
||||
AclEntry = dbModels.AclEntry;
|
||||
User = dbModels.User;
|
||||
|
||||
// Create methods from data-schemas (includes file methods)
|
||||
methods = createMethods(mongoose);
|
||||
getFiles = methods.getFiles;
|
||||
createFile = methods.createFile;
|
||||
seedDefaultRoles = methods.seedDefaultRoles;
|
||||
|
||||
// Seed default roles
|
||||
await seedDefaultRoles();
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
// Clean up all collections before disconnecting
|
||||
const collections = mongoose.connection.collections;
|
||||
for (const key in collections) {
|
||||
await collections[key].deleteMany({});
|
||||
}
|
||||
|
||||
// Clear only the models we added
|
||||
for (const modelName of modelsToCleanup) {
|
||||
if (mongoose.models[modelName]) {
|
||||
delete mongoose.models[modelName];
|
||||
}
|
||||
}
|
||||
|
||||
await mongoose.disconnect();
|
||||
await mongoServer.stop();
|
||||
});
|
||||
|
||||
beforeEach(async () => {
|
||||
await File.deleteMany({});
|
||||
await Agent.deleteMany({});
|
||||
await AclEntry.deleteMany({});
|
||||
await User.deleteMany({});
|
||||
// Don't delete AccessRole as they are seeded defaults needed for tests
|
||||
});
|
||||
|
||||
describe('hasAccessToFilesViaAgent', () => {
|
||||
it('should efficiently check access for multiple files at once', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const agentId = uuidv4();
|
||||
const fileIds = [uuidv4(), uuidv4(), uuidv4(), uuidv4()];
|
||||
|
||||
// Create users
|
||||
await User.create({
|
||||
_id: userId,
|
||||
email: 'user@example.com',
|
||||
emailVerified: true,
|
||||
provider: 'local',
|
||||
});
|
||||
|
||||
await User.create({
|
||||
_id: authorId,
|
||||
email: 'author@example.com',
|
||||
emailVerified: true,
|
||||
provider: 'local',
|
||||
});
|
||||
|
||||
// Create files
|
||||
for (const fileId of fileIds) {
|
||||
await createFile({
|
||||
user: authorId,
|
||||
file_id: fileId,
|
||||
filename: `file-${fileId}.txt`,
|
||||
filepath: `/uploads/${fileId}`,
|
||||
});
|
||||
}
|
||||
|
||||
// Create agent with only first two files attached
|
||||
const agent = await createAgent({
|
||||
id: agentId,
|
||||
name: 'Test Agent',
|
||||
author: authorId,
|
||||
model: 'gpt-4',
|
||||
provider: 'openai',
|
||||
tool_resources: {
|
||||
file_search: {
|
||||
file_ids: [fileIds[0], fileIds[1]],
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// Grant EDIT permission to user on the agent
|
||||
await grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: userId,
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceId: agent._id,
|
||||
accessRoleId: AccessRoleIds.AGENT_EDITOR,
|
||||
grantedBy: authorId,
|
||||
});
|
||||
|
||||
// Check access for all files
|
||||
const { hasAccessToFilesViaAgent } = require('~/server/services/Files/permissions');
|
||||
const accessMap = await hasAccessToFilesViaAgent({
|
||||
userId: userId,
|
||||
role: SystemRoles.USER,
|
||||
fileIds,
|
||||
agentId: agent.id, // Use agent.id which is the custom UUID
|
||||
});
|
||||
|
||||
// Should have access only to the first two files
|
||||
expect(accessMap.get(fileIds[0])).toBe(true);
|
||||
expect(accessMap.get(fileIds[1])).toBe(true);
|
||||
expect(accessMap.get(fileIds[2])).toBe(false);
|
||||
expect(accessMap.get(fileIds[3])).toBe(false);
|
||||
});
|
||||
|
||||
it('should grant access to all files when user is the agent author', async () => {
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const agentId = uuidv4();
|
||||
const fileIds = [uuidv4(), uuidv4(), uuidv4()];
|
||||
|
||||
// Create author user
|
||||
await User.create({
|
||||
_id: authorId,
|
||||
email: 'author@example.com',
|
||||
emailVerified: true,
|
||||
provider: 'local',
|
||||
});
|
||||
|
||||
// Create agent
|
||||
await createAgent({
|
||||
id: agentId,
|
||||
name: 'Test Agent',
|
||||
author: authorId,
|
||||
model: 'gpt-4',
|
||||
provider: 'openai',
|
||||
tool_resources: {
|
||||
file_search: {
|
||||
file_ids: [fileIds[0]], // Only one file attached
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// Check access as the author
|
||||
const { hasAccessToFilesViaAgent } = require('~/server/services/Files/permissions');
|
||||
const accessMap = await hasAccessToFilesViaAgent({
|
||||
userId: authorId,
|
||||
role: SystemRoles.USER,
|
||||
fileIds,
|
||||
agentId,
|
||||
});
|
||||
|
||||
// Author should have access to all files
|
||||
expect(accessMap.get(fileIds[0])).toBe(true);
|
||||
expect(accessMap.get(fileIds[1])).toBe(true);
|
||||
expect(accessMap.get(fileIds[2])).toBe(true);
|
||||
});
|
||||
|
||||
it('should handle non-existent agent gracefully', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const fileIds = [uuidv4(), uuidv4()];
|
||||
|
||||
// Create user
|
||||
await User.create({
|
||||
_id: userId,
|
||||
email: 'user@example.com',
|
||||
emailVerified: true,
|
||||
provider: 'local',
|
||||
});
|
||||
|
||||
const { hasAccessToFilesViaAgent } = require('~/server/services/Files/permissions');
|
||||
const accessMap = await hasAccessToFilesViaAgent({
|
||||
userId: userId,
|
||||
role: SystemRoles.USER,
|
||||
fileIds,
|
||||
agentId: 'non-existent-agent',
|
||||
});
|
||||
|
||||
// Should have no access to any files
|
||||
expect(accessMap.get(fileIds[0])).toBe(false);
|
||||
expect(accessMap.get(fileIds[1])).toBe(false);
|
||||
});
|
||||
|
||||
it('should deny access when user only has VIEW permission and needs access for deletion', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const agentId = uuidv4();
|
||||
const fileIds = [uuidv4(), uuidv4()];
|
||||
|
||||
// Create users
|
||||
await User.create({
|
||||
_id: userId,
|
||||
email: 'user@example.com',
|
||||
emailVerified: true,
|
||||
provider: 'local',
|
||||
});
|
||||
|
||||
await User.create({
|
||||
_id: authorId,
|
||||
email: 'author@example.com',
|
||||
emailVerified: true,
|
||||
provider: 'local',
|
||||
});
|
||||
|
||||
// Create agent with files
|
||||
const agent = await createAgent({
|
||||
id: agentId,
|
||||
name: 'View-Only Agent',
|
||||
author: authorId,
|
||||
model: 'gpt-4',
|
||||
provider: 'openai',
|
||||
tool_resources: {
|
||||
file_search: {
|
||||
file_ids: fileIds,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// Grant only VIEW permission to user on the agent
|
||||
await grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: userId,
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceId: agent._id,
|
||||
accessRoleId: AccessRoleIds.AGENT_VIEWER,
|
||||
grantedBy: authorId,
|
||||
});
|
||||
|
||||
// Check access for files
|
||||
const { hasAccessToFilesViaAgent } = require('~/server/services/Files/permissions');
|
||||
const accessMap = await hasAccessToFilesViaAgent({
|
||||
userId: userId,
|
||||
role: SystemRoles.USER,
|
||||
fileIds,
|
||||
agentId,
|
||||
isDelete: true,
|
||||
});
|
||||
|
||||
// Should have no access to any files when only VIEW permission
|
||||
expect(accessMap.get(fileIds[0])).toBe(false);
|
||||
expect(accessMap.get(fileIds[1])).toBe(false);
|
||||
});
|
||||
|
||||
it('should grant access when user has VIEW permission', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const agentId = uuidv4();
|
||||
const fileIds = [uuidv4(), uuidv4()];
|
||||
|
||||
// Create users
|
||||
await User.create({
|
||||
_id: userId,
|
||||
email: 'user@example.com',
|
||||
emailVerified: true,
|
||||
provider: 'local',
|
||||
});
|
||||
|
||||
await User.create({
|
||||
_id: authorId,
|
||||
email: 'author@example.com',
|
||||
emailVerified: true,
|
||||
provider: 'local',
|
||||
});
|
||||
|
||||
// Create agent with files
|
||||
const agent = await createAgent({
|
||||
id: agentId,
|
||||
name: 'View-Only Agent',
|
||||
author: authorId,
|
||||
model: 'gpt-4',
|
||||
provider: 'openai',
|
||||
tool_resources: {
|
||||
file_search: {
|
||||
file_ids: fileIds,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// Grant only VIEW permission to user on the agent
|
||||
await grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: userId,
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceId: agent._id,
|
||||
accessRoleId: AccessRoleIds.AGENT_VIEWER,
|
||||
grantedBy: authorId,
|
||||
});
|
||||
|
||||
// Check access for files
|
||||
const { hasAccessToFilesViaAgent } = require('~/server/services/Files/permissions');
|
||||
const accessMap = await hasAccessToFilesViaAgent({
|
||||
userId: userId,
|
||||
role: SystemRoles.USER,
|
||||
fileIds,
|
||||
agentId,
|
||||
});
|
||||
|
||||
expect(accessMap.get(fileIds[0])).toBe(true);
|
||||
expect(accessMap.get(fileIds[1])).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getFiles with agent access control', () => {
|
||||
test('should return files owned by user and files accessible through agent', async () => {
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const agentId = `agent_${uuidv4()}`;
|
||||
const ownedFileId = `file_${uuidv4()}`;
|
||||
const sharedFileId = `file_${uuidv4()}`;
|
||||
const inaccessibleFileId = `file_${uuidv4()}`;
|
||||
|
||||
// Create users
|
||||
await User.create({
|
||||
_id: userId,
|
||||
email: 'user@example.com',
|
||||
emailVerified: true,
|
||||
provider: 'local',
|
||||
});
|
||||
|
||||
await User.create({
|
||||
_id: authorId,
|
||||
email: 'author@example.com',
|
||||
emailVerified: true,
|
||||
provider: 'local',
|
||||
});
|
||||
|
||||
// Create agent with shared file
|
||||
const agent = await createAgent({
|
||||
id: agentId,
|
||||
name: 'Shared Agent',
|
||||
provider: 'test',
|
||||
model: 'test-model',
|
||||
author: authorId,
|
||||
tool_resources: {
|
||||
file_search: {
|
||||
file_ids: [sharedFileId],
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// Grant EDIT permission to user on the agent
|
||||
await grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: userId,
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceId: agent._id,
|
||||
accessRoleId: AccessRoleIds.AGENT_EDITOR,
|
||||
grantedBy: authorId,
|
||||
});
|
||||
|
||||
// Create files
|
||||
await createFile({
|
||||
file_id: ownedFileId,
|
||||
user: userId,
|
||||
filename: 'owned.txt',
|
||||
filepath: '/uploads/owned.txt',
|
||||
type: 'text/plain',
|
||||
bytes: 100,
|
||||
});
|
||||
|
||||
await createFile({
|
||||
file_id: sharedFileId,
|
||||
user: authorId,
|
||||
filename: 'shared.txt',
|
||||
filepath: '/uploads/shared.txt',
|
||||
type: 'text/plain',
|
||||
bytes: 200,
|
||||
embedded: true,
|
||||
});
|
||||
|
||||
await createFile({
|
||||
file_id: inaccessibleFileId,
|
||||
user: authorId,
|
||||
filename: 'inaccessible.txt',
|
||||
filepath: '/uploads/inaccessible.txt',
|
||||
type: 'text/plain',
|
||||
bytes: 300,
|
||||
});
|
||||
|
||||
// Get all files first
|
||||
const allFiles = await getFiles(
|
||||
{ file_id: { $in: [ownedFileId, sharedFileId, inaccessibleFileId] } },
|
||||
null,
|
||||
{ text: 0 },
|
||||
);
|
||||
|
||||
// Then filter by access control
|
||||
const { filterFilesByAgentAccess } = require('~/server/services/Files/permissions');
|
||||
const files = await filterFilesByAgentAccess({
|
||||
files: allFiles,
|
||||
userId: userId,
|
||||
role: SystemRoles.USER,
|
||||
agentId,
|
||||
});
|
||||
|
||||
expect(files).toHaveLength(2);
|
||||
expect(files.map((f) => f.file_id)).toContain(ownedFileId);
|
||||
expect(files.map((f) => f.file_id)).toContain(sharedFileId);
|
||||
expect(files.map((f) => f.file_id)).not.toContain(inaccessibleFileId);
|
||||
});
|
||||
|
||||
test('should return all files when no userId/agentId provided', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const fileId1 = `file_${uuidv4()}`;
|
||||
const fileId2 = `file_${uuidv4()}`;
|
||||
|
||||
await createFile({
|
||||
file_id: fileId1,
|
||||
user: userId,
|
||||
filename: 'file1.txt',
|
||||
filepath: '/uploads/file1.txt',
|
||||
type: 'text/plain',
|
||||
bytes: 100,
|
||||
});
|
||||
|
||||
await createFile({
|
||||
file_id: fileId2,
|
||||
user: new mongoose.Types.ObjectId(),
|
||||
filename: 'file2.txt',
|
||||
filepath: '/uploads/file2.txt',
|
||||
type: 'text/plain',
|
||||
bytes: 200,
|
||||
});
|
||||
|
||||
const files = await getFiles({ file_id: { $in: [fileId1, fileId2] } });
|
||||
expect(files).toHaveLength(2);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Role-based file permissions', () => {
|
||||
it('should optimize permission checks when role is provided', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const agentId = uuidv4();
|
||||
const fileIds = [uuidv4(), uuidv4()];
|
||||
|
||||
// Create users
|
||||
await User.create({
|
||||
_id: userId,
|
||||
email: 'user@example.com',
|
||||
emailVerified: true,
|
||||
provider: 'local',
|
||||
role: 'ADMIN', // User has ADMIN role
|
||||
});
|
||||
|
||||
await User.create({
|
||||
_id: authorId,
|
||||
email: 'author@example.com',
|
||||
emailVerified: true,
|
||||
provider: 'local',
|
||||
});
|
||||
|
||||
// Create files
|
||||
for (const fileId of fileIds) {
|
||||
await createFile({
|
||||
file_id: fileId,
|
||||
user: authorId,
|
||||
filename: `${fileId}.txt`,
|
||||
filepath: `/uploads/${fileId}.txt`,
|
||||
type: 'text/plain',
|
||||
bytes: 100,
|
||||
});
|
||||
}
|
||||
|
||||
// Create agent with files
|
||||
const agent = await createAgent({
|
||||
id: agentId,
|
||||
name: 'Test Agent',
|
||||
author: authorId,
|
||||
model: 'gpt-4',
|
||||
provider: 'openai',
|
||||
tool_resources: {
|
||||
file_search: {
|
||||
file_ids: fileIds,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// Grant permission to ADMIN role
|
||||
await grantPermission({
|
||||
principalType: PrincipalType.ROLE,
|
||||
principalId: 'ADMIN',
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceId: agent._id,
|
||||
accessRoleId: AccessRoleIds.AGENT_EDITOR,
|
||||
grantedBy: authorId,
|
||||
});
|
||||
|
||||
// Check access with role provided (should avoid DB query)
|
||||
const { hasAccessToFilesViaAgent } = require('~/server/services/Files/permissions');
|
||||
const accessMapWithRole = await hasAccessToFilesViaAgent({
|
||||
userId: userId,
|
||||
role: 'ADMIN',
|
||||
fileIds,
|
||||
agentId: agent.id,
|
||||
});
|
||||
|
||||
// User should have access through their ADMIN role
|
||||
expect(accessMapWithRole.get(fileIds[0])).toBe(true);
|
||||
expect(accessMapWithRole.get(fileIds[1])).toBe(true);
|
||||
|
||||
// Check access without role (will query DB to get user's role)
|
||||
const accessMapWithoutRole = await hasAccessToFilesViaAgent({
|
||||
userId: userId,
|
||||
fileIds,
|
||||
agentId: agent.id,
|
||||
});
|
||||
|
||||
// Should have same result
|
||||
expect(accessMapWithoutRole.get(fileIds[0])).toBe(true);
|
||||
expect(accessMapWithoutRole.get(fileIds[1])).toBe(true);
|
||||
});
|
||||
|
||||
it('should deny access when user role changes', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const authorId = new mongoose.Types.ObjectId();
|
||||
const agentId = uuidv4();
|
||||
const fileId = uuidv4();
|
||||
|
||||
// Create users
|
||||
await User.create({
|
||||
_id: userId,
|
||||
email: 'user@example.com',
|
||||
emailVerified: true,
|
||||
provider: 'local',
|
||||
role: 'EDITOR',
|
||||
});
|
||||
|
||||
await User.create({
|
||||
_id: authorId,
|
||||
email: 'author@example.com',
|
||||
emailVerified: true,
|
||||
provider: 'local',
|
||||
});
|
||||
|
||||
// Create file
|
||||
await createFile({
|
||||
file_id: fileId,
|
||||
user: authorId,
|
||||
filename: 'test.txt',
|
||||
filepath: '/uploads/test.txt',
|
||||
type: 'text/plain',
|
||||
bytes: 100,
|
||||
});
|
||||
|
||||
// Create agent
|
||||
const agent = await createAgent({
|
||||
id: agentId,
|
||||
name: 'Test Agent',
|
||||
author: authorId,
|
||||
model: 'gpt-4',
|
||||
provider: 'openai',
|
||||
tool_resources: {
|
||||
file_search: {
|
||||
file_ids: [fileId],
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// Grant permission to EDITOR role only
|
||||
await grantPermission({
|
||||
principalType: PrincipalType.ROLE,
|
||||
principalId: 'EDITOR',
|
||||
resourceType: ResourceType.AGENT,
|
||||
resourceId: agent._id,
|
||||
accessRoleId: AccessRoleIds.AGENT_EDITOR,
|
||||
grantedBy: authorId,
|
||||
});
|
||||
|
||||
const { hasAccessToFilesViaAgent } = require('~/server/services/Files/permissions');
|
||||
|
||||
// Check with EDITOR role - should have access
|
||||
const accessAsEditor = await hasAccessToFilesViaAgent({
|
||||
userId: userId,
|
||||
role: 'EDITOR',
|
||||
fileIds: [fileId],
|
||||
agentId: agent.id,
|
||||
});
|
||||
expect(accessAsEditor.get(fileId)).toBe(true);
|
||||
|
||||
// Simulate role change to USER - should lose access
|
||||
const accessAsUser = await hasAccessToFilesViaAgent({
|
||||
userId: userId,
|
||||
role: SystemRoles.USER,
|
||||
fileIds: [fileId],
|
||||
agentId: agent.id,
|
||||
});
|
||||
expect(accessAsUser.get(fileId)).toBe(false);
|
||||
});
|
||||
});
|
||||
});
|
||||
372
api/models/Message.js
Normal file
372
api/models/Message.js
Normal file
@ -0,0 +1,372 @@
|
||||
const { z } = require('zod');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { createTempChatExpirationDate } = require('@librechat/api');
|
||||
const { Message } = require('~/db/models');
|
||||
|
||||
const idSchema = z.string().uuid();
|
||||
|
||||
/**
|
||||
* Saves a message in the database.
|
||||
*
|
||||
* @async
|
||||
* @function saveMessage
|
||||
* @param {ServerRequest} req - The request object containing user information.
|
||||
* @param {Object} params - The message data object.
|
||||
* @param {string} params.endpoint - The endpoint where the message originated.
|
||||
* @param {string} params.iconURL - The URL of the sender's icon.
|
||||
* @param {string} params.messageId - The unique identifier for the message.
|
||||
* @param {string} params.newMessageId - The new unique identifier for the message (if applicable).
|
||||
* @param {string} params.conversationId - The identifier of the conversation.
|
||||
* @param {string} [params.parentMessageId] - The identifier of the parent message, if any.
|
||||
* @param {string} params.sender - The identifier of the sender.
|
||||
* @param {string} params.text - The text content of the message.
|
||||
* @param {boolean} params.isCreatedByUser - Indicates if the message was created by the user.
|
||||
* @param {string} [params.error] - Any error associated with the message.
|
||||
* @param {boolean} [params.unfinished] - Indicates if the message is unfinished.
|
||||
* @param {Object[]} [params.files] - An array of files associated with the message.
|
||||
* @param {string} [params.finish_reason] - Reason for finishing the message.
|
||||
* @param {number} [params.tokenCount] - The number of tokens in the message.
|
||||
* @param {string} [params.plugin] - Plugin associated with the message.
|
||||
* @param {string[]} [params.plugins] - An array of plugins associated with the message.
|
||||
* @param {string} [params.model] - The model used to generate the message.
|
||||
* @param {Object} [metadata] - Additional metadata for this operation
|
||||
* @param {string} [metadata.context] - The context of the operation
|
||||
* @returns {Promise<TMessage>} The updated or newly inserted message document.
|
||||
* @throws {Error} If there is an error in saving the message.
|
||||
*/
|
||||
async function saveMessage(req, params, metadata) {
|
||||
if (!req?.user?.id) {
|
||||
throw new Error('User not authenticated');
|
||||
}
|
||||
|
||||
const validConvoId = idSchema.safeParse(params.conversationId);
|
||||
if (!validConvoId.success) {
|
||||
logger.warn(`Invalid conversation ID: ${params.conversationId}`);
|
||||
logger.info(`---\`saveMessage\` context: ${metadata?.context}`);
|
||||
logger.info(`---Invalid conversation ID Params: ${JSON.stringify(params, null, 2)}`);
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const update = {
|
||||
...params,
|
||||
user: req.user.id,
|
||||
messageId: params.newMessageId || params.messageId,
|
||||
};
|
||||
|
||||
if (req?.body?.isTemporary) {
|
||||
try {
|
||||
const appConfig = req.config;
|
||||
update.expiredAt = createTempChatExpirationDate(appConfig?.interfaceConfig);
|
||||
} catch (err) {
|
||||
logger.error('Error creating temporary chat expiration date:', err);
|
||||
logger.info(`---\`saveMessage\` context: ${metadata?.context}`);
|
||||
update.expiredAt = null;
|
||||
}
|
||||
} else {
|
||||
update.expiredAt = null;
|
||||
}
|
||||
|
||||
if (update.tokenCount != null && isNaN(update.tokenCount)) {
|
||||
logger.warn(
|
||||
`Resetting invalid \`tokenCount\` for message \`${params.messageId}\`: ${update.tokenCount}`,
|
||||
);
|
||||
logger.info(`---\`saveMessage\` context: ${metadata?.context}`);
|
||||
update.tokenCount = 0;
|
||||
}
|
||||
const message = await Message.findOneAndUpdate(
|
||||
{ messageId: params.messageId, user: req.user.id },
|
||||
update,
|
||||
{ upsert: true, new: true },
|
||||
);
|
||||
|
||||
return message.toObject();
|
||||
} catch (err) {
|
||||
logger.error('Error saving message:', err);
|
||||
logger.info(`---\`saveMessage\` context: ${metadata?.context}`);
|
||||
|
||||
// Check if this is a duplicate key error (MongoDB error code 11000)
|
||||
if (err.code === 11000 && err.message.includes('duplicate key error')) {
|
||||
// Log the duplicate key error but don't crash the application
|
||||
logger.warn(`Duplicate messageId detected: ${params.messageId}. Continuing execution.`);
|
||||
|
||||
try {
|
||||
// Try to find the existing message with this ID
|
||||
const existingMessage = await Message.findOne({
|
||||
messageId: params.messageId,
|
||||
user: req.user.id,
|
||||
});
|
||||
|
||||
// If we found it, return it
|
||||
if (existingMessage) {
|
||||
return existingMessage.toObject();
|
||||
}
|
||||
|
||||
// If we can't find it (unlikely but possible in race conditions)
|
||||
return {
|
||||
...params,
|
||||
messageId: params.messageId,
|
||||
user: req.user.id,
|
||||
};
|
||||
} catch (findError) {
|
||||
// If the findOne also fails, log it but don't crash
|
||||
logger.warn(
|
||||
`Could not retrieve existing message with ID ${params.messageId}: ${findError.message}`,
|
||||
);
|
||||
return {
|
||||
...params,
|
||||
messageId: params.messageId,
|
||||
user: req.user.id,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
throw err; // Re-throw other errors
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Saves multiple messages in the database in bulk.
|
||||
*
|
||||
* @async
|
||||
* @function bulkSaveMessages
|
||||
* @param {Object[]} messages - An array of message objects to save.
|
||||
* @param {boolean} [overrideTimestamp=false] - Indicates whether to override the timestamps of the messages. Defaults to false.
|
||||
* @returns {Promise<Object>} The result of the bulk write operation.
|
||||
* @throws {Error} If there is an error in saving messages in bulk.
|
||||
*/
|
||||
async function bulkSaveMessages(messages, overrideTimestamp = false) {
|
||||
try {
|
||||
const bulkOps = messages.map((message) => ({
|
||||
updateOne: {
|
||||
filter: { messageId: message.messageId },
|
||||
update: message,
|
||||
timestamps: !overrideTimestamp,
|
||||
upsert: true,
|
||||
},
|
||||
}));
|
||||
const result = await Message.bulkWrite(bulkOps);
|
||||
return result;
|
||||
} catch (err) {
|
||||
logger.error('Error saving messages in bulk:', err);
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Records a message in the database.
|
||||
*
|
||||
* @async
|
||||
* @function recordMessage
|
||||
* @param {Object} params - The message data object.
|
||||
* @param {string} params.user - The identifier of the user.
|
||||
* @param {string} params.endpoint - The endpoint where the message originated.
|
||||
* @param {string} params.messageId - The unique identifier for the message.
|
||||
* @param {string} params.conversationId - The identifier of the conversation.
|
||||
* @param {string} [params.parentMessageId] - The identifier of the parent message, if any.
|
||||
* @param {Partial<TMessage>} rest - Any additional properties from the TMessage typedef not explicitly listed.
|
||||
* @returns {Promise<Object>} The updated or newly inserted message document.
|
||||
* @throws {Error} If there is an error in saving the message.
|
||||
*/
|
||||
async function recordMessage({
|
||||
user,
|
||||
endpoint,
|
||||
messageId,
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
...rest
|
||||
}) {
|
||||
try {
|
||||
// No parsing of convoId as may use threadId
|
||||
const message = {
|
||||
user,
|
||||
endpoint,
|
||||
messageId,
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
...rest,
|
||||
};
|
||||
|
||||
return await Message.findOneAndUpdate({ user, messageId }, message, {
|
||||
upsert: true,
|
||||
new: true,
|
||||
});
|
||||
} catch (err) {
|
||||
logger.error('Error recording message:', err);
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Updates the text of a message.
|
||||
*
|
||||
* @async
|
||||
* @function updateMessageText
|
||||
* @param {Object} params - The update data object.
|
||||
* @param {Object} req - The request object.
|
||||
* @param {string} params.messageId - The unique identifier for the message.
|
||||
* @param {string} params.text - The new text content of the message.
|
||||
* @returns {Promise<void>}
|
||||
* @throws {Error} If there is an error in updating the message text.
|
||||
*/
|
||||
async function updateMessageText(req, { messageId, text }) {
|
||||
try {
|
||||
await Message.updateOne({ messageId, user: req.user.id }, { text });
|
||||
} catch (err) {
|
||||
logger.error('Error updating message text:', err);
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Updates a message.
|
||||
*
|
||||
* @async
|
||||
* @function updateMessage
|
||||
* @param {Object} req - The request object.
|
||||
* @param {Object} message - The message object containing update data.
|
||||
* @param {string} message.messageId - The unique identifier for the message.
|
||||
* @param {string} [message.text] - The new text content of the message.
|
||||
* @param {Object[]} [message.files] - The files associated with the message.
|
||||
* @param {boolean} [message.isCreatedByUser] - Indicates if the message was created by the user.
|
||||
* @param {string} [message.sender] - The identifier of the sender.
|
||||
* @param {number} [message.tokenCount] - The number of tokens in the message.
|
||||
* @param {Object} [metadata] - The operation metadata
|
||||
* @param {string} [metadata.context] - The operation metadata
|
||||
* @returns {Promise<TMessage>} The updated message document.
|
||||
* @throws {Error} If there is an error in updating the message or if the message is not found.
|
||||
*/
|
||||
async function updateMessage(req, message, metadata) {
|
||||
try {
|
||||
const { messageId, ...update } = message;
|
||||
const updatedMessage = await Message.findOneAndUpdate(
|
||||
{ messageId, user: req.user.id },
|
||||
update,
|
||||
{
|
||||
new: true,
|
||||
},
|
||||
);
|
||||
|
||||
if (!updatedMessage) {
|
||||
throw new Error('Message not found or user not authorized.');
|
||||
}
|
||||
|
||||
return {
|
||||
messageId: updatedMessage.messageId,
|
||||
conversationId: updatedMessage.conversationId,
|
||||
parentMessageId: updatedMessage.parentMessageId,
|
||||
sender: updatedMessage.sender,
|
||||
text: updatedMessage.text,
|
||||
isCreatedByUser: updatedMessage.isCreatedByUser,
|
||||
tokenCount: updatedMessage.tokenCount,
|
||||
feedback: updatedMessage.feedback,
|
||||
};
|
||||
} catch (err) {
|
||||
logger.error('Error updating message:', err);
|
||||
if (metadata && metadata?.context) {
|
||||
logger.info(`---\`updateMessage\` context: ${metadata.context}`);
|
||||
}
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Deletes messages in a conversation since a specific message.
|
||||
*
|
||||
* @async
|
||||
* @function deleteMessagesSince
|
||||
* @param {Object} params - The parameters object.
|
||||
* @param {Object} req - The request object.
|
||||
* @param {string} params.messageId - The unique identifier for the message.
|
||||
* @param {string} params.conversationId - The identifier of the conversation.
|
||||
* @returns {Promise<Number>} The number of deleted messages.
|
||||
* @throws {Error} If there is an error in deleting messages.
|
||||
*/
|
||||
async function deleteMessagesSince(req, { messageId, conversationId }) {
|
||||
try {
|
||||
const message = await Message.findOne({ messageId, user: req.user.id }).lean();
|
||||
|
||||
if (message) {
|
||||
const query = Message.find({ conversationId, user: req.user.id });
|
||||
return await query.deleteMany({
|
||||
createdAt: { $gt: message.createdAt },
|
||||
});
|
||||
}
|
||||
return undefined;
|
||||
} catch (err) {
|
||||
logger.error('Error deleting messages:', err);
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieves messages from the database.
|
||||
* @async
|
||||
* @function getMessages
|
||||
* @param {Record<string, unknown>} filter - The filter criteria.
|
||||
* @param {string | undefined} [select] - The fields to select.
|
||||
* @returns {Promise<TMessage[]>} The messages that match the filter criteria.
|
||||
* @throws {Error} If there is an error in retrieving messages.
|
||||
*/
|
||||
async function getMessages(filter, select) {
|
||||
try {
|
||||
if (select) {
|
||||
return await Message.find(filter).select(select).sort({ createdAt: 1 }).lean();
|
||||
}
|
||||
|
||||
return await Message.find(filter).sort({ createdAt: 1 }).lean();
|
||||
} catch (err) {
|
||||
logger.error('Error getting messages:', err);
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieves a single message from the database.
|
||||
* @async
|
||||
* @function getMessage
|
||||
* @param {{ user: string, messageId: string }} params - The search parameters
|
||||
* @returns {Promise<TMessage | null>} The message that matches the criteria or null if not found
|
||||
* @throws {Error} If there is an error in retrieving the message
|
||||
*/
|
||||
async function getMessage({ user, messageId }) {
|
||||
try {
|
||||
return await Message.findOne({
|
||||
user,
|
||||
messageId,
|
||||
}).lean();
|
||||
} catch (err) {
|
||||
logger.error('Error getting message:', err);
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Deletes messages from the database.
|
||||
*
|
||||
* @async
|
||||
* @function deleteMessages
|
||||
* @param {import('mongoose').FilterQuery<import('mongoose').Document>} filter - The filter criteria to find messages to delete.
|
||||
* @returns {Promise<import('mongoose').DeleteResult>} The metadata with count of deleted messages.
|
||||
* @throws {Error} If there is an error in deleting messages.
|
||||
*/
|
||||
async function deleteMessages(filter) {
|
||||
try {
|
||||
return await Message.deleteMany(filter);
|
||||
} catch (err) {
|
||||
logger.error('Error deleting messages:', err);
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
saveMessage,
|
||||
bulkSaveMessages,
|
||||
recordMessage,
|
||||
updateMessageText,
|
||||
updateMessage,
|
||||
deleteMessagesSince,
|
||||
getMessages,
|
||||
getMessage,
|
||||
deleteMessages,
|
||||
};
|
||||
898
api/models/Message.spec.js
Normal file
898
api/models/Message.spec.js
Normal file
@ -0,0 +1,898 @@
|
||||
const mongoose = require('mongoose');
|
||||
const { v4: uuidv4 } = require('uuid');
|
||||
const { messageSchema } = require('@librechat/data-schemas');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
|
||||
const {
|
||||
saveMessage,
|
||||
getMessages,
|
||||
updateMessage,
|
||||
deleteMessages,
|
||||
bulkSaveMessages,
|
||||
updateMessageText,
|
||||
deleteMessagesSince,
|
||||
} = require('./Message');
|
||||
|
||||
jest.mock('~/server/services/Config/app');
|
||||
|
||||
/**
|
||||
* @type {import('mongoose').Model<import('@librechat/data-schemas').IMessage>}
|
||||
*/
|
||||
let Message;
|
||||
|
||||
describe('Message Operations', () => {
|
||||
let mongoServer;
|
||||
let mockReq;
|
||||
let mockMessageData;
|
||||
|
||||
beforeAll(async () => {
|
||||
mongoServer = await MongoMemoryServer.create();
|
||||
const mongoUri = mongoServer.getUri();
|
||||
Message = mongoose.models.Message || mongoose.model('Message', messageSchema);
|
||||
await mongoose.connect(mongoUri);
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await mongoose.disconnect();
|
||||
await mongoServer.stop();
|
||||
});
|
||||
|
||||
beforeEach(async () => {
|
||||
// Clear database
|
||||
await Message.deleteMany({});
|
||||
|
||||
mockReq = {
|
||||
user: { id: 'user123' },
|
||||
config: {
|
||||
interfaceConfig: {
|
||||
temporaryChatRetention: 24, // Default 24 hours
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
mockMessageData = {
|
||||
messageId: 'msg123',
|
||||
conversationId: uuidv4(),
|
||||
text: 'Hello, world!',
|
||||
user: 'user123',
|
||||
};
|
||||
});
|
||||
|
||||
describe('saveMessage', () => {
|
||||
it('should save a message for an authenticated user', async () => {
|
||||
const result = await saveMessage(mockReq, mockMessageData);
|
||||
|
||||
expect(result.messageId).toBe('msg123');
|
||||
expect(result.user).toBe('user123');
|
||||
expect(result.text).toBe('Hello, world!');
|
||||
|
||||
// Verify the message was actually saved to the database
|
||||
const savedMessage = await Message.findOne({ messageId: 'msg123', user: 'user123' });
|
||||
expect(savedMessage).toBeTruthy();
|
||||
expect(savedMessage.text).toBe('Hello, world!');
|
||||
});
|
||||
|
||||
it('should throw an error for unauthenticated user', async () => {
|
||||
mockReq.user = null;
|
||||
await expect(saveMessage(mockReq, mockMessageData)).rejects.toThrow('User not authenticated');
|
||||
});
|
||||
|
||||
it('should handle invalid conversation ID gracefully', async () => {
|
||||
mockMessageData.conversationId = 'invalid-id';
|
||||
const result = await saveMessage(mockReq, mockMessageData);
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('updateMessageText', () => {
|
||||
it('should update message text for the authenticated user', async () => {
|
||||
// First save a message
|
||||
await saveMessage(mockReq, mockMessageData);
|
||||
|
||||
// Then update it
|
||||
await updateMessageText(mockReq, { messageId: 'msg123', text: 'Updated text' });
|
||||
|
||||
// Verify the update
|
||||
const updatedMessage = await Message.findOne({ messageId: 'msg123', user: 'user123' });
|
||||
expect(updatedMessage.text).toBe('Updated text');
|
||||
});
|
||||
});
|
||||
|
||||
describe('updateMessage', () => {
|
||||
it('should update a message for the authenticated user', async () => {
|
||||
// First save a message
|
||||
await saveMessage(mockReq, mockMessageData);
|
||||
|
||||
const result = await updateMessage(mockReq, { messageId: 'msg123', text: 'Updated text' });
|
||||
|
||||
expect(result.messageId).toBe('msg123');
|
||||
expect(result.text).toBe('Updated text');
|
||||
|
||||
// Verify in database
|
||||
const updatedMessage = await Message.findOne({ messageId: 'msg123', user: 'user123' });
|
||||
expect(updatedMessage.text).toBe('Updated text');
|
||||
});
|
||||
|
||||
it('should throw an error if message is not found', async () => {
|
||||
await expect(
|
||||
updateMessage(mockReq, { messageId: 'nonexistent', text: 'Test' }),
|
||||
).rejects.toThrow('Message not found or user not authorized.');
|
||||
});
|
||||
});
|
||||
|
||||
describe('deleteMessagesSince', () => {
|
||||
it('should delete messages only for the authenticated user', async () => {
|
||||
const conversationId = uuidv4();
|
||||
|
||||
// Create multiple messages in the same conversation
|
||||
await saveMessage(mockReq, {
|
||||
messageId: 'msg1',
|
||||
conversationId,
|
||||
text: 'First message',
|
||||
user: 'user123',
|
||||
});
|
||||
|
||||
await saveMessage(mockReq, {
|
||||
messageId: 'msg2',
|
||||
conversationId,
|
||||
text: 'Second message',
|
||||
user: 'user123',
|
||||
});
|
||||
|
||||
await saveMessage(mockReq, {
|
||||
messageId: 'msg3',
|
||||
conversationId,
|
||||
text: 'Third message',
|
||||
user: 'user123',
|
||||
});
|
||||
|
||||
// Delete messages since message2 (this should only delete messages created AFTER msg2)
|
||||
await deleteMessagesSince(mockReq, {
|
||||
messageId: 'msg2',
|
||||
conversationId,
|
||||
});
|
||||
|
||||
// Verify msg1 and msg2 remain, msg3 is deleted
|
||||
const remainingMessages = await Message.find({ conversationId, user: 'user123' });
|
||||
expect(remainingMessages).toHaveLength(2);
|
||||
expect(remainingMessages.map((m) => m.messageId)).toContain('msg1');
|
||||
expect(remainingMessages.map((m) => m.messageId)).toContain('msg2');
|
||||
expect(remainingMessages.map((m) => m.messageId)).not.toContain('msg3');
|
||||
});
|
||||
|
||||
it('should return undefined if no message is found', async () => {
|
||||
const result = await deleteMessagesSince(mockReq, {
|
||||
messageId: 'nonexistent',
|
||||
conversationId: 'convo123',
|
||||
});
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('getMessages', () => {
|
||||
it('should retrieve messages with the correct filter', async () => {
|
||||
const conversationId = uuidv4();
|
||||
|
||||
// Save some messages
|
||||
await saveMessage(mockReq, {
|
||||
messageId: 'msg1',
|
||||
conversationId,
|
||||
text: 'First message',
|
||||
user: 'user123',
|
||||
});
|
||||
|
||||
await saveMessage(mockReq, {
|
||||
messageId: 'msg2',
|
||||
conversationId,
|
||||
text: 'Second message',
|
||||
user: 'user123',
|
||||
});
|
||||
|
||||
const messages = await getMessages({ conversationId });
|
||||
expect(messages).toHaveLength(2);
|
||||
expect(messages[0].text).toBe('First message');
|
||||
expect(messages[1].text).toBe('Second message');
|
||||
});
|
||||
});
|
||||
|
||||
describe('deleteMessages', () => {
|
||||
it('should delete messages with the correct filter', async () => {
|
||||
// Save some messages for different users
|
||||
await saveMessage(mockReq, mockMessageData);
|
||||
await saveMessage(
|
||||
{ user: { id: 'user456' } },
|
||||
{
|
||||
messageId: 'msg456',
|
||||
conversationId: uuidv4(),
|
||||
text: 'Other user message',
|
||||
user: 'user456',
|
||||
},
|
||||
);
|
||||
|
||||
await deleteMessages({ user: 'user123' });
|
||||
|
||||
// Verify only user123's messages were deleted
|
||||
const user123Messages = await Message.find({ user: 'user123' });
|
||||
const user456Messages = await Message.find({ user: 'user456' });
|
||||
|
||||
expect(user123Messages).toHaveLength(0);
|
||||
expect(user456Messages).toHaveLength(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Conversation Hijacking Prevention', () => {
|
||||
it("should not allow editing a message in another user's conversation", async () => {
|
||||
const attackerReq = { user: { id: 'attacker123' } };
|
||||
const victimConversationId = uuidv4();
|
||||
const victimMessageId = 'victim-msg-123';
|
||||
|
||||
// First, save a message as the victim (but we'll try to edit as attacker)
|
||||
const victimReq = { user: { id: 'victim123' } };
|
||||
await saveMessage(victimReq, {
|
||||
messageId: victimMessageId,
|
||||
conversationId: victimConversationId,
|
||||
text: 'Victim message',
|
||||
user: 'victim123',
|
||||
});
|
||||
|
||||
// Attacker tries to edit the victim's message
|
||||
await expect(
|
||||
updateMessage(attackerReq, {
|
||||
messageId: victimMessageId,
|
||||
conversationId: victimConversationId,
|
||||
text: 'Hacked message',
|
||||
}),
|
||||
).rejects.toThrow('Message not found or user not authorized.');
|
||||
|
||||
// Verify the original message is unchanged
|
||||
const originalMessage = await Message.findOne({
|
||||
messageId: victimMessageId,
|
||||
user: 'victim123',
|
||||
});
|
||||
expect(originalMessage.text).toBe('Victim message');
|
||||
});
|
||||
|
||||
it("should not allow deleting messages from another user's conversation", async () => {
|
||||
const attackerReq = { user: { id: 'attacker123' } };
|
||||
const victimConversationId = uuidv4();
|
||||
const victimMessageId = 'victim-msg-123';
|
||||
|
||||
// Save a message as the victim
|
||||
const victimReq = { user: { id: 'victim123' } };
|
||||
await saveMessage(victimReq, {
|
||||
messageId: victimMessageId,
|
||||
conversationId: victimConversationId,
|
||||
text: 'Victim message',
|
||||
user: 'victim123',
|
||||
});
|
||||
|
||||
// Attacker tries to delete from victim's conversation
|
||||
const result = await deleteMessagesSince(attackerReq, {
|
||||
messageId: victimMessageId,
|
||||
conversationId: victimConversationId,
|
||||
});
|
||||
|
||||
expect(result).toBeUndefined();
|
||||
|
||||
// Verify the victim's message still exists
|
||||
const victimMessage = await Message.findOne({
|
||||
messageId: victimMessageId,
|
||||
user: 'victim123',
|
||||
});
|
||||
expect(victimMessage).toBeTruthy();
|
||||
expect(victimMessage.text).toBe('Victim message');
|
||||
});
|
||||
|
||||
it("should not allow inserting a new message into another user's conversation", async () => {
|
||||
const attackerReq = { user: { id: 'attacker123' } };
|
||||
const victimConversationId = uuidv4();
|
||||
|
||||
// Attacker tries to save a message - this should succeed but with attacker's user ID
|
||||
const result = await saveMessage(attackerReq, {
|
||||
conversationId: victimConversationId,
|
||||
text: 'Inserted malicious message',
|
||||
messageId: 'new-msg-123',
|
||||
user: 'attacker123',
|
||||
});
|
||||
|
||||
expect(result).toBeTruthy();
|
||||
expect(result.user).toBe('attacker123');
|
||||
|
||||
// Verify the message was saved with the attacker's user ID, not as an anonymous message
|
||||
const savedMessage = await Message.findOne({ messageId: 'new-msg-123' });
|
||||
expect(savedMessage.user).toBe('attacker123');
|
||||
expect(savedMessage.conversationId).toBe(victimConversationId);
|
||||
});
|
||||
|
||||
it('should allow retrieving messages from any conversation', async () => {
|
||||
const victimConversationId = uuidv4();
|
||||
|
||||
// Save a message in the victim's conversation
|
||||
const victimReq = { user: { id: 'victim123' } };
|
||||
await saveMessage(victimReq, {
|
||||
messageId: 'victim-msg',
|
||||
conversationId: victimConversationId,
|
||||
text: 'Victim message',
|
||||
user: 'victim123',
|
||||
});
|
||||
|
||||
// Anyone should be able to retrieve messages by conversation ID
|
||||
const messages = await getMessages({ conversationId: victimConversationId });
|
||||
expect(messages).toHaveLength(1);
|
||||
expect(messages[0].text).toBe('Victim message');
|
||||
});
|
||||
});
|
||||
|
||||
describe('isTemporary message handling', () => {
|
||||
beforeEach(() => {
|
||||
// Reset mocks before each test
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
it('should save a message with expiredAt when isTemporary is true', async () => {
|
||||
// Mock app config with 24 hour retention
|
||||
mockReq.config.interfaceConfig.temporaryChatRetention = 24;
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
|
||||
const beforeSave = new Date();
|
||||
const result = await saveMessage(mockReq, mockMessageData);
|
||||
const afterSave = new Date();
|
||||
|
||||
expect(result.messageId).toBe('msg123');
|
||||
expect(result.expiredAt).toBeDefined();
|
||||
expect(result.expiredAt).toBeInstanceOf(Date);
|
||||
|
||||
// Verify expiredAt is approximately 24 hours in the future
|
||||
const expectedExpirationTime = new Date(beforeSave.getTime() + 24 * 60 * 60 * 1000);
|
||||
const actualExpirationTime = new Date(result.expiredAt);
|
||||
|
||||
expect(actualExpirationTime.getTime()).toBeGreaterThanOrEqual(
|
||||
expectedExpirationTime.getTime() - 1000,
|
||||
);
|
||||
expect(actualExpirationTime.getTime()).toBeLessThanOrEqual(
|
||||
new Date(afterSave.getTime() + 24 * 60 * 60 * 1000 + 1000).getTime(),
|
||||
);
|
||||
});
|
||||
|
||||
it('should save a message without expiredAt when isTemporary is false', async () => {
|
||||
mockReq.body = { isTemporary: false };
|
||||
|
||||
const result = await saveMessage(mockReq, mockMessageData);
|
||||
|
||||
expect(result.messageId).toBe('msg123');
|
||||
expect(result.expiredAt).toBeNull();
|
||||
});
|
||||
|
||||
it('should save a message without expiredAt when isTemporary is not provided', async () => {
|
||||
// No isTemporary in body
|
||||
mockReq.body = {};
|
||||
|
||||
const result = await saveMessage(mockReq, mockMessageData);
|
||||
|
||||
expect(result.messageId).toBe('msg123');
|
||||
expect(result.expiredAt).toBeNull();
|
||||
});
|
||||
|
||||
it('should use custom retention period from config', async () => {
|
||||
// Mock app config with 48 hour retention
|
||||
mockReq.config.interfaceConfig.temporaryChatRetention = 48;
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
|
||||
const beforeSave = new Date();
|
||||
const result = await saveMessage(mockReq, mockMessageData);
|
||||
|
||||
expect(result.expiredAt).toBeDefined();
|
||||
|
||||
// Verify expiredAt is approximately 48 hours in the future
|
||||
const expectedExpirationTime = new Date(beforeSave.getTime() + 48 * 60 * 60 * 1000);
|
||||
const actualExpirationTime = new Date(result.expiredAt);
|
||||
|
||||
expect(actualExpirationTime.getTime()).toBeGreaterThanOrEqual(
|
||||
expectedExpirationTime.getTime() - 1000,
|
||||
);
|
||||
expect(actualExpirationTime.getTime()).toBeLessThanOrEqual(
|
||||
expectedExpirationTime.getTime() + 1000,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle minimum retention period (1 hour)', async () => {
|
||||
// Mock app config with less than minimum retention
|
||||
mockReq.config.interfaceConfig.temporaryChatRetention = 0.5; // Half hour - should be clamped to 1 hour
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
|
||||
const beforeSave = new Date();
|
||||
const result = await saveMessage(mockReq, mockMessageData);
|
||||
|
||||
expect(result.expiredAt).toBeDefined();
|
||||
|
||||
// Verify expiredAt is approximately 1 hour in the future (minimum)
|
||||
const expectedExpirationTime = new Date(beforeSave.getTime() + 1 * 60 * 60 * 1000);
|
||||
const actualExpirationTime = new Date(result.expiredAt);
|
||||
|
||||
expect(actualExpirationTime.getTime()).toBeGreaterThanOrEqual(
|
||||
expectedExpirationTime.getTime() - 1000,
|
||||
);
|
||||
expect(actualExpirationTime.getTime()).toBeLessThanOrEqual(
|
||||
expectedExpirationTime.getTime() + 1000,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle maximum retention period (8760 hours)', async () => {
|
||||
// Mock app config with more than maximum retention
|
||||
mockReq.config.interfaceConfig.temporaryChatRetention = 10000; // Should be clamped to 8760 hours
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
|
||||
const beforeSave = new Date();
|
||||
const result = await saveMessage(mockReq, mockMessageData);
|
||||
|
||||
expect(result.expiredAt).toBeDefined();
|
||||
|
||||
// Verify expiredAt is approximately 8760 hours (1 year) in the future
|
||||
const expectedExpirationTime = new Date(beforeSave.getTime() + 8760 * 60 * 60 * 1000);
|
||||
const actualExpirationTime = new Date(result.expiredAt);
|
||||
|
||||
expect(actualExpirationTime.getTime()).toBeGreaterThanOrEqual(
|
||||
expectedExpirationTime.getTime() - 1000,
|
||||
);
|
||||
expect(actualExpirationTime.getTime()).toBeLessThanOrEqual(
|
||||
expectedExpirationTime.getTime() + 1000,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle missing config gracefully', async () => {
|
||||
// Simulate missing config - should use default retention period
|
||||
delete mockReq.config;
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
|
||||
const beforeSave = new Date();
|
||||
const result = await saveMessage(mockReq, mockMessageData);
|
||||
const afterSave = new Date();
|
||||
|
||||
// Should still save the message with default retention period (30 days)
|
||||
expect(result.messageId).toBe('msg123');
|
||||
expect(result.expiredAt).toBeDefined();
|
||||
expect(result.expiredAt).toBeInstanceOf(Date);
|
||||
|
||||
// Verify expiredAt is approximately 30 days in the future (720 hours)
|
||||
const expectedExpirationTime = new Date(beforeSave.getTime() + 720 * 60 * 60 * 1000);
|
||||
const actualExpirationTime = new Date(result.expiredAt);
|
||||
|
||||
expect(actualExpirationTime.getTime()).toBeGreaterThanOrEqual(
|
||||
expectedExpirationTime.getTime() - 1000,
|
||||
);
|
||||
expect(actualExpirationTime.getTime()).toBeLessThanOrEqual(
|
||||
new Date(afterSave.getTime() + 720 * 60 * 60 * 1000 + 1000).getTime(),
|
||||
);
|
||||
});
|
||||
|
||||
it('should use default retention when config is not provided', async () => {
|
||||
// Mock getAppConfig to return empty config
|
||||
mockReq.config = {}; // Empty config
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
|
||||
const beforeSave = new Date();
|
||||
const result = await saveMessage(mockReq, mockMessageData);
|
||||
|
||||
expect(result.expiredAt).toBeDefined();
|
||||
|
||||
// Default retention is 30 days (720 hours)
|
||||
const expectedExpirationTime = new Date(beforeSave.getTime() + 30 * 24 * 60 * 60 * 1000);
|
||||
const actualExpirationTime = new Date(result.expiredAt);
|
||||
|
||||
expect(actualExpirationTime.getTime()).toBeGreaterThanOrEqual(
|
||||
expectedExpirationTime.getTime() - 1000,
|
||||
);
|
||||
expect(actualExpirationTime.getTime()).toBeLessThanOrEqual(
|
||||
expectedExpirationTime.getTime() + 1000,
|
||||
);
|
||||
});
|
||||
|
||||
it('should not update expiredAt on message update', async () => {
|
||||
// First save a temporary message
|
||||
mockReq.config.interfaceConfig.temporaryChatRetention = 24;
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
const savedMessage = await saveMessage(mockReq, mockMessageData);
|
||||
const originalExpiredAt = savedMessage.expiredAt;
|
||||
|
||||
// Now update the message without isTemporary flag
|
||||
mockReq.body = {};
|
||||
const updatedMessage = await updateMessage(mockReq, {
|
||||
messageId: 'msg123',
|
||||
text: 'Updated text',
|
||||
});
|
||||
|
||||
// expiredAt should not be in the returned updated message object
|
||||
expect(updatedMessage.expiredAt).toBeUndefined();
|
||||
|
||||
// Verify in database that expiredAt wasn't changed
|
||||
const dbMessage = await Message.findOne({ messageId: 'msg123', user: 'user123' });
|
||||
expect(dbMessage.expiredAt).toEqual(originalExpiredAt);
|
||||
});
|
||||
|
||||
it('should preserve expiredAt when saving existing temporary message', async () => {
|
||||
// First save a temporary message
|
||||
mockReq.config.interfaceConfig.temporaryChatRetention = 24;
|
||||
|
||||
mockReq.body = { isTemporary: true };
|
||||
const firstSave = await saveMessage(mockReq, mockMessageData);
|
||||
const originalExpiredAt = firstSave.expiredAt;
|
||||
|
||||
// Wait a bit to ensure time difference
|
||||
await new Promise((resolve) => setTimeout(resolve, 100));
|
||||
|
||||
// Save again with same messageId but different text
|
||||
const updatedData = { ...mockMessageData, text: 'Updated text' };
|
||||
const secondSave = await saveMessage(mockReq, updatedData);
|
||||
|
||||
// Should update text but create new expiredAt
|
||||
expect(secondSave.text).toBe('Updated text');
|
||||
expect(secondSave.expiredAt).toBeDefined();
|
||||
expect(new Date(secondSave.expiredAt).getTime()).toBeGreaterThan(
|
||||
new Date(originalExpiredAt).getTime(),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle bulk operations with temporary messages', async () => {
|
||||
// This test verifies bulkSaveMessages doesn't interfere with expiredAt
|
||||
const messages = [
|
||||
{
|
||||
messageId: 'bulk1',
|
||||
conversationId: uuidv4(),
|
||||
text: 'Bulk message 1',
|
||||
user: 'user123',
|
||||
expiredAt: new Date(Date.now() + 24 * 60 * 60 * 1000),
|
||||
},
|
||||
{
|
||||
messageId: 'bulk2',
|
||||
conversationId: uuidv4(),
|
||||
text: 'Bulk message 2',
|
||||
user: 'user123',
|
||||
expiredAt: null,
|
||||
},
|
||||
];
|
||||
|
||||
await bulkSaveMessages(messages);
|
||||
|
||||
const savedMessages = await Message.find({
|
||||
messageId: { $in: ['bulk1', 'bulk2'] },
|
||||
}).lean();
|
||||
|
||||
expect(savedMessages).toHaveLength(2);
|
||||
|
||||
const bulk1 = savedMessages.find((m) => m.messageId === 'bulk1');
|
||||
const bulk2 = savedMessages.find((m) => m.messageId === 'bulk2');
|
||||
|
||||
expect(bulk1.expiredAt).toBeDefined();
|
||||
expect(bulk2.expiredAt).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Message cursor pagination', () => {
|
||||
/**
|
||||
* Helper to create messages with specific timestamps
|
||||
* Uses collection.insertOne to bypass Mongoose timestamps
|
||||
*/
|
||||
const createMessageWithTimestamp = async (index, conversationId, createdAt) => {
|
||||
const messageId = uuidv4();
|
||||
await Message.collection.insertOne({
|
||||
messageId,
|
||||
conversationId,
|
||||
user: 'user123',
|
||||
text: `Message ${index}`,
|
||||
isCreatedByUser: index % 2 === 0,
|
||||
createdAt,
|
||||
updatedAt: createdAt,
|
||||
});
|
||||
return Message.findOne({ messageId }).lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Simulates the pagination logic from api/server/routes/messages.js
|
||||
* This tests the exact query pattern used in the route
|
||||
*/
|
||||
const getMessagesByCursor = async ({
|
||||
conversationId,
|
||||
user,
|
||||
pageSize = 25,
|
||||
cursor = null,
|
||||
sortBy = 'createdAt',
|
||||
sortDirection = 'desc',
|
||||
}) => {
|
||||
const sortOrder = sortDirection === 'asc' ? 1 : -1;
|
||||
const sortField = ['createdAt', 'updatedAt'].includes(sortBy) ? sortBy : 'createdAt';
|
||||
const cursorOperator = sortDirection === 'asc' ? '$gt' : '$lt';
|
||||
|
||||
const filter = { conversationId, user };
|
||||
if (cursor) {
|
||||
filter[sortField] = { [cursorOperator]: new Date(cursor) };
|
||||
}
|
||||
|
||||
const messages = await Message.find(filter)
|
||||
.sort({ [sortField]: sortOrder })
|
||||
.limit(pageSize + 1)
|
||||
.lean();
|
||||
|
||||
let nextCursor = null;
|
||||
if (messages.length > pageSize) {
|
||||
messages.pop(); // Remove extra item used to detect next page
|
||||
// Create cursor from the last RETURNED item (not the popped one)
|
||||
nextCursor = messages[messages.length - 1][sortField];
|
||||
}
|
||||
|
||||
return { messages, nextCursor };
|
||||
};
|
||||
|
||||
it('should return messages for a conversation with pagination', async () => {
|
||||
const conversationId = uuidv4();
|
||||
const baseTime = new Date('2026-01-01T00:00:00.000Z');
|
||||
|
||||
// Create 30 messages to test pagination
|
||||
for (let i = 0; i < 30; i++) {
|
||||
const createdAt = new Date(baseTime.getTime() - i * 60000); // Each 1 minute apart
|
||||
await createMessageWithTimestamp(i, conversationId, createdAt);
|
||||
}
|
||||
|
||||
// Fetch first page (pageSize 25)
|
||||
const page1 = await getMessagesByCursor({
|
||||
conversationId,
|
||||
user: 'user123',
|
||||
pageSize: 25,
|
||||
});
|
||||
|
||||
expect(page1.messages).toHaveLength(25);
|
||||
expect(page1.nextCursor).toBeTruthy();
|
||||
|
||||
// Fetch second page using cursor
|
||||
const page2 = await getMessagesByCursor({
|
||||
conversationId,
|
||||
user: 'user123',
|
||||
pageSize: 25,
|
||||
cursor: page1.nextCursor,
|
||||
});
|
||||
|
||||
// Should get remaining 5 messages
|
||||
expect(page2.messages).toHaveLength(5);
|
||||
expect(page2.nextCursor).toBeNull();
|
||||
|
||||
// Verify no duplicates and no gaps
|
||||
const allMessageIds = [
|
||||
...page1.messages.map((m) => m.messageId),
|
||||
...page2.messages.map((m) => m.messageId),
|
||||
];
|
||||
const uniqueIds = new Set(allMessageIds);
|
||||
|
||||
expect(uniqueIds.size).toBe(30); // All 30 messages accounted for
|
||||
expect(allMessageIds.length).toBe(30); // No duplicates
|
||||
});
|
||||
|
||||
it('should not skip message at page boundary (item 26 bug fix)', async () => {
|
||||
const conversationId = uuidv4();
|
||||
const baseTime = new Date('2026-01-01T12:00:00.000Z');
|
||||
|
||||
// Create exactly 26 messages
|
||||
const messages = [];
|
||||
for (let i = 0; i < 26; i++) {
|
||||
const createdAt = new Date(baseTime.getTime() - i * 60000);
|
||||
const msg = await createMessageWithTimestamp(i, conversationId, createdAt);
|
||||
messages.push(msg);
|
||||
}
|
||||
|
||||
// The 26th message (index 25) should be on page 2
|
||||
const item26 = messages[25];
|
||||
|
||||
// Fetch first page with pageSize 25
|
||||
const page1 = await getMessagesByCursor({
|
||||
conversationId,
|
||||
user: 'user123',
|
||||
pageSize: 25,
|
||||
});
|
||||
|
||||
expect(page1.messages).toHaveLength(25);
|
||||
expect(page1.nextCursor).toBeTruthy();
|
||||
|
||||
// Item 26 should NOT be in page 1
|
||||
const page1Ids = page1.messages.map((m) => m.messageId);
|
||||
expect(page1Ids).not.toContain(item26.messageId);
|
||||
|
||||
// Fetch second page
|
||||
const page2 = await getMessagesByCursor({
|
||||
conversationId,
|
||||
user: 'user123',
|
||||
pageSize: 25,
|
||||
cursor: page1.nextCursor,
|
||||
});
|
||||
|
||||
// Item 26 MUST be in page 2 (this was the bug - it was being skipped)
|
||||
expect(page2.messages).toHaveLength(1);
|
||||
expect(page2.messages[0].messageId).toBe(item26.messageId);
|
||||
});
|
||||
|
||||
it('should sort by createdAt DESC by default', async () => {
|
||||
const conversationId = uuidv4();
|
||||
|
||||
// Create messages with specific timestamps
|
||||
const msg1 = await createMessageWithTimestamp(
|
||||
1,
|
||||
conversationId,
|
||||
new Date('2026-01-01T00:00:00.000Z'),
|
||||
);
|
||||
const msg2 = await createMessageWithTimestamp(
|
||||
2,
|
||||
conversationId,
|
||||
new Date('2026-01-02T00:00:00.000Z'),
|
||||
);
|
||||
const msg3 = await createMessageWithTimestamp(
|
||||
3,
|
||||
conversationId,
|
||||
new Date('2026-01-03T00:00:00.000Z'),
|
||||
);
|
||||
|
||||
const result = await getMessagesByCursor({
|
||||
conversationId,
|
||||
user: 'user123',
|
||||
});
|
||||
|
||||
// Should be sorted by createdAt DESC (newest first) by default
|
||||
expect(result.messages).toHaveLength(3);
|
||||
expect(result.messages[0].messageId).toBe(msg3.messageId);
|
||||
expect(result.messages[1].messageId).toBe(msg2.messageId);
|
||||
expect(result.messages[2].messageId).toBe(msg1.messageId);
|
||||
});
|
||||
|
||||
it('should support ascending sort direction', async () => {
|
||||
const conversationId = uuidv4();
|
||||
|
||||
const msg1 = await createMessageWithTimestamp(
|
||||
1,
|
||||
conversationId,
|
||||
new Date('2026-01-01T00:00:00.000Z'),
|
||||
);
|
||||
const msg2 = await createMessageWithTimestamp(
|
||||
2,
|
||||
conversationId,
|
||||
new Date('2026-01-02T00:00:00.000Z'),
|
||||
);
|
||||
|
||||
const result = await getMessagesByCursor({
|
||||
conversationId,
|
||||
user: 'user123',
|
||||
sortDirection: 'asc',
|
||||
});
|
||||
|
||||
// Should be sorted by createdAt ASC (oldest first)
|
||||
expect(result.messages).toHaveLength(2);
|
||||
expect(result.messages[0].messageId).toBe(msg1.messageId);
|
||||
expect(result.messages[1].messageId).toBe(msg2.messageId);
|
||||
});
|
||||
|
||||
it('should handle empty conversation', async () => {
|
||||
const conversationId = uuidv4();
|
||||
|
||||
const result = await getMessagesByCursor({
|
||||
conversationId,
|
||||
user: 'user123',
|
||||
});
|
||||
|
||||
expect(result.messages).toHaveLength(0);
|
||||
expect(result.nextCursor).toBeNull();
|
||||
});
|
||||
|
||||
it('should only return messages for the specified user', async () => {
|
||||
const conversationId = uuidv4();
|
||||
const createdAt = new Date();
|
||||
|
||||
// Create a message for user123
|
||||
await Message.collection.insertOne({
|
||||
messageId: uuidv4(),
|
||||
conversationId,
|
||||
user: 'user123',
|
||||
text: 'User message',
|
||||
createdAt,
|
||||
updatedAt: createdAt,
|
||||
});
|
||||
|
||||
// Create a message for a different user
|
||||
await Message.collection.insertOne({
|
||||
messageId: uuidv4(),
|
||||
conversationId,
|
||||
user: 'otherUser',
|
||||
text: 'Other user message',
|
||||
createdAt,
|
||||
updatedAt: createdAt,
|
||||
});
|
||||
|
||||
const result = await getMessagesByCursor({
|
||||
conversationId,
|
||||
user: 'user123',
|
||||
});
|
||||
|
||||
// Should only return user123's message
|
||||
expect(result.messages).toHaveLength(1);
|
||||
expect(result.messages[0].user).toBe('user123');
|
||||
});
|
||||
|
||||
it('should handle exactly pageSize number of messages (no next page)', async () => {
|
||||
const conversationId = uuidv4();
|
||||
const baseTime = new Date('2026-01-01T00:00:00.000Z');
|
||||
|
||||
// Create exactly 25 messages (equal to default pageSize)
|
||||
for (let i = 0; i < 25; i++) {
|
||||
const createdAt = new Date(baseTime.getTime() - i * 60000);
|
||||
await createMessageWithTimestamp(i, conversationId, createdAt);
|
||||
}
|
||||
|
||||
const result = await getMessagesByCursor({
|
||||
conversationId,
|
||||
user: 'user123',
|
||||
pageSize: 25,
|
||||
});
|
||||
|
||||
expect(result.messages).toHaveLength(25);
|
||||
expect(result.nextCursor).toBeNull(); // No next page
|
||||
});
|
||||
|
||||
it('should handle pageSize of 1', async () => {
|
||||
const conversationId = uuidv4();
|
||||
const baseTime = new Date('2026-01-01T00:00:00.000Z');
|
||||
|
||||
// Create 3 messages
|
||||
for (let i = 0; i < 3; i++) {
|
||||
const createdAt = new Date(baseTime.getTime() - i * 60000);
|
||||
await createMessageWithTimestamp(i, conversationId, createdAt);
|
||||
}
|
||||
|
||||
// Fetch with pageSize 1
|
||||
let cursor = null;
|
||||
const allMessages = [];
|
||||
|
||||
for (let page = 0; page < 5; page++) {
|
||||
const result = await getMessagesByCursor({
|
||||
conversationId,
|
||||
user: 'user123',
|
||||
pageSize: 1,
|
||||
cursor,
|
||||
});
|
||||
|
||||
allMessages.push(...result.messages);
|
||||
cursor = result.nextCursor;
|
||||
|
||||
if (!cursor) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Should get all 3 messages without duplicates
|
||||
expect(allMessages).toHaveLength(3);
|
||||
const uniqueIds = new Set(allMessages.map((m) => m.messageId));
|
||||
expect(uniqueIds.size).toBe(3);
|
||||
});
|
||||
|
||||
it('should handle messages with same createdAt timestamp', async () => {
|
||||
const conversationId = uuidv4();
|
||||
const sameTime = new Date('2026-01-01T12:00:00.000Z');
|
||||
|
||||
// Create multiple messages with the exact same timestamp
|
||||
const messages = [];
|
||||
for (let i = 0; i < 5; i++) {
|
||||
const msg = await createMessageWithTimestamp(i, conversationId, sameTime);
|
||||
messages.push(msg);
|
||||
}
|
||||
|
||||
const result = await getMessagesByCursor({
|
||||
conversationId,
|
||||
user: 'user123',
|
||||
pageSize: 10,
|
||||
});
|
||||
|
||||
// All messages should be returned
|
||||
expect(result.messages).toHaveLength(5);
|
||||
});
|
||||
});
|
||||
});
|
||||
82
api/models/Preset.js
Normal file
82
api/models/Preset.js
Normal file
@ -0,0 +1,82 @@
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { Preset } = require('~/db/models');
|
||||
|
||||
const getPreset = async (user, presetId) => {
|
||||
try {
|
||||
return await Preset.findOne({ user, presetId }).lean();
|
||||
} catch (error) {
|
||||
logger.error('[getPreset] Error getting single preset', error);
|
||||
return { message: 'Error getting single preset' };
|
||||
}
|
||||
};
|
||||
|
||||
module.exports = {
|
||||
getPreset,
|
||||
getPresets: async (user, filter) => {
|
||||
try {
|
||||
const presets = await Preset.find({ ...filter, user }).lean();
|
||||
const defaultValue = 10000;
|
||||
|
||||
presets.sort((a, b) => {
|
||||
let orderA = a.order !== undefined ? a.order : defaultValue;
|
||||
let orderB = b.order !== undefined ? b.order : defaultValue;
|
||||
|
||||
if (orderA !== orderB) {
|
||||
return orderA - orderB;
|
||||
}
|
||||
|
||||
return b.updatedAt - a.updatedAt;
|
||||
});
|
||||
|
||||
return presets;
|
||||
} catch (error) {
|
||||
logger.error('[getPresets] Error getting presets', error);
|
||||
return { message: 'Error retrieving presets' };
|
||||
}
|
||||
},
|
||||
savePreset: async (user, { presetId, newPresetId, defaultPreset, ...preset }) => {
|
||||
try {
|
||||
const setter = { $set: {} };
|
||||
const { user: _, ...cleanPreset } = preset;
|
||||
const update = { presetId, ...cleanPreset };
|
||||
if (preset.tools && Array.isArray(preset.tools)) {
|
||||
update.tools =
|
||||
preset.tools
|
||||
.map((tool) => tool?.pluginKey ?? tool)
|
||||
.filter((toolName) => typeof toolName === 'string') ?? [];
|
||||
}
|
||||
if (newPresetId) {
|
||||
update.presetId = newPresetId;
|
||||
}
|
||||
|
||||
if (defaultPreset) {
|
||||
update.defaultPreset = defaultPreset;
|
||||
update.order = 0;
|
||||
|
||||
const currentDefault = await Preset.findOne({ defaultPreset: true, user });
|
||||
|
||||
if (currentDefault && currentDefault.presetId !== presetId) {
|
||||
await Preset.findByIdAndUpdate(currentDefault._id, {
|
||||
$unset: { defaultPreset: '', order: '' },
|
||||
});
|
||||
}
|
||||
} else if (defaultPreset === false) {
|
||||
update.defaultPreset = undefined;
|
||||
update.order = undefined;
|
||||
setter['$unset'] = { defaultPreset: '', order: '' };
|
||||
}
|
||||
|
||||
setter.$set = update;
|
||||
return await Preset.findOneAndUpdate({ presetId, user }, setter, { new: true, upsert: true });
|
||||
} catch (error) {
|
||||
logger.error('[savePreset] Error saving preset', error);
|
||||
return { message: 'Error saving preset' };
|
||||
}
|
||||
},
|
||||
deletePresets: async (user, filter) => {
|
||||
// let toRemove = await Preset.find({ ...filter, user }).select('presetId');
|
||||
// const ids = toRemove.map((instance) => instance.presetId);
|
||||
let deleteCount = await Preset.deleteMany({ ...filter, user });
|
||||
return deleteCount;
|
||||
},
|
||||
};
|
||||
133
api/models/Project.js
Normal file
133
api/models/Project.js
Normal file
@ -0,0 +1,133 @@
|
||||
const { GLOBAL_PROJECT_NAME } = require('librechat-data-provider').Constants;
|
||||
const { Project } = require('~/db/models');
|
||||
|
||||
/**
|
||||
* Retrieve a project by ID and convert the found project document to a plain object.
|
||||
*
|
||||
* @param {string} projectId - The ID of the project to find and return as a plain object.
|
||||
* @param {string|string[]} [fieldsToSelect] - The fields to include or exclude in the returned document.
|
||||
* @returns {Promise<IMongoProject>} A plain object representing the project document, or `null` if no project is found.
|
||||
*/
|
||||
const getProjectById = async function (projectId, fieldsToSelect = null) {
|
||||
const query = Project.findById(projectId);
|
||||
|
||||
if (fieldsToSelect) {
|
||||
query.select(fieldsToSelect);
|
||||
}
|
||||
|
||||
return await query.lean();
|
||||
};
|
||||
|
||||
/**
|
||||
* Retrieve a project by name and convert the found project document to a plain object.
|
||||
* If the project with the given name doesn't exist and the name is "instance", create it and return the lean version.
|
||||
*
|
||||
* @param {string} projectName - The name of the project to find or create.
|
||||
* @param {string|string[]} [fieldsToSelect] - The fields to include or exclude in the returned document.
|
||||
* @returns {Promise<IMongoProject>} A plain object representing the project document.
|
||||
*/
|
||||
const getProjectByName = async function (projectName, fieldsToSelect = null) {
|
||||
const query = { name: projectName };
|
||||
const update = { $setOnInsert: { name: projectName } };
|
||||
const options = {
|
||||
new: true,
|
||||
upsert: projectName === GLOBAL_PROJECT_NAME,
|
||||
lean: true,
|
||||
select: fieldsToSelect,
|
||||
};
|
||||
|
||||
return await Project.findOneAndUpdate(query, update, options);
|
||||
};
|
||||
|
||||
/**
|
||||
* Add an array of prompt group IDs to a project's promptGroupIds array, ensuring uniqueness.
|
||||
*
|
||||
* @param {string} projectId - The ID of the project to update.
|
||||
* @param {string[]} promptGroupIds - The array of prompt group IDs to add to the project.
|
||||
* @returns {Promise<IMongoProject>} The updated project document.
|
||||
*/
|
||||
const addGroupIdsToProject = async function (projectId, promptGroupIds) {
|
||||
return await Project.findByIdAndUpdate(
|
||||
projectId,
|
||||
{ $addToSet: { promptGroupIds: { $each: promptGroupIds } } },
|
||||
{ new: true },
|
||||
);
|
||||
};
|
||||
|
||||
/**
|
||||
* Remove an array of prompt group IDs from a project's promptGroupIds array.
|
||||
*
|
||||
* @param {string} projectId - The ID of the project to update.
|
||||
* @param {string[]} promptGroupIds - The array of prompt group IDs to remove from the project.
|
||||
* @returns {Promise<IMongoProject>} The updated project document.
|
||||
*/
|
||||
const removeGroupIdsFromProject = async function (projectId, promptGroupIds) {
|
||||
return await Project.findByIdAndUpdate(
|
||||
projectId,
|
||||
{ $pull: { promptGroupIds: { $in: promptGroupIds } } },
|
||||
{ new: true },
|
||||
);
|
||||
};
|
||||
|
||||
/**
|
||||
* Remove a prompt group ID from all projects.
|
||||
*
|
||||
* @param {string} promptGroupId - The ID of the prompt group to remove from projects.
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
const removeGroupFromAllProjects = async (promptGroupId) => {
|
||||
await Project.updateMany({}, { $pull: { promptGroupIds: promptGroupId } });
|
||||
};
|
||||
|
||||
/**
|
||||
* Add an array of agent IDs to a project's agentIds array, ensuring uniqueness.
|
||||
*
|
||||
* @param {string} projectId - The ID of the project to update.
|
||||
* @param {string[]} agentIds - The array of agent IDs to add to the project.
|
||||
* @returns {Promise<IMongoProject>} The updated project document.
|
||||
*/
|
||||
const addAgentIdsToProject = async function (projectId, agentIds) {
|
||||
return await Project.findByIdAndUpdate(
|
||||
projectId,
|
||||
{ $addToSet: { agentIds: { $each: agentIds } } },
|
||||
{ new: true },
|
||||
);
|
||||
};
|
||||
|
||||
/**
|
||||
* Remove an array of agent IDs from a project's agentIds array.
|
||||
*
|
||||
* @param {string} projectId - The ID of the project to update.
|
||||
* @param {string[]} agentIds - The array of agent IDs to remove from the project.
|
||||
* @returns {Promise<IMongoProject>} The updated project document.
|
||||
*/
|
||||
const removeAgentIdsFromProject = async function (projectId, agentIds) {
|
||||
return await Project.findByIdAndUpdate(
|
||||
projectId,
|
||||
{ $pull: { agentIds: { $in: agentIds } } },
|
||||
{ new: true },
|
||||
);
|
||||
};
|
||||
|
||||
/**
|
||||
* Remove an agent ID from all projects.
|
||||
*
|
||||
* @param {string} agentId - The ID of the agent to remove from projects.
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
const removeAgentFromAllProjects = async (agentId) => {
|
||||
await Project.updateMany({}, { $pull: { agentIds: agentId } });
|
||||
};
|
||||
|
||||
module.exports = {
|
||||
getProjectById,
|
||||
getProjectByName,
|
||||
/* prompts */
|
||||
addGroupIdsToProject,
|
||||
removeGroupIdsFromProject,
|
||||
removeGroupFromAllProjects,
|
||||
/* agents */
|
||||
addAgentIdsToProject,
|
||||
removeAgentIdsFromProject,
|
||||
removeAgentFromAllProjects,
|
||||
};
|
||||
708
api/models/Prompt.js
Normal file
708
api/models/Prompt.js
Normal file
@ -0,0 +1,708 @@
|
||||
const { ObjectId } = require('mongodb');
|
||||
const { escapeRegExp } = require('@librechat/api');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const {
|
||||
Constants,
|
||||
SystemRoles,
|
||||
ResourceType,
|
||||
SystemCategories,
|
||||
} = require('librechat-data-provider');
|
||||
const {
|
||||
removeGroupFromAllProjects,
|
||||
removeGroupIdsFromProject,
|
||||
addGroupIdsToProject,
|
||||
getProjectByName,
|
||||
} = require('./Project');
|
||||
const { removeAllPermissions } = require('~/server/services/PermissionService');
|
||||
const { PromptGroup, Prompt, AclEntry } = require('~/db/models');
|
||||
|
||||
/**
|
||||
* Create a pipeline for the aggregation to get prompt groups
|
||||
* @param {Object} query
|
||||
* @param {number} skip
|
||||
* @param {number} limit
|
||||
* @returns {[Object]} - The pipeline for the aggregation
|
||||
*/
|
||||
const createGroupPipeline = (query, skip, limit) => {
|
||||
return [
|
||||
{ $match: query },
|
||||
{ $sort: { createdAt: -1 } },
|
||||
{ $skip: skip },
|
||||
{ $limit: limit },
|
||||
{
|
||||
$lookup: {
|
||||
from: 'prompts',
|
||||
localField: 'productionId',
|
||||
foreignField: '_id',
|
||||
as: 'productionPrompt',
|
||||
},
|
||||
},
|
||||
{ $unwind: { path: '$productionPrompt', preserveNullAndEmptyArrays: true } },
|
||||
{
|
||||
$project: {
|
||||
name: 1,
|
||||
numberOfGenerations: 1,
|
||||
oneliner: 1,
|
||||
category: 1,
|
||||
projectIds: 1,
|
||||
productionId: 1,
|
||||
author: 1,
|
||||
authorName: 1,
|
||||
createdAt: 1,
|
||||
updatedAt: 1,
|
||||
'productionPrompt.prompt': 1,
|
||||
// 'productionPrompt._id': 1,
|
||||
// 'productionPrompt.type': 1,
|
||||
},
|
||||
},
|
||||
];
|
||||
};
|
||||
|
||||
/**
|
||||
* Create a pipeline for the aggregation to get all prompt groups
|
||||
* @param {Object} query
|
||||
* @param {Partial<MongoPromptGroup>} $project
|
||||
* @returns {[Object]} - The pipeline for the aggregation
|
||||
*/
|
||||
const createAllGroupsPipeline = (
|
||||
query,
|
||||
$project = {
|
||||
name: 1,
|
||||
oneliner: 1,
|
||||
category: 1,
|
||||
author: 1,
|
||||
authorName: 1,
|
||||
createdAt: 1,
|
||||
updatedAt: 1,
|
||||
command: 1,
|
||||
'productionPrompt.prompt': 1,
|
||||
},
|
||||
) => {
|
||||
return [
|
||||
{ $match: query },
|
||||
{ $sort: { createdAt: -1 } },
|
||||
{
|
||||
$lookup: {
|
||||
from: 'prompts',
|
||||
localField: 'productionId',
|
||||
foreignField: '_id',
|
||||
as: 'productionPrompt',
|
||||
},
|
||||
},
|
||||
{ $unwind: { path: '$productionPrompt', preserveNullAndEmptyArrays: true } },
|
||||
{
|
||||
$project,
|
||||
},
|
||||
];
|
||||
};
|
||||
|
||||
/**
|
||||
* Get all prompt groups with filters
|
||||
* @param {ServerRequest} req
|
||||
* @param {TPromptGroupsWithFilterRequest} filter
|
||||
* @returns {Promise<PromptGroupListResponse>}
|
||||
*/
|
||||
const getAllPromptGroups = async (req, filter) => {
|
||||
try {
|
||||
const { name, ...query } = filter;
|
||||
|
||||
let searchShared = true;
|
||||
let searchSharedOnly = false;
|
||||
if (name) {
|
||||
query.name = new RegExp(escapeRegExp(name), 'i');
|
||||
}
|
||||
if (!query.category) {
|
||||
delete query.category;
|
||||
} else if (query.category === SystemCategories.MY_PROMPTS) {
|
||||
searchShared = false;
|
||||
delete query.category;
|
||||
} else if (query.category === SystemCategories.NO_CATEGORY) {
|
||||
query.category = '';
|
||||
} else if (query.category === SystemCategories.SHARED_PROMPTS) {
|
||||
searchSharedOnly = true;
|
||||
delete query.category;
|
||||
}
|
||||
|
||||
let combinedQuery = query;
|
||||
|
||||
if (searchShared) {
|
||||
const project = await getProjectByName(Constants.GLOBAL_PROJECT_NAME, 'promptGroupIds');
|
||||
if (project && project.promptGroupIds && project.promptGroupIds.length > 0) {
|
||||
const projectQuery = { _id: { $in: project.promptGroupIds }, ...query };
|
||||
delete projectQuery.author;
|
||||
combinedQuery = searchSharedOnly ? projectQuery : { $or: [projectQuery, query] };
|
||||
}
|
||||
}
|
||||
|
||||
const promptGroupsPipeline = createAllGroupsPipeline(combinedQuery);
|
||||
return await PromptGroup.aggregate(promptGroupsPipeline).exec();
|
||||
} catch (error) {
|
||||
console.error('Error getting all prompt groups', error);
|
||||
return { message: 'Error getting all prompt groups' };
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Get prompt groups with filters
|
||||
* @param {ServerRequest} req
|
||||
* @param {TPromptGroupsWithFilterRequest} filter
|
||||
* @returns {Promise<PromptGroupListResponse>}
|
||||
*/
|
||||
const getPromptGroups = async (req, filter) => {
|
||||
try {
|
||||
const { pageNumber = 1, pageSize = 10, name, ...query } = filter;
|
||||
|
||||
const validatedPageNumber = Math.max(parseInt(pageNumber, 10), 1);
|
||||
const validatedPageSize = Math.max(parseInt(pageSize, 10), 1);
|
||||
|
||||
let searchShared = true;
|
||||
let searchSharedOnly = false;
|
||||
if (name) {
|
||||
query.name = new RegExp(escapeRegExp(name), 'i');
|
||||
}
|
||||
if (!query.category) {
|
||||
delete query.category;
|
||||
} else if (query.category === SystemCategories.MY_PROMPTS) {
|
||||
searchShared = false;
|
||||
delete query.category;
|
||||
} else if (query.category === SystemCategories.NO_CATEGORY) {
|
||||
query.category = '';
|
||||
} else if (query.category === SystemCategories.SHARED_PROMPTS) {
|
||||
searchSharedOnly = true;
|
||||
delete query.category;
|
||||
}
|
||||
|
||||
let combinedQuery = query;
|
||||
|
||||
if (searchShared) {
|
||||
// const projects = req.user.projects || []; // TODO: handle multiple projects
|
||||
const project = await getProjectByName(Constants.GLOBAL_PROJECT_NAME, 'promptGroupIds');
|
||||
if (project && project.promptGroupIds && project.promptGroupIds.length > 0) {
|
||||
const projectQuery = { _id: { $in: project.promptGroupIds }, ...query };
|
||||
delete projectQuery.author;
|
||||
combinedQuery = searchSharedOnly ? projectQuery : { $or: [projectQuery, query] };
|
||||
}
|
||||
}
|
||||
|
||||
const skip = (validatedPageNumber - 1) * validatedPageSize;
|
||||
const limit = validatedPageSize;
|
||||
|
||||
const promptGroupsPipeline = createGroupPipeline(combinedQuery, skip, limit);
|
||||
const totalPromptGroupsPipeline = [{ $match: combinedQuery }, { $count: 'total' }];
|
||||
|
||||
const [promptGroupsResults, totalPromptGroupsResults] = await Promise.all([
|
||||
PromptGroup.aggregate(promptGroupsPipeline).exec(),
|
||||
PromptGroup.aggregate(totalPromptGroupsPipeline).exec(),
|
||||
]);
|
||||
|
||||
const promptGroups = promptGroupsResults;
|
||||
const totalPromptGroups =
|
||||
totalPromptGroupsResults.length > 0 ? totalPromptGroupsResults[0].total : 0;
|
||||
|
||||
return {
|
||||
promptGroups,
|
||||
pageNumber: validatedPageNumber.toString(),
|
||||
pageSize: validatedPageSize.toString(),
|
||||
pages: Math.ceil(totalPromptGroups / validatedPageSize).toString(),
|
||||
};
|
||||
} catch (error) {
|
||||
console.error('Error getting prompt groups', error);
|
||||
return { message: 'Error getting prompt groups' };
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* @param {Object} fields
|
||||
* @param {string} fields._id
|
||||
* @param {string} fields.author
|
||||
* @param {string} fields.role
|
||||
* @returns {Promise<TDeletePromptGroupResponse>}
|
||||
*/
|
||||
const deletePromptGroup = async ({ _id, author, role }) => {
|
||||
// Build query - with ACL, author is optional
|
||||
const query = { _id };
|
||||
const groupQuery = { groupId: new ObjectId(_id) };
|
||||
|
||||
// Legacy: Add author filter if provided (backward compatibility)
|
||||
if (author && role !== SystemRoles.ADMIN) {
|
||||
query.author = author;
|
||||
groupQuery.author = author;
|
||||
}
|
||||
|
||||
const response = await PromptGroup.deleteOne(query);
|
||||
|
||||
if (!response || response.deletedCount === 0) {
|
||||
throw new Error('Prompt group not found');
|
||||
}
|
||||
|
||||
await Prompt.deleteMany(groupQuery);
|
||||
await removeGroupFromAllProjects(_id);
|
||||
|
||||
try {
|
||||
await removeAllPermissions({ resourceType: ResourceType.PROMPTGROUP, resourceId: _id });
|
||||
} catch (error) {
|
||||
logger.error('Error removing promptGroup permissions:', error);
|
||||
}
|
||||
|
||||
return { message: 'Prompt group deleted successfully' };
|
||||
};
|
||||
|
||||
/**
|
||||
* Get prompt groups by accessible IDs with optional cursor-based pagination.
|
||||
* @param {Object} params - The parameters for getting accessible prompt groups.
|
||||
* @param {Array} [params.accessibleIds] - Array of prompt group ObjectIds the user has ACL access to.
|
||||
* @param {Object} [params.otherParams] - Additional query parameters (including author filter).
|
||||
* @param {number} [params.limit] - Number of prompt groups to return (max 100). If not provided, returns all prompt groups.
|
||||
* @param {string} [params.after] - Cursor for pagination - get prompt groups after this cursor. // base64 encoded JSON string with updatedAt and _id.
|
||||
* @returns {Promise<Object>} A promise that resolves to an object containing the prompt groups data and pagination info.
|
||||
*/
|
||||
async function getListPromptGroupsByAccess({
|
||||
accessibleIds = [],
|
||||
otherParams = {},
|
||||
limit = null,
|
||||
after = null,
|
||||
}) {
|
||||
const isPaginated = limit !== null && limit !== undefined;
|
||||
const normalizedLimit = isPaginated ? Math.min(Math.max(1, parseInt(limit) || 20), 100) : null;
|
||||
|
||||
// Build base query combining ACL accessible prompt groups with other filters
|
||||
const baseQuery = { ...otherParams, _id: { $in: accessibleIds } };
|
||||
|
||||
// Add cursor condition
|
||||
if (after && typeof after === 'string' && after !== 'undefined' && after !== 'null') {
|
||||
try {
|
||||
const cursor = JSON.parse(Buffer.from(after, 'base64').toString('utf8'));
|
||||
const { updatedAt, _id } = cursor;
|
||||
|
||||
const cursorCondition = {
|
||||
$or: [
|
||||
{ updatedAt: { $lt: new Date(updatedAt) } },
|
||||
{ updatedAt: new Date(updatedAt), _id: { $gt: new ObjectId(_id) } },
|
||||
],
|
||||
};
|
||||
|
||||
// Merge cursor condition with base query
|
||||
if (Object.keys(baseQuery).length > 0) {
|
||||
baseQuery.$and = [{ ...baseQuery }, cursorCondition];
|
||||
// Remove the original conditions from baseQuery to avoid duplication
|
||||
Object.keys(baseQuery).forEach((key) => {
|
||||
if (key !== '$and') delete baseQuery[key];
|
||||
});
|
||||
} else {
|
||||
Object.assign(baseQuery, cursorCondition);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.warn('Invalid cursor:', error.message);
|
||||
}
|
||||
}
|
||||
|
||||
// Build aggregation pipeline
|
||||
const pipeline = [{ $match: baseQuery }, { $sort: { updatedAt: -1, _id: 1 } }];
|
||||
|
||||
// Only apply limit if pagination is requested
|
||||
if (isPaginated) {
|
||||
pipeline.push({ $limit: normalizedLimit + 1 });
|
||||
}
|
||||
|
||||
// Add lookup for production prompt
|
||||
pipeline.push(
|
||||
{
|
||||
$lookup: {
|
||||
from: 'prompts',
|
||||
localField: 'productionId',
|
||||
foreignField: '_id',
|
||||
as: 'productionPrompt',
|
||||
},
|
||||
},
|
||||
{ $unwind: { path: '$productionPrompt', preserveNullAndEmptyArrays: true } },
|
||||
{
|
||||
$project: {
|
||||
name: 1,
|
||||
numberOfGenerations: 1,
|
||||
oneliner: 1,
|
||||
category: 1,
|
||||
projectIds: 1,
|
||||
productionId: 1,
|
||||
author: 1,
|
||||
authorName: 1,
|
||||
createdAt: 1,
|
||||
updatedAt: 1,
|
||||
'productionPrompt.prompt': 1,
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
const promptGroups = await PromptGroup.aggregate(pipeline).exec();
|
||||
|
||||
const hasMore = isPaginated ? promptGroups.length > normalizedLimit : false;
|
||||
const data = (isPaginated ? promptGroups.slice(0, normalizedLimit) : promptGroups).map(
|
||||
(group) => {
|
||||
if (group.author) {
|
||||
group.author = group.author.toString();
|
||||
}
|
||||
return group;
|
||||
},
|
||||
);
|
||||
|
||||
// Generate next cursor only if paginated
|
||||
let nextCursor = null;
|
||||
if (isPaginated && hasMore && data.length > 0) {
|
||||
const lastGroup = promptGroups[normalizedLimit - 1];
|
||||
nextCursor = Buffer.from(
|
||||
JSON.stringify({
|
||||
updatedAt: lastGroup.updatedAt.toISOString(),
|
||||
_id: lastGroup._id.toString(),
|
||||
}),
|
||||
).toString('base64');
|
||||
}
|
||||
|
||||
return {
|
||||
object: 'list',
|
||||
data,
|
||||
first_id: data.length > 0 ? data[0]._id.toString() : null,
|
||||
last_id: data.length > 0 ? data[data.length - 1]._id.toString() : null,
|
||||
has_more: hasMore,
|
||||
after: nextCursor,
|
||||
};
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
getPromptGroups,
|
||||
deletePromptGroup,
|
||||
getAllPromptGroups,
|
||||
getListPromptGroupsByAccess,
|
||||
/**
|
||||
* Create a prompt and its respective group
|
||||
* @param {TCreatePromptRecord} saveData
|
||||
* @returns {Promise<TCreatePromptResponse>}
|
||||
*/
|
||||
createPromptGroup: async (saveData) => {
|
||||
try {
|
||||
const { prompt, group, author, authorName } = saveData;
|
||||
|
||||
let newPromptGroup = await PromptGroup.findOneAndUpdate(
|
||||
{ ...group, author, authorName, productionId: null },
|
||||
{ $setOnInsert: { ...group, author, authorName, productionId: null } },
|
||||
{ new: true, upsert: true },
|
||||
)
|
||||
.lean()
|
||||
.select('-__v')
|
||||
.exec();
|
||||
|
||||
const newPrompt = await Prompt.findOneAndUpdate(
|
||||
{ ...prompt, author, groupId: newPromptGroup._id },
|
||||
{ $setOnInsert: { ...prompt, author, groupId: newPromptGroup._id } },
|
||||
{ new: true, upsert: true },
|
||||
)
|
||||
.lean()
|
||||
.select('-__v')
|
||||
.exec();
|
||||
|
||||
newPromptGroup = await PromptGroup.findByIdAndUpdate(
|
||||
newPromptGroup._id,
|
||||
{ productionId: newPrompt._id },
|
||||
{ new: true },
|
||||
)
|
||||
.lean()
|
||||
.select('-__v')
|
||||
.exec();
|
||||
|
||||
return {
|
||||
prompt: newPrompt,
|
||||
group: {
|
||||
...newPromptGroup,
|
||||
productionPrompt: { prompt: newPrompt.prompt },
|
||||
},
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error('Error saving prompt group', error);
|
||||
throw new Error('Error saving prompt group');
|
||||
}
|
||||
},
|
||||
/**
|
||||
* Save a prompt
|
||||
* @param {TCreatePromptRecord} saveData
|
||||
* @returns {Promise<TCreatePromptResponse>}
|
||||
*/
|
||||
savePrompt: async (saveData) => {
|
||||
try {
|
||||
const { prompt, author } = saveData;
|
||||
const newPromptData = {
|
||||
...prompt,
|
||||
author,
|
||||
};
|
||||
|
||||
/** @type {TPrompt} */
|
||||
let newPrompt;
|
||||
try {
|
||||
newPrompt = await Prompt.create(newPromptData);
|
||||
} catch (error) {
|
||||
if (error?.message?.includes('groupId_1_version_1')) {
|
||||
await Prompt.db.collection('prompts').dropIndex('groupId_1_version_1');
|
||||
} else {
|
||||
throw error;
|
||||
}
|
||||
newPrompt = await Prompt.create(newPromptData);
|
||||
}
|
||||
|
||||
return { prompt: newPrompt };
|
||||
} catch (error) {
|
||||
logger.error('Error saving prompt', error);
|
||||
return { message: 'Error saving prompt' };
|
||||
}
|
||||
},
|
||||
getPrompts: async (filter) => {
|
||||
try {
|
||||
return await Prompt.find(filter).sort({ createdAt: -1 }).lean();
|
||||
} catch (error) {
|
||||
logger.error('Error getting prompts', error);
|
||||
return { message: 'Error getting prompts' };
|
||||
}
|
||||
},
|
||||
getPrompt: async (filter) => {
|
||||
try {
|
||||
if (filter.groupId) {
|
||||
filter.groupId = new ObjectId(filter.groupId);
|
||||
}
|
||||
return await Prompt.findOne(filter).lean();
|
||||
} catch (error) {
|
||||
logger.error('Error getting prompt', error);
|
||||
return { message: 'Error getting prompt' };
|
||||
}
|
||||
},
|
||||
/**
|
||||
* Get prompt groups with filters
|
||||
* @param {TGetRandomPromptsRequest} filter
|
||||
* @returns {Promise<TGetRandomPromptsResponse>}
|
||||
*/
|
||||
getRandomPromptGroups: async (filter) => {
|
||||
try {
|
||||
const result = await PromptGroup.aggregate([
|
||||
{
|
||||
$match: {
|
||||
category: { $ne: '' },
|
||||
},
|
||||
},
|
||||
{
|
||||
$group: {
|
||||
_id: '$category',
|
||||
promptGroup: { $first: '$$ROOT' },
|
||||
},
|
||||
},
|
||||
{
|
||||
$replaceRoot: { newRoot: '$promptGroup' },
|
||||
},
|
||||
{
|
||||
$sample: { size: +filter.limit + +filter.skip },
|
||||
},
|
||||
{
|
||||
$skip: +filter.skip,
|
||||
},
|
||||
{
|
||||
$limit: +filter.limit,
|
||||
},
|
||||
]);
|
||||
return { prompts: result };
|
||||
} catch (error) {
|
||||
logger.error('Error getting prompt groups', error);
|
||||
return { message: 'Error getting prompt groups' };
|
||||
}
|
||||
},
|
||||
getPromptGroupsWithPrompts: async (filter) => {
|
||||
try {
|
||||
return await PromptGroup.findOne(filter)
|
||||
.populate({
|
||||
path: 'prompts',
|
||||
select: '-_id -__v -user',
|
||||
})
|
||||
.select('-_id -__v -user')
|
||||
.lean();
|
||||
} catch (error) {
|
||||
logger.error('Error getting prompt groups', error);
|
||||
return { message: 'Error getting prompt groups' };
|
||||
}
|
||||
},
|
||||
getPromptGroup: async (filter) => {
|
||||
try {
|
||||
return await PromptGroup.findOne(filter).lean();
|
||||
} catch (error) {
|
||||
logger.error('Error getting prompt group', error);
|
||||
return { message: 'Error getting prompt group' };
|
||||
}
|
||||
},
|
||||
/**
|
||||
* Deletes a prompt and its corresponding prompt group if it is the last prompt in the group.
|
||||
*
|
||||
* @param {Object} options - The options for deleting the prompt.
|
||||
* @param {ObjectId|string} options.promptId - The ID of the prompt to delete.
|
||||
* @param {ObjectId|string} options.groupId - The ID of the prompt's group.
|
||||
* @param {ObjectId|string} options.author - The ID of the prompt's author.
|
||||
* @param {string} options.role - The role of the prompt's author.
|
||||
* @return {Promise<TDeletePromptResponse>} An object containing the result of the deletion.
|
||||
* If the prompt was deleted successfully, the object will have a property 'prompt' with the value 'Prompt deleted successfully'.
|
||||
* If the prompt group was deleted successfully, the object will have a property 'promptGroup' with the message 'Prompt group deleted successfully' and id of the deleted group.
|
||||
* If there was an error deleting the prompt, the object will have a property 'message' with the value 'Error deleting prompt'.
|
||||
*/
|
||||
deletePrompt: async ({ promptId, groupId, author, role }) => {
|
||||
const query = { _id: promptId, groupId, author };
|
||||
if (role === SystemRoles.ADMIN) {
|
||||
delete query.author;
|
||||
}
|
||||
const { deletedCount } = await Prompt.deleteOne(query);
|
||||
if (deletedCount === 0) {
|
||||
throw new Error('Failed to delete the prompt');
|
||||
}
|
||||
|
||||
const remainingPrompts = await Prompt.find({ groupId })
|
||||
.select('_id')
|
||||
.sort({ createdAt: 1 })
|
||||
.lean();
|
||||
|
||||
if (remainingPrompts.length === 0) {
|
||||
// Remove all ACL entries for the promptGroup when deleting the last prompt
|
||||
try {
|
||||
await removeAllPermissions({
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: groupId,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error('Error removing promptGroup permissions:', error);
|
||||
}
|
||||
|
||||
await PromptGroup.deleteOne({ _id: groupId });
|
||||
await removeGroupFromAllProjects(groupId);
|
||||
|
||||
return {
|
||||
prompt: 'Prompt deleted successfully',
|
||||
promptGroup: {
|
||||
message: 'Prompt group deleted successfully',
|
||||
id: groupId,
|
||||
},
|
||||
};
|
||||
} else {
|
||||
const promptGroup = await PromptGroup.findById(groupId).lean();
|
||||
if (promptGroup.productionId.toString() === promptId.toString()) {
|
||||
await PromptGroup.updateOne(
|
||||
{ _id: groupId },
|
||||
{ productionId: remainingPrompts[remainingPrompts.length - 1]._id },
|
||||
);
|
||||
}
|
||||
|
||||
return { prompt: 'Prompt deleted successfully' };
|
||||
}
|
||||
},
|
||||
/**
|
||||
* Delete all prompts and prompt groups created by a specific user.
|
||||
* @param {ServerRequest} req - The server request object.
|
||||
* @param {string} userId - The ID of the user whose prompts and prompt groups are to be deleted.
|
||||
*/
|
||||
deleteUserPrompts: async (req, userId) => {
|
||||
try {
|
||||
const promptGroups = await getAllPromptGroups(req, { author: new ObjectId(userId) });
|
||||
|
||||
if (promptGroups.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const groupIds = promptGroups.map((group) => group._id);
|
||||
|
||||
for (const groupId of groupIds) {
|
||||
await removeGroupFromAllProjects(groupId);
|
||||
}
|
||||
|
||||
await AclEntry.deleteMany({
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: { $in: groupIds },
|
||||
});
|
||||
|
||||
await PromptGroup.deleteMany({ author: new ObjectId(userId) });
|
||||
await Prompt.deleteMany({ author: new ObjectId(userId) });
|
||||
} catch (error) {
|
||||
logger.error('[deleteUserPrompts] General error:', error);
|
||||
}
|
||||
},
|
||||
/**
|
||||
* Update prompt group
|
||||
* @param {Partial<MongoPromptGroup>} filter - Filter to find prompt group
|
||||
* @param {Partial<MongoPromptGroup>} data - Data to update
|
||||
* @returns {Promise<TUpdatePromptGroupResponse>}
|
||||
*/
|
||||
updatePromptGroup: async (filter, data) => {
|
||||
try {
|
||||
const updateOps = {};
|
||||
if (data.removeProjectIds) {
|
||||
for (const projectId of data.removeProjectIds) {
|
||||
await removeGroupIdsFromProject(projectId, [filter._id]);
|
||||
}
|
||||
|
||||
updateOps.$pull = { projectIds: { $in: data.removeProjectIds } };
|
||||
delete data.removeProjectIds;
|
||||
}
|
||||
|
||||
if (data.projectIds) {
|
||||
for (const projectId of data.projectIds) {
|
||||
await addGroupIdsToProject(projectId, [filter._id]);
|
||||
}
|
||||
|
||||
updateOps.$addToSet = { projectIds: { $each: data.projectIds } };
|
||||
delete data.projectIds;
|
||||
}
|
||||
|
||||
const updateData = { ...data, ...updateOps };
|
||||
const updatedDoc = await PromptGroup.findOneAndUpdate(filter, updateData, {
|
||||
new: true,
|
||||
upsert: false,
|
||||
});
|
||||
|
||||
if (!updatedDoc) {
|
||||
throw new Error('Prompt group not found');
|
||||
}
|
||||
|
||||
return updatedDoc;
|
||||
} catch (error) {
|
||||
logger.error('Error updating prompt group', error);
|
||||
return { message: 'Error updating prompt group' };
|
||||
}
|
||||
},
|
||||
/**
|
||||
* Function to make a prompt production based on its ID.
|
||||
* @param {String} promptId - The ID of the prompt to make production.
|
||||
* @returns {Object} The result of the production operation.
|
||||
*/
|
||||
makePromptProduction: async (promptId) => {
|
||||
try {
|
||||
const prompt = await Prompt.findById(promptId).lean();
|
||||
|
||||
if (!prompt) {
|
||||
throw new Error('Prompt not found');
|
||||
}
|
||||
|
||||
await PromptGroup.findByIdAndUpdate(
|
||||
prompt.groupId,
|
||||
{ productionId: prompt._id },
|
||||
{ new: true },
|
||||
)
|
||||
.lean()
|
||||
.exec();
|
||||
|
||||
return {
|
||||
message: 'Prompt production made successfully',
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error('Error making prompt production', error);
|
||||
return { message: 'Error making prompt production' };
|
||||
}
|
||||
},
|
||||
updatePromptLabels: async (_id, labels) => {
|
||||
try {
|
||||
const response = await Prompt.updateOne({ _id }, { $set: { labels } });
|
||||
if (response.matchedCount === 0) {
|
||||
return { message: 'Prompt not found' };
|
||||
}
|
||||
return { message: 'Prompt labels updated successfully' };
|
||||
} catch (error) {
|
||||
logger.error('Error updating prompt labels', error);
|
||||
return { message: 'Error updating prompt labels' };
|
||||
}
|
||||
},
|
||||
};
|
||||
564
api/models/Prompt.spec.js
Normal file
564
api/models/Prompt.spec.js
Normal file
@ -0,0 +1,564 @@
|
||||
const mongoose = require('mongoose');
|
||||
const { ObjectId } = require('mongodb');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
const {
|
||||
SystemRoles,
|
||||
ResourceType,
|
||||
AccessRoleIds,
|
||||
PrincipalType,
|
||||
PermissionBits,
|
||||
} = require('librechat-data-provider');
|
||||
|
||||
// Mock the config/connect module to prevent connection attempts during tests
|
||||
jest.mock('../../config/connect', () => jest.fn().mockResolvedValue(true));
|
||||
|
||||
const dbModels = require('~/db/models');
|
||||
|
||||
// Disable console for tests
|
||||
logger.silent = true;
|
||||
|
||||
let mongoServer;
|
||||
let Prompt, PromptGroup, AclEntry, AccessRole, User, Group, Project;
|
||||
let promptFns, permissionService;
|
||||
let testUsers, testGroups, testRoles;
|
||||
|
||||
beforeAll(async () => {
|
||||
// Set up MongoDB memory server
|
||||
mongoServer = await MongoMemoryServer.create();
|
||||
const mongoUri = mongoServer.getUri();
|
||||
await mongoose.connect(mongoUri);
|
||||
|
||||
// Initialize models
|
||||
Prompt = dbModels.Prompt;
|
||||
PromptGroup = dbModels.PromptGroup;
|
||||
AclEntry = dbModels.AclEntry;
|
||||
AccessRole = dbModels.AccessRole;
|
||||
User = dbModels.User;
|
||||
Group = dbModels.Group;
|
||||
Project = dbModels.Project;
|
||||
|
||||
promptFns = require('~/models/Prompt');
|
||||
permissionService = require('~/server/services/PermissionService');
|
||||
|
||||
// Create test data
|
||||
await setupTestData();
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await mongoose.disconnect();
|
||||
await mongoServer.stop();
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
async function setupTestData() {
|
||||
// Create access roles for promptGroups
|
||||
testRoles = {
|
||||
viewer: await AccessRole.create({
|
||||
accessRoleId: AccessRoleIds.PROMPTGROUP_VIEWER,
|
||||
name: 'Viewer',
|
||||
description: 'Can view promptGroups',
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
permBits: PermissionBits.VIEW,
|
||||
}),
|
||||
editor: await AccessRole.create({
|
||||
accessRoleId: AccessRoleIds.PROMPTGROUP_EDITOR,
|
||||
name: 'Editor',
|
||||
description: 'Can view and edit promptGroups',
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
permBits: PermissionBits.VIEW | PermissionBits.EDIT,
|
||||
}),
|
||||
owner: await AccessRole.create({
|
||||
accessRoleId: AccessRoleIds.PROMPTGROUP_OWNER,
|
||||
name: 'Owner',
|
||||
description: 'Full control over promptGroups',
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
permBits:
|
||||
PermissionBits.VIEW | PermissionBits.EDIT | PermissionBits.DELETE | PermissionBits.SHARE,
|
||||
}),
|
||||
};
|
||||
|
||||
// Create test users
|
||||
testUsers = {
|
||||
owner: await User.create({
|
||||
name: 'Prompt Owner',
|
||||
email: 'owner@example.com',
|
||||
role: SystemRoles.USER,
|
||||
}),
|
||||
editor: await User.create({
|
||||
name: 'Prompt Editor',
|
||||
email: 'editor@example.com',
|
||||
role: SystemRoles.USER,
|
||||
}),
|
||||
viewer: await User.create({
|
||||
name: 'Prompt Viewer',
|
||||
email: 'viewer@example.com',
|
||||
role: SystemRoles.USER,
|
||||
}),
|
||||
admin: await User.create({
|
||||
name: 'Admin User',
|
||||
email: 'admin@example.com',
|
||||
role: SystemRoles.ADMIN,
|
||||
}),
|
||||
noAccess: await User.create({
|
||||
name: 'No Access User',
|
||||
email: 'noaccess@example.com',
|
||||
role: SystemRoles.USER,
|
||||
}),
|
||||
};
|
||||
|
||||
// Create test groups
|
||||
testGroups = {
|
||||
editors: await Group.create({
|
||||
name: 'Prompt Editors',
|
||||
description: 'Group with editor access',
|
||||
}),
|
||||
viewers: await Group.create({
|
||||
name: 'Prompt Viewers',
|
||||
description: 'Group with viewer access',
|
||||
}),
|
||||
};
|
||||
|
||||
await Project.create({
|
||||
name: 'Global',
|
||||
description: 'Global project',
|
||||
promptGroupIds: [],
|
||||
});
|
||||
}
|
||||
|
||||
describe('Prompt ACL Permissions', () => {
|
||||
describe('Creating Prompts with Permissions', () => {
|
||||
it('should grant owner permissions when creating a prompt', async () => {
|
||||
// First create a group
|
||||
const testGroup = await PromptGroup.create({
|
||||
name: 'Test Group',
|
||||
category: 'testing',
|
||||
author: testUsers.owner._id,
|
||||
authorName: testUsers.owner.name,
|
||||
productionId: new mongoose.Types.ObjectId(),
|
||||
});
|
||||
|
||||
const promptData = {
|
||||
prompt: {
|
||||
prompt: 'Test prompt content',
|
||||
name: 'Test Prompt',
|
||||
type: 'text',
|
||||
groupId: testGroup._id,
|
||||
},
|
||||
author: testUsers.owner._id,
|
||||
};
|
||||
|
||||
await promptFns.savePrompt(promptData);
|
||||
|
||||
// Manually grant permissions as would happen in the route
|
||||
await permissionService.grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: testUsers.owner._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testGroup._id,
|
||||
accessRoleId: AccessRoleIds.PROMPTGROUP_OWNER,
|
||||
grantedBy: testUsers.owner._id,
|
||||
});
|
||||
|
||||
// Check ACL entry
|
||||
const aclEntry = await AclEntry.findOne({
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testGroup._id,
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: testUsers.owner._id,
|
||||
});
|
||||
|
||||
expect(aclEntry).toBeTruthy();
|
||||
expect(aclEntry.permBits).toBe(testRoles.owner.permBits);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Accessing Prompts', () => {
|
||||
let testPromptGroup;
|
||||
|
||||
beforeEach(async () => {
|
||||
// Create a prompt group
|
||||
testPromptGroup = await PromptGroup.create({
|
||||
name: 'Test Group',
|
||||
author: testUsers.owner._id,
|
||||
authorName: testUsers.owner.name,
|
||||
productionId: new ObjectId(),
|
||||
});
|
||||
|
||||
// Create a prompt
|
||||
await Prompt.create({
|
||||
prompt: 'Test prompt for access control',
|
||||
name: 'Access Test Prompt',
|
||||
author: testUsers.owner._id,
|
||||
groupId: testPromptGroup._id,
|
||||
type: 'text',
|
||||
});
|
||||
|
||||
// Grant owner permissions
|
||||
await permissionService.grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: testUsers.owner._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testPromptGroup._id,
|
||||
accessRoleId: AccessRoleIds.PROMPTGROUP_OWNER,
|
||||
grantedBy: testUsers.owner._id,
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
await Prompt.deleteMany({});
|
||||
await PromptGroup.deleteMany({});
|
||||
await AclEntry.deleteMany({});
|
||||
});
|
||||
|
||||
it('owner should have full access to their prompt', async () => {
|
||||
const hasAccess = await permissionService.checkPermission({
|
||||
userId: testUsers.owner._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testPromptGroup._id,
|
||||
requiredPermission: PermissionBits.VIEW,
|
||||
});
|
||||
|
||||
expect(hasAccess).toBe(true);
|
||||
|
||||
const canEdit = await permissionService.checkPermission({
|
||||
userId: testUsers.owner._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testPromptGroup._id,
|
||||
requiredPermission: PermissionBits.EDIT,
|
||||
});
|
||||
|
||||
expect(canEdit).toBe(true);
|
||||
});
|
||||
|
||||
it('user with viewer role should only have view access', async () => {
|
||||
// Grant viewer permissions
|
||||
await permissionService.grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: testUsers.viewer._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testPromptGroup._id,
|
||||
accessRoleId: AccessRoleIds.PROMPTGROUP_VIEWER,
|
||||
grantedBy: testUsers.owner._id,
|
||||
});
|
||||
|
||||
const canView = await permissionService.checkPermission({
|
||||
userId: testUsers.viewer._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testPromptGroup._id,
|
||||
requiredPermission: PermissionBits.VIEW,
|
||||
});
|
||||
|
||||
const canEdit = await permissionService.checkPermission({
|
||||
userId: testUsers.viewer._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testPromptGroup._id,
|
||||
requiredPermission: PermissionBits.EDIT,
|
||||
});
|
||||
|
||||
expect(canView).toBe(true);
|
||||
expect(canEdit).toBe(false);
|
||||
});
|
||||
|
||||
it('user without permissions should have no access', async () => {
|
||||
const hasAccess = await permissionService.checkPermission({
|
||||
userId: testUsers.noAccess._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testPromptGroup._id,
|
||||
requiredPermission: PermissionBits.VIEW,
|
||||
});
|
||||
|
||||
expect(hasAccess).toBe(false);
|
||||
});
|
||||
|
||||
it('admin should have access regardless of permissions', async () => {
|
||||
// Admin users should work through normal permission system
|
||||
// The middleware layer handles admin bypass, not the permission service
|
||||
const hasAccess = await permissionService.checkPermission({
|
||||
userId: testUsers.admin._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testPromptGroup._id,
|
||||
requiredPermission: PermissionBits.VIEW,
|
||||
});
|
||||
|
||||
// Without explicit permissions, even admin won't have access at this layer
|
||||
expect(hasAccess).toBe(false);
|
||||
|
||||
// The actual admin bypass happens in the middleware layer (`canAccessPromptViaGroup`/`canAccessPromptGroupResource`)
|
||||
// which checks req.user.role === SystemRoles.ADMIN
|
||||
});
|
||||
});
|
||||
|
||||
describe('Group-based Access', () => {
|
||||
let testPromptGroup;
|
||||
|
||||
beforeEach(async () => {
|
||||
// Create a prompt group first
|
||||
testPromptGroup = await PromptGroup.create({
|
||||
name: 'Group Access Test Group',
|
||||
author: testUsers.owner._id,
|
||||
authorName: testUsers.owner.name,
|
||||
productionId: new ObjectId(),
|
||||
});
|
||||
|
||||
await Prompt.create({
|
||||
prompt: 'Group access test prompt',
|
||||
name: 'Group Test',
|
||||
author: testUsers.owner._id,
|
||||
groupId: testPromptGroup._id,
|
||||
type: 'text',
|
||||
});
|
||||
|
||||
// Add users to groups
|
||||
await User.findByIdAndUpdate(testUsers.editor._id, {
|
||||
$push: { groups: testGroups.editors._id },
|
||||
});
|
||||
|
||||
await User.findByIdAndUpdate(testUsers.viewer._id, {
|
||||
$push: { groups: testGroups.viewers._id },
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
await Prompt.deleteMany({});
|
||||
await AclEntry.deleteMany({});
|
||||
await User.updateMany({}, { $set: { groups: [] } });
|
||||
});
|
||||
|
||||
it('group members should inherit group permissions', async () => {
|
||||
// Create a prompt group
|
||||
const testPromptGroup = await PromptGroup.create({
|
||||
name: 'Group Test Group',
|
||||
author: testUsers.owner._id,
|
||||
authorName: testUsers.owner.name,
|
||||
productionId: new ObjectId(),
|
||||
});
|
||||
|
||||
const { addUserToGroup } = require('~/models');
|
||||
await addUserToGroup(testUsers.editor._id, testGroups.editors._id);
|
||||
|
||||
const prompt = await promptFns.savePrompt({
|
||||
author: testUsers.owner._id,
|
||||
prompt: {
|
||||
prompt: 'Group test prompt',
|
||||
name: 'Group Test',
|
||||
groupId: testPromptGroup._id,
|
||||
type: 'text',
|
||||
},
|
||||
});
|
||||
|
||||
// Check if savePrompt returned an error
|
||||
if (!prompt || !prompt.prompt) {
|
||||
throw new Error(`Failed to save prompt: ${prompt?.message || 'Unknown error'}`);
|
||||
}
|
||||
|
||||
// Grant edit permissions to the group
|
||||
await permissionService.grantPermission({
|
||||
principalType: PrincipalType.GROUP,
|
||||
principalId: testGroups.editors._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testPromptGroup._id,
|
||||
accessRoleId: AccessRoleIds.PROMPTGROUP_EDITOR,
|
||||
grantedBy: testUsers.owner._id,
|
||||
});
|
||||
|
||||
// Check if group member has access
|
||||
const hasAccess = await permissionService.checkPermission({
|
||||
userId: testUsers.editor._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testPromptGroup._id,
|
||||
requiredPermission: PermissionBits.EDIT,
|
||||
});
|
||||
|
||||
expect(hasAccess).toBe(true);
|
||||
|
||||
// Check that non-member doesn't have access
|
||||
const nonMemberAccess = await permissionService.checkPermission({
|
||||
userId: testUsers.viewer._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testPromptGroup._id,
|
||||
requiredPermission: PermissionBits.EDIT,
|
||||
});
|
||||
|
||||
expect(nonMemberAccess).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Public Access', () => {
|
||||
let publicPromptGroup, privatePromptGroup;
|
||||
|
||||
beforeEach(async () => {
|
||||
// Create separate prompt groups for public and private access
|
||||
publicPromptGroup = await PromptGroup.create({
|
||||
name: 'Public Access Test Group',
|
||||
author: testUsers.owner._id,
|
||||
authorName: testUsers.owner.name,
|
||||
productionId: new ObjectId(),
|
||||
});
|
||||
|
||||
privatePromptGroup = await PromptGroup.create({
|
||||
name: 'Private Access Test Group',
|
||||
author: testUsers.owner._id,
|
||||
authorName: testUsers.owner.name,
|
||||
productionId: new ObjectId(),
|
||||
});
|
||||
|
||||
// Create prompts in their respective groups
|
||||
await Prompt.create({
|
||||
prompt: 'Public prompt',
|
||||
name: 'Public',
|
||||
author: testUsers.owner._id,
|
||||
groupId: publicPromptGroup._id,
|
||||
type: 'text',
|
||||
});
|
||||
|
||||
await Prompt.create({
|
||||
prompt: 'Private prompt',
|
||||
name: 'Private',
|
||||
author: testUsers.owner._id,
|
||||
groupId: privatePromptGroup._id,
|
||||
type: 'text',
|
||||
});
|
||||
|
||||
// Grant public view access to publicPromptGroup
|
||||
await permissionService.grantPermission({
|
||||
principalType: PrincipalType.PUBLIC,
|
||||
principalId: null,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: publicPromptGroup._id,
|
||||
accessRoleId: AccessRoleIds.PROMPTGROUP_VIEWER,
|
||||
grantedBy: testUsers.owner._id,
|
||||
});
|
||||
|
||||
// Grant only owner access to privatePromptGroup
|
||||
await permissionService.grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: testUsers.owner._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: privatePromptGroup._id,
|
||||
accessRoleId: AccessRoleIds.PROMPTGROUP_OWNER,
|
||||
grantedBy: testUsers.owner._id,
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
await Prompt.deleteMany({});
|
||||
await PromptGroup.deleteMany({});
|
||||
await AclEntry.deleteMany({});
|
||||
});
|
||||
|
||||
it('public prompt should be accessible to any user', async () => {
|
||||
const hasAccess = await permissionService.checkPermission({
|
||||
userId: testUsers.noAccess._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: publicPromptGroup._id,
|
||||
requiredPermission: PermissionBits.VIEW,
|
||||
includePublic: true,
|
||||
});
|
||||
|
||||
expect(hasAccess).toBe(true);
|
||||
});
|
||||
|
||||
it('private prompt should not be accessible to unauthorized users', async () => {
|
||||
const hasAccess = await permissionService.checkPermission({
|
||||
userId: testUsers.noAccess._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: privatePromptGroup._id,
|
||||
requiredPermission: PermissionBits.VIEW,
|
||||
includePublic: true,
|
||||
});
|
||||
|
||||
expect(hasAccess).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Prompt Deletion', () => {
|
||||
let testPromptGroup;
|
||||
|
||||
it('should remove ACL entries when prompt is deleted', async () => {
|
||||
testPromptGroup = await PromptGroup.create({
|
||||
name: 'Deletion Test Group',
|
||||
author: testUsers.owner._id,
|
||||
authorName: testUsers.owner.name,
|
||||
productionId: new ObjectId(),
|
||||
});
|
||||
|
||||
const prompt = await promptFns.savePrompt({
|
||||
author: testUsers.owner._id,
|
||||
prompt: {
|
||||
prompt: 'To be deleted',
|
||||
name: 'Delete Test',
|
||||
groupId: testPromptGroup._id,
|
||||
type: 'text',
|
||||
},
|
||||
});
|
||||
|
||||
// Check if savePrompt returned an error
|
||||
if (!prompt || !prompt.prompt) {
|
||||
throw new Error(`Failed to save prompt: ${prompt?.message || 'Unknown error'}`);
|
||||
}
|
||||
|
||||
const testPromptId = prompt.prompt._id;
|
||||
const promptGroupId = testPromptGroup._id;
|
||||
|
||||
// Grant permission
|
||||
await permissionService.grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: testUsers.owner._id,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testPromptGroup._id,
|
||||
accessRoleId: AccessRoleIds.PROMPTGROUP_OWNER,
|
||||
grantedBy: testUsers.owner._id,
|
||||
});
|
||||
|
||||
// Verify ACL entry exists
|
||||
const beforeDelete = await AclEntry.find({
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testPromptGroup._id,
|
||||
});
|
||||
expect(beforeDelete).toHaveLength(1);
|
||||
|
||||
// Delete the prompt
|
||||
await promptFns.deletePrompt({
|
||||
promptId: testPromptId,
|
||||
groupId: promptGroupId,
|
||||
author: testUsers.owner._id,
|
||||
role: SystemRoles.USER,
|
||||
});
|
||||
|
||||
// Verify ACL entries are removed
|
||||
const aclEntries = await AclEntry.find({
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: testPromptGroup._id,
|
||||
});
|
||||
|
||||
expect(aclEntries).toHaveLength(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Backwards Compatibility', () => {
|
||||
it('should handle prompts without ACL entries gracefully', async () => {
|
||||
// Create a prompt group first
|
||||
const promptGroup = await PromptGroup.create({
|
||||
name: 'Legacy Test Group',
|
||||
author: testUsers.owner._id,
|
||||
authorName: testUsers.owner.name,
|
||||
productionId: new ObjectId(),
|
||||
});
|
||||
|
||||
// Create a prompt without ACL entries (legacy prompt)
|
||||
const legacyPrompt = await Prompt.create({
|
||||
prompt: 'Legacy prompt without ACL',
|
||||
name: 'Legacy',
|
||||
author: testUsers.owner._id,
|
||||
groupId: promptGroup._id,
|
||||
type: 'text',
|
||||
});
|
||||
|
||||
// The system should handle this gracefully
|
||||
const prompt = await promptFns.getPrompt({ _id: legacyPrompt._id });
|
||||
expect(prompt).toBeTruthy();
|
||||
expect(prompt._id.toString()).toBe(legacyPrompt._id.toString());
|
||||
});
|
||||
});
|
||||
});
|
||||
280
api/models/PromptGroupMigration.spec.js
Normal file
280
api/models/PromptGroupMigration.spec.js
Normal file
@ -0,0 +1,280 @@
|
||||
const mongoose = require('mongoose');
|
||||
const { ObjectId } = require('mongodb');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
const {
|
||||
Constants,
|
||||
ResourceType,
|
||||
AccessRoleIds,
|
||||
PrincipalType,
|
||||
PrincipalModel,
|
||||
PermissionBits,
|
||||
} = require('librechat-data-provider');
|
||||
|
||||
// Mock the config/connect module to prevent connection attempts during tests
|
||||
jest.mock('../../config/connect', () => jest.fn().mockResolvedValue(true));
|
||||
|
||||
// Disable console for tests
|
||||
logger.silent = true;
|
||||
|
||||
describe('PromptGroup Migration Script', () => {
|
||||
let mongoServer;
|
||||
let Prompt, PromptGroup, AclEntry, AccessRole, User, Project;
|
||||
let migrateToPromptGroupPermissions;
|
||||
let testOwner, testProject;
|
||||
let ownerRole, viewerRole;
|
||||
|
||||
beforeAll(async () => {
|
||||
// Set up MongoDB memory server
|
||||
mongoServer = await MongoMemoryServer.create();
|
||||
const mongoUri = mongoServer.getUri();
|
||||
await mongoose.connect(mongoUri);
|
||||
|
||||
// Initialize models
|
||||
const dbModels = require('~/db/models');
|
||||
Prompt = dbModels.Prompt;
|
||||
PromptGroup = dbModels.PromptGroup;
|
||||
AclEntry = dbModels.AclEntry;
|
||||
AccessRole = dbModels.AccessRole;
|
||||
User = dbModels.User;
|
||||
Project = dbModels.Project;
|
||||
|
||||
// Create test user
|
||||
testOwner = await User.create({
|
||||
name: 'Test Owner',
|
||||
email: 'owner@test.com',
|
||||
role: 'USER',
|
||||
});
|
||||
|
||||
// Create test project with the proper name
|
||||
const projectName = Constants.GLOBAL_PROJECT_NAME || 'instance';
|
||||
testProject = await Project.create({
|
||||
name: projectName,
|
||||
description: 'Global project',
|
||||
promptGroupIds: [],
|
||||
});
|
||||
|
||||
// Create promptGroup access roles
|
||||
ownerRole = await AccessRole.create({
|
||||
accessRoleId: AccessRoleIds.PROMPTGROUP_OWNER,
|
||||
name: 'Owner',
|
||||
description: 'Full control over promptGroups',
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
permBits:
|
||||
PermissionBits.VIEW | PermissionBits.EDIT | PermissionBits.DELETE | PermissionBits.SHARE,
|
||||
});
|
||||
|
||||
viewerRole = await AccessRole.create({
|
||||
accessRoleId: AccessRoleIds.PROMPTGROUP_VIEWER,
|
||||
name: 'Viewer',
|
||||
description: 'Can view promptGroups',
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
permBits: PermissionBits.VIEW,
|
||||
});
|
||||
|
||||
await AccessRole.create({
|
||||
accessRoleId: AccessRoleIds.PROMPTGROUP_EDITOR,
|
||||
name: 'Editor',
|
||||
description: 'Can view and edit promptGroups',
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
permBits: PermissionBits.VIEW | PermissionBits.EDIT,
|
||||
});
|
||||
|
||||
// Import migration function
|
||||
const migration = require('../../config/migrate-prompt-permissions');
|
||||
migrateToPromptGroupPermissions = migration.migrateToPromptGroupPermissions;
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await mongoose.disconnect();
|
||||
await mongoServer.stop();
|
||||
});
|
||||
|
||||
beforeEach(async () => {
|
||||
// Clean up before each test
|
||||
await Prompt.deleteMany({});
|
||||
await PromptGroup.deleteMany({});
|
||||
await AclEntry.deleteMany({});
|
||||
// Reset the project's promptGroupIds array
|
||||
testProject.promptGroupIds = [];
|
||||
await testProject.save();
|
||||
});
|
||||
|
||||
it('should categorize promptGroups correctly in dry run', async () => {
|
||||
// Create global prompt group (in Global project)
|
||||
const globalPromptGroup = await PromptGroup.create({
|
||||
name: 'Global Group',
|
||||
author: testOwner._id,
|
||||
authorName: testOwner.name,
|
||||
productionId: new ObjectId(),
|
||||
});
|
||||
|
||||
// Create private prompt group (not in any project)
|
||||
await PromptGroup.create({
|
||||
name: 'Private Group',
|
||||
author: testOwner._id,
|
||||
authorName: testOwner.name,
|
||||
productionId: new ObjectId(),
|
||||
});
|
||||
|
||||
// Add global group to project's promptGroupIds array
|
||||
testProject.promptGroupIds = [globalPromptGroup._id];
|
||||
await testProject.save();
|
||||
|
||||
const result = await migrateToPromptGroupPermissions({ dryRun: true });
|
||||
|
||||
expect(result.dryRun).toBe(true);
|
||||
expect(result.summary.total).toBe(2);
|
||||
expect(result.summary.globalViewAccess).toBe(1);
|
||||
expect(result.summary.privateGroups).toBe(1);
|
||||
});
|
||||
|
||||
it('should grant appropriate permissions during migration', async () => {
|
||||
// Create prompt groups
|
||||
const globalPromptGroup = await PromptGroup.create({
|
||||
name: 'Global Group',
|
||||
author: testOwner._id,
|
||||
authorName: testOwner.name,
|
||||
productionId: new ObjectId(),
|
||||
});
|
||||
|
||||
const privatePromptGroup = await PromptGroup.create({
|
||||
name: 'Private Group',
|
||||
author: testOwner._id,
|
||||
authorName: testOwner.name,
|
||||
productionId: new ObjectId(),
|
||||
});
|
||||
|
||||
// Add global group to project's promptGroupIds array
|
||||
testProject.promptGroupIds = [globalPromptGroup._id];
|
||||
await testProject.save();
|
||||
|
||||
const result = await migrateToPromptGroupPermissions({ dryRun: false });
|
||||
|
||||
expect(result.migrated).toBe(2);
|
||||
expect(result.errors).toBe(0);
|
||||
expect(result.ownerGrants).toBe(2);
|
||||
expect(result.publicViewGrants).toBe(1);
|
||||
|
||||
// Check global promptGroup permissions
|
||||
const globalOwnerEntry = await AclEntry.findOne({
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: globalPromptGroup._id,
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: testOwner._id,
|
||||
});
|
||||
expect(globalOwnerEntry).toBeTruthy();
|
||||
expect(globalOwnerEntry.permBits).toBe(ownerRole.permBits);
|
||||
|
||||
const globalPublicEntry = await AclEntry.findOne({
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: globalPromptGroup._id,
|
||||
principalType: PrincipalType.PUBLIC,
|
||||
});
|
||||
expect(globalPublicEntry).toBeTruthy();
|
||||
expect(globalPublicEntry.permBits).toBe(viewerRole.permBits);
|
||||
|
||||
// Check private promptGroup permissions
|
||||
const privateOwnerEntry = await AclEntry.findOne({
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: privatePromptGroup._id,
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: testOwner._id,
|
||||
});
|
||||
expect(privateOwnerEntry).toBeTruthy();
|
||||
expect(privateOwnerEntry.permBits).toBe(ownerRole.permBits);
|
||||
|
||||
const privatePublicEntry = await AclEntry.findOne({
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: privatePromptGroup._id,
|
||||
principalType: PrincipalType.PUBLIC,
|
||||
});
|
||||
expect(privatePublicEntry).toBeNull();
|
||||
});
|
||||
|
||||
it('should skip promptGroups that already have ACL entries', async () => {
|
||||
// Create prompt groups
|
||||
const promptGroup1 = await PromptGroup.create({
|
||||
name: 'Group 1',
|
||||
author: testOwner._id,
|
||||
authorName: testOwner.name,
|
||||
productionId: new ObjectId(),
|
||||
});
|
||||
|
||||
const promptGroup2 = await PromptGroup.create({
|
||||
name: 'Group 2',
|
||||
author: testOwner._id,
|
||||
authorName: testOwner.name,
|
||||
productionId: new ObjectId(),
|
||||
});
|
||||
|
||||
// Grant permission to one promptGroup manually (simulating it already has ACL)
|
||||
await AclEntry.create({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: testOwner._id,
|
||||
principalModel: PrincipalModel.USER,
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: promptGroup1._id,
|
||||
permBits: ownerRole.permBits,
|
||||
roleId: ownerRole._id,
|
||||
grantedBy: testOwner._id,
|
||||
grantedAt: new Date(),
|
||||
});
|
||||
|
||||
const result = await migrateToPromptGroupPermissions({ dryRun: false });
|
||||
|
||||
// Should only migrate promptGroup2, skip promptGroup1
|
||||
expect(result.migrated).toBe(1);
|
||||
expect(result.errors).toBe(0);
|
||||
|
||||
// Verify promptGroup2 now has permissions
|
||||
const group2Entry = await AclEntry.findOne({
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: promptGroup2._id,
|
||||
});
|
||||
expect(group2Entry).toBeTruthy();
|
||||
});
|
||||
|
||||
it('should handle promptGroups with prompts correctly', async () => {
|
||||
// Create a promptGroup with some prompts
|
||||
const promptGroup = await PromptGroup.create({
|
||||
name: 'Group with Prompts',
|
||||
author: testOwner._id,
|
||||
authorName: testOwner.name,
|
||||
productionId: new ObjectId(),
|
||||
});
|
||||
|
||||
// Create some prompts in this group
|
||||
await Prompt.create({
|
||||
prompt: 'First prompt',
|
||||
author: testOwner._id,
|
||||
groupId: promptGroup._id,
|
||||
type: 'text',
|
||||
});
|
||||
|
||||
await Prompt.create({
|
||||
prompt: 'Second prompt',
|
||||
author: testOwner._id,
|
||||
groupId: promptGroup._id,
|
||||
type: 'text',
|
||||
});
|
||||
|
||||
const result = await migrateToPromptGroupPermissions({ dryRun: false });
|
||||
|
||||
expect(result.migrated).toBe(1);
|
||||
expect(result.errors).toBe(0);
|
||||
|
||||
// Verify the promptGroup has permissions
|
||||
const groupEntry = await AclEntry.findOne({
|
||||
resourceType: ResourceType.PROMPTGROUP,
|
||||
resourceId: promptGroup._id,
|
||||
});
|
||||
expect(groupEntry).toBeTruthy();
|
||||
|
||||
// Verify no prompt-level permissions were created
|
||||
const promptEntries = await AclEntry.find({
|
||||
resourceType: 'prompt',
|
||||
});
|
||||
expect(promptEntries).toHaveLength(0);
|
||||
});
|
||||
});
|
||||
256
api/models/Role.js
Normal file
256
api/models/Role.js
Normal file
@ -0,0 +1,256 @@
|
||||
const {
|
||||
CacheKeys,
|
||||
SystemRoles,
|
||||
roleDefaults,
|
||||
permissionsSchema,
|
||||
removeNullishValues,
|
||||
} = require('librechat-data-provider');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const getLogStores = require('~/cache/getLogStores');
|
||||
const { Role } = require('~/db/models');
|
||||
|
||||
/**
|
||||
* Retrieve a role by name and convert the found role document to a plain object.
|
||||
* If the role with the given name doesn't exist and the name is a system defined role,
|
||||
* create it and return the lean version.
|
||||
*
|
||||
* @param {string} roleName - The name of the role to find or create.
|
||||
* @param {string|string[]} [fieldsToSelect] - The fields to include or exclude in the returned document.
|
||||
* @returns {Promise<IRole>} Role document.
|
||||
*/
|
||||
const getRoleByName = async function (roleName, fieldsToSelect = null) {
|
||||
const cache = getLogStores(CacheKeys.ROLES);
|
||||
try {
|
||||
const cachedRole = await cache.get(roleName);
|
||||
if (cachedRole) {
|
||||
return cachedRole;
|
||||
}
|
||||
let query = Role.findOne({ name: roleName });
|
||||
if (fieldsToSelect) {
|
||||
query = query.select(fieldsToSelect);
|
||||
}
|
||||
let role = await query.lean().exec();
|
||||
|
||||
if (!role && SystemRoles[roleName]) {
|
||||
role = await new Role(roleDefaults[roleName]).save();
|
||||
await cache.set(roleName, role);
|
||||
return role.toObject();
|
||||
}
|
||||
await cache.set(roleName, role);
|
||||
return role;
|
||||
} catch (error) {
|
||||
throw new Error(`Failed to retrieve or create role: ${error.message}`);
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Update role values by name.
|
||||
*
|
||||
* @param {string} roleName - The name of the role to update.
|
||||
* @param {Partial<TRole>} updates - The fields to update.
|
||||
* @returns {Promise<TRole>} Updated role document.
|
||||
*/
|
||||
const updateRoleByName = async function (roleName, updates) {
|
||||
const cache = getLogStores(CacheKeys.ROLES);
|
||||
try {
|
||||
const role = await Role.findOneAndUpdate(
|
||||
{ name: roleName },
|
||||
{ $set: updates },
|
||||
{ new: true, lean: true },
|
||||
)
|
||||
.select('-__v')
|
||||
.lean()
|
||||
.exec();
|
||||
await cache.set(roleName, role);
|
||||
return role;
|
||||
} catch (error) {
|
||||
throw new Error(`Failed to update role: ${error.message}`);
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Updates access permissions for a specific role and multiple permission types.
|
||||
* @param {string} roleName - The role to update.
|
||||
* @param {Object.<PermissionTypes, Object.<Permissions, boolean>>} permissionsUpdate - Permissions to update and their values.
|
||||
* @param {IRole} [roleData] - Optional role data to use instead of fetching from the database.
|
||||
*/
|
||||
async function updateAccessPermissions(roleName, permissionsUpdate, roleData) {
|
||||
// Filter and clean the permission updates based on our schema definition.
|
||||
const updates = {};
|
||||
for (const [permissionType, permissions] of Object.entries(permissionsUpdate)) {
|
||||
if (permissionsSchema.shape && permissionsSchema.shape[permissionType]) {
|
||||
updates[permissionType] = removeNullishValues(permissions);
|
||||
}
|
||||
}
|
||||
if (!Object.keys(updates).length) {
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const role = roleData ?? (await getRoleByName(roleName));
|
||||
if (!role) {
|
||||
return;
|
||||
}
|
||||
|
||||
const currentPermissions = role.permissions || {};
|
||||
const updatedPermissions = { ...currentPermissions };
|
||||
let hasChanges = false;
|
||||
|
||||
const unsetFields = {};
|
||||
const permissionTypes = Object.keys(permissionsSchema.shape || {});
|
||||
for (const permType of permissionTypes) {
|
||||
if (role[permType] && typeof role[permType] === 'object') {
|
||||
logger.info(
|
||||
`Migrating '${roleName}' role from old schema: found '${permType}' at top level`,
|
||||
);
|
||||
|
||||
updatedPermissions[permType] = {
|
||||
...updatedPermissions[permType],
|
||||
...role[permType],
|
||||
};
|
||||
|
||||
unsetFields[permType] = 1;
|
||||
hasChanges = true;
|
||||
}
|
||||
}
|
||||
|
||||
for (const [permissionType, permissions] of Object.entries(updates)) {
|
||||
const currentTypePermissions = currentPermissions[permissionType] || {};
|
||||
updatedPermissions[permissionType] = { ...currentTypePermissions };
|
||||
|
||||
for (const [permission, value] of Object.entries(permissions)) {
|
||||
if (currentTypePermissions[permission] !== value) {
|
||||
updatedPermissions[permissionType][permission] = value;
|
||||
hasChanges = true;
|
||||
logger.info(
|
||||
`Updating '${roleName}' role permission '${permissionType}' '${permission}' from ${currentTypePermissions[permission]} to: ${value}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (hasChanges) {
|
||||
const updateObj = { permissions: updatedPermissions };
|
||||
|
||||
if (Object.keys(unsetFields).length > 0) {
|
||||
logger.info(
|
||||
`Unsetting old schema fields for '${roleName}' role: ${Object.keys(unsetFields).join(', ')}`,
|
||||
);
|
||||
|
||||
try {
|
||||
await Role.updateOne(
|
||||
{ name: roleName },
|
||||
{
|
||||
$set: updateObj,
|
||||
$unset: unsetFields,
|
||||
},
|
||||
);
|
||||
|
||||
const cache = getLogStores(CacheKeys.ROLES);
|
||||
const updatedRole = await Role.findOne({ name: roleName }).select('-__v').lean().exec();
|
||||
await cache.set(roleName, updatedRole);
|
||||
|
||||
logger.info(`Updated role '${roleName}' and removed old schema fields`);
|
||||
} catch (updateError) {
|
||||
logger.error(`Error during role migration update: ${updateError.message}`);
|
||||
throw updateError;
|
||||
}
|
||||
} else {
|
||||
// Standard update if no migration needed
|
||||
await updateRoleByName(roleName, updateObj);
|
||||
}
|
||||
|
||||
logger.info(`Updated '${roleName}' role permissions`);
|
||||
} else {
|
||||
logger.info(`No changes needed for '${roleName}' role permissions`);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error(`Failed to update ${roleName} role permissions:`, error);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Migrates roles from old schema to new schema structure.
|
||||
* This can be called directly to fix existing roles.
|
||||
*
|
||||
* @param {string} [roleName] - Optional specific role to migrate. If not provided, migrates all roles.
|
||||
* @returns {Promise<number>} Number of roles migrated.
|
||||
*/
|
||||
const migrateRoleSchema = async function (roleName) {
|
||||
try {
|
||||
// Get roles to migrate
|
||||
let roles;
|
||||
if (roleName) {
|
||||
const role = await Role.findOne({ name: roleName });
|
||||
roles = role ? [role] : [];
|
||||
} else {
|
||||
roles = await Role.find({});
|
||||
}
|
||||
|
||||
logger.info(`Migrating ${roles.length} roles to new schema structure`);
|
||||
let migratedCount = 0;
|
||||
|
||||
for (const role of roles) {
|
||||
const permissionTypes = Object.keys(permissionsSchema.shape || {});
|
||||
const unsetFields = {};
|
||||
let hasOldSchema = false;
|
||||
|
||||
// Check for old schema fields
|
||||
for (const permType of permissionTypes) {
|
||||
if (role[permType] && typeof role[permType] === 'object') {
|
||||
hasOldSchema = true;
|
||||
|
||||
// Ensure permissions object exists
|
||||
role.permissions = role.permissions || {};
|
||||
|
||||
// Migrate permissions from old location to new
|
||||
role.permissions[permType] = {
|
||||
...role.permissions[permType],
|
||||
...role[permType],
|
||||
};
|
||||
|
||||
// Mark field for removal
|
||||
unsetFields[permType] = 1;
|
||||
}
|
||||
}
|
||||
|
||||
if (hasOldSchema) {
|
||||
try {
|
||||
logger.info(`Migrating role '${role.name}' from old schema structure`);
|
||||
|
||||
// Simple update operation
|
||||
await Role.updateOne(
|
||||
{ _id: role._id },
|
||||
{
|
||||
$set: { permissions: role.permissions },
|
||||
$unset: unsetFields,
|
||||
},
|
||||
);
|
||||
|
||||
// Refresh cache
|
||||
const cache = getLogStores(CacheKeys.ROLES);
|
||||
const updatedRole = await Role.findById(role._id).lean().exec();
|
||||
await cache.set(role.name, updatedRole);
|
||||
|
||||
migratedCount++;
|
||||
logger.info(`Migrated role '${role.name}'`);
|
||||
} catch (error) {
|
||||
logger.error(`Failed to migrate role '${role.name}': ${error.message}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.info(`Migration complete: ${migratedCount} roles migrated`);
|
||||
return migratedCount;
|
||||
} catch (error) {
|
||||
logger.error(`Role schema migration failed: ${error.message}`);
|
||||
throw error;
|
||||
}
|
||||
};
|
||||
|
||||
module.exports = {
|
||||
getRoleByName,
|
||||
updateRoleByName,
|
||||
migrateRoleSchema,
|
||||
updateAccessPermissions,
|
||||
};
|
||||
405
api/models/Role.spec.js
Normal file
405
api/models/Role.spec.js
Normal file
@ -0,0 +1,405 @@
|
||||
const mongoose = require('mongoose');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
const {
|
||||
SystemRoles,
|
||||
Permissions,
|
||||
roleDefaults,
|
||||
PermissionTypes,
|
||||
} = require('librechat-data-provider');
|
||||
const { getRoleByName, updateAccessPermissions } = require('~/models/Role');
|
||||
const getLogStores = require('~/cache/getLogStores');
|
||||
const { initializeRoles } = require('~/models');
|
||||
const { Role } = require('~/db/models');
|
||||
|
||||
// Mock the cache
|
||||
jest.mock('~/cache/getLogStores', () =>
|
||||
jest.fn().mockReturnValue({
|
||||
get: jest.fn(),
|
||||
set: jest.fn(),
|
||||
del: jest.fn(),
|
||||
}),
|
||||
);
|
||||
|
||||
let mongoServer;
|
||||
|
||||
beforeAll(async () => {
|
||||
mongoServer = await MongoMemoryServer.create();
|
||||
const mongoUri = mongoServer.getUri();
|
||||
await mongoose.connect(mongoUri);
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await mongoose.disconnect();
|
||||
await mongoServer.stop();
|
||||
});
|
||||
|
||||
beforeEach(async () => {
|
||||
await Role.deleteMany({});
|
||||
getLogStores.mockClear();
|
||||
});
|
||||
|
||||
describe('updateAccessPermissions', () => {
|
||||
it('should update permissions when changes are needed', async () => {
|
||||
await new Role({
|
||||
name: SystemRoles.USER,
|
||||
permissions: {
|
||||
[PermissionTypes.PROMPTS]: {
|
||||
CREATE: true,
|
||||
USE: true,
|
||||
SHARE: false,
|
||||
},
|
||||
},
|
||||
}).save();
|
||||
|
||||
await updateAccessPermissions(SystemRoles.USER, {
|
||||
[PermissionTypes.PROMPTS]: {
|
||||
CREATE: true,
|
||||
USE: true,
|
||||
SHARE: true,
|
||||
},
|
||||
});
|
||||
|
||||
const updatedRole = await getRoleByName(SystemRoles.USER);
|
||||
expect(updatedRole.permissions[PermissionTypes.PROMPTS]).toEqual({
|
||||
CREATE: true,
|
||||
USE: true,
|
||||
SHARE: true,
|
||||
});
|
||||
});
|
||||
|
||||
it('should not update permissions when no changes are needed', async () => {
|
||||
await new Role({
|
||||
name: SystemRoles.USER,
|
||||
permissions: {
|
||||
[PermissionTypes.PROMPTS]: {
|
||||
CREATE: true,
|
||||
USE: true,
|
||||
SHARE: false,
|
||||
},
|
||||
},
|
||||
}).save();
|
||||
|
||||
await updateAccessPermissions(SystemRoles.USER, {
|
||||
[PermissionTypes.PROMPTS]: {
|
||||
CREATE: true,
|
||||
USE: true,
|
||||
SHARE: false,
|
||||
},
|
||||
});
|
||||
|
||||
const updatedRole = await getRoleByName(SystemRoles.USER);
|
||||
expect(updatedRole.permissions[PermissionTypes.PROMPTS]).toEqual({
|
||||
CREATE: true,
|
||||
USE: true,
|
||||
SHARE: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle non-existent roles', async () => {
|
||||
await updateAccessPermissions('NON_EXISTENT_ROLE', {
|
||||
[PermissionTypes.PROMPTS]: { CREATE: true },
|
||||
});
|
||||
const role = await Role.findOne({ name: 'NON_EXISTENT_ROLE' });
|
||||
expect(role).toBeNull();
|
||||
});
|
||||
|
||||
it('should update only specified permissions', async () => {
|
||||
await new Role({
|
||||
name: SystemRoles.USER,
|
||||
permissions: {
|
||||
[PermissionTypes.PROMPTS]: {
|
||||
CREATE: true,
|
||||
USE: true,
|
||||
SHARE: false,
|
||||
},
|
||||
},
|
||||
}).save();
|
||||
|
||||
await updateAccessPermissions(SystemRoles.USER, {
|
||||
[PermissionTypes.PROMPTS]: { SHARE: true },
|
||||
});
|
||||
|
||||
const updatedRole = await getRoleByName(SystemRoles.USER);
|
||||
expect(updatedRole.permissions[PermissionTypes.PROMPTS]).toEqual({
|
||||
CREATE: true,
|
||||
USE: true,
|
||||
SHARE: true,
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle partial updates', async () => {
|
||||
await new Role({
|
||||
name: SystemRoles.USER,
|
||||
permissions: {
|
||||
[PermissionTypes.PROMPTS]: {
|
||||
CREATE: true,
|
||||
USE: true,
|
||||
SHARE: false,
|
||||
},
|
||||
},
|
||||
}).save();
|
||||
|
||||
await updateAccessPermissions(SystemRoles.USER, {
|
||||
[PermissionTypes.PROMPTS]: { USE: false },
|
||||
});
|
||||
|
||||
const updatedRole = await getRoleByName(SystemRoles.USER);
|
||||
expect(updatedRole.permissions[PermissionTypes.PROMPTS]).toEqual({
|
||||
CREATE: true,
|
||||
USE: false,
|
||||
SHARE: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('should update multiple permission types at once', async () => {
|
||||
await new Role({
|
||||
name: SystemRoles.USER,
|
||||
permissions: {
|
||||
[PermissionTypes.PROMPTS]: { CREATE: true, USE: true, SHARE: false },
|
||||
[PermissionTypes.BOOKMARKS]: { USE: true },
|
||||
},
|
||||
}).save();
|
||||
|
||||
await updateAccessPermissions(SystemRoles.USER, {
|
||||
[PermissionTypes.PROMPTS]: { USE: false, SHARE: true },
|
||||
[PermissionTypes.BOOKMARKS]: { USE: false },
|
||||
});
|
||||
|
||||
const updatedRole = await getRoleByName(SystemRoles.USER);
|
||||
expect(updatedRole.permissions[PermissionTypes.PROMPTS]).toEqual({
|
||||
CREATE: true,
|
||||
USE: false,
|
||||
SHARE: true,
|
||||
});
|
||||
expect(updatedRole.permissions[PermissionTypes.BOOKMARKS]).toEqual({ USE: false });
|
||||
});
|
||||
|
||||
it('should handle updates for a single permission type', async () => {
|
||||
await new Role({
|
||||
name: SystemRoles.USER,
|
||||
permissions: {
|
||||
[PermissionTypes.PROMPTS]: { CREATE: true, USE: true, SHARE: false },
|
||||
},
|
||||
}).save();
|
||||
|
||||
await updateAccessPermissions(SystemRoles.USER, {
|
||||
[PermissionTypes.PROMPTS]: { USE: false, SHARE: true },
|
||||
});
|
||||
|
||||
const updatedRole = await getRoleByName(SystemRoles.USER);
|
||||
expect(updatedRole.permissions[PermissionTypes.PROMPTS]).toEqual({
|
||||
CREATE: true,
|
||||
USE: false,
|
||||
SHARE: true,
|
||||
});
|
||||
});
|
||||
|
||||
it('should update MULTI_CONVO permissions', async () => {
|
||||
await new Role({
|
||||
name: SystemRoles.USER,
|
||||
permissions: {
|
||||
[PermissionTypes.MULTI_CONVO]: { USE: false },
|
||||
},
|
||||
}).save();
|
||||
|
||||
await updateAccessPermissions(SystemRoles.USER, {
|
||||
[PermissionTypes.MULTI_CONVO]: { USE: true },
|
||||
});
|
||||
|
||||
const updatedRole = await getRoleByName(SystemRoles.USER);
|
||||
expect(updatedRole.permissions[PermissionTypes.MULTI_CONVO]).toEqual({ USE: true });
|
||||
});
|
||||
|
||||
it('should update MULTI_CONVO permissions along with other permission types', async () => {
|
||||
await new Role({
|
||||
name: SystemRoles.USER,
|
||||
permissions: {
|
||||
[PermissionTypes.PROMPTS]: { CREATE: true, USE: true, SHARE: false },
|
||||
[PermissionTypes.MULTI_CONVO]: { USE: false },
|
||||
},
|
||||
}).save();
|
||||
|
||||
await updateAccessPermissions(SystemRoles.USER, {
|
||||
[PermissionTypes.PROMPTS]: { SHARE: true },
|
||||
[PermissionTypes.MULTI_CONVO]: { USE: true },
|
||||
});
|
||||
|
||||
const updatedRole = await getRoleByName(SystemRoles.USER);
|
||||
expect(updatedRole.permissions[PermissionTypes.PROMPTS]).toEqual({
|
||||
CREATE: true,
|
||||
USE: true,
|
||||
SHARE: true,
|
||||
});
|
||||
expect(updatedRole.permissions[PermissionTypes.MULTI_CONVO]).toEqual({ USE: true });
|
||||
});
|
||||
|
||||
it('should not update MULTI_CONVO permissions when no changes are needed', async () => {
|
||||
await new Role({
|
||||
name: SystemRoles.USER,
|
||||
permissions: {
|
||||
[PermissionTypes.MULTI_CONVO]: { USE: true },
|
||||
},
|
||||
}).save();
|
||||
|
||||
await updateAccessPermissions(SystemRoles.USER, {
|
||||
[PermissionTypes.MULTI_CONVO]: { USE: true },
|
||||
});
|
||||
|
||||
const updatedRole = await getRoleByName(SystemRoles.USER);
|
||||
expect(updatedRole.permissions[PermissionTypes.MULTI_CONVO]).toEqual({ USE: true });
|
||||
});
|
||||
});
|
||||
|
||||
describe('initializeRoles', () => {
|
||||
beforeEach(async () => {
|
||||
await Role.deleteMany({});
|
||||
});
|
||||
|
||||
it('should create default roles if they do not exist', async () => {
|
||||
await initializeRoles();
|
||||
|
||||
const adminRole = await getRoleByName(SystemRoles.ADMIN);
|
||||
const userRole = await getRoleByName(SystemRoles.USER);
|
||||
|
||||
expect(adminRole).toBeTruthy();
|
||||
expect(userRole).toBeTruthy();
|
||||
|
||||
// Check if all permission types exist in the permissions field
|
||||
Object.values(PermissionTypes).forEach((permType) => {
|
||||
expect(adminRole.permissions[permType]).toBeDefined();
|
||||
expect(userRole.permissions[permType]).toBeDefined();
|
||||
});
|
||||
|
||||
// Example: Check default values for ADMIN role
|
||||
expect(adminRole.permissions[PermissionTypes.PROMPTS].SHARE).toBe(true);
|
||||
expect(adminRole.permissions[PermissionTypes.BOOKMARKS].USE).toBe(true);
|
||||
expect(adminRole.permissions[PermissionTypes.AGENTS].CREATE).toBe(true);
|
||||
});
|
||||
|
||||
it('should not modify existing permissions for existing roles', async () => {
|
||||
const customUserRole = {
|
||||
name: SystemRoles.USER,
|
||||
permissions: {
|
||||
[PermissionTypes.PROMPTS]: {
|
||||
[Permissions.USE]: false,
|
||||
[Permissions.CREATE]: true,
|
||||
[Permissions.SHARE]: true,
|
||||
},
|
||||
[PermissionTypes.BOOKMARKS]: { [Permissions.USE]: false },
|
||||
},
|
||||
};
|
||||
|
||||
await new Role(customUserRole).save();
|
||||
await initializeRoles();
|
||||
|
||||
const userRole = await getRoleByName(SystemRoles.USER);
|
||||
expect(userRole.permissions[PermissionTypes.PROMPTS]).toEqual(
|
||||
customUserRole.permissions[PermissionTypes.PROMPTS],
|
||||
);
|
||||
expect(userRole.permissions[PermissionTypes.BOOKMARKS]).toEqual(
|
||||
customUserRole.permissions[PermissionTypes.BOOKMARKS],
|
||||
);
|
||||
expect(userRole.permissions[PermissionTypes.AGENTS]).toBeDefined();
|
||||
});
|
||||
|
||||
it('should add new permission types to existing roles', async () => {
|
||||
const partialUserRole = {
|
||||
name: SystemRoles.USER,
|
||||
permissions: {
|
||||
[PermissionTypes.PROMPTS]:
|
||||
roleDefaults[SystemRoles.USER].permissions[PermissionTypes.PROMPTS],
|
||||
[PermissionTypes.BOOKMARKS]:
|
||||
roleDefaults[SystemRoles.USER].permissions[PermissionTypes.BOOKMARKS],
|
||||
},
|
||||
};
|
||||
|
||||
await new Role(partialUserRole).save();
|
||||
await initializeRoles();
|
||||
|
||||
const userRole = await getRoleByName(SystemRoles.USER);
|
||||
expect(userRole.permissions[PermissionTypes.AGENTS]).toBeDefined();
|
||||
expect(userRole.permissions[PermissionTypes.AGENTS].CREATE).toBeDefined();
|
||||
expect(userRole.permissions[PermissionTypes.AGENTS].USE).toBeDefined();
|
||||
expect(userRole.permissions[PermissionTypes.AGENTS].SHARE).toBeDefined();
|
||||
});
|
||||
|
||||
it('should handle multiple runs without duplicating or modifying data', async () => {
|
||||
await initializeRoles();
|
||||
await initializeRoles();
|
||||
|
||||
const adminRoles = await Role.find({ name: SystemRoles.ADMIN });
|
||||
const userRoles = await Role.find({ name: SystemRoles.USER });
|
||||
|
||||
expect(adminRoles).toHaveLength(1);
|
||||
expect(userRoles).toHaveLength(1);
|
||||
|
||||
const adminPerms = adminRoles[0].toObject().permissions;
|
||||
const userPerms = userRoles[0].toObject().permissions;
|
||||
Object.values(PermissionTypes).forEach((permType) => {
|
||||
expect(adminPerms[permType]).toBeDefined();
|
||||
expect(userPerms[permType]).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
it('should update roles with missing permission types from roleDefaults', async () => {
|
||||
const partialAdminRole = {
|
||||
name: SystemRoles.ADMIN,
|
||||
permissions: {
|
||||
[PermissionTypes.PROMPTS]: {
|
||||
[Permissions.USE]: false,
|
||||
[Permissions.CREATE]: false,
|
||||
[Permissions.SHARE]: false,
|
||||
},
|
||||
[PermissionTypes.BOOKMARKS]:
|
||||
roleDefaults[SystemRoles.ADMIN].permissions[PermissionTypes.BOOKMARKS],
|
||||
},
|
||||
};
|
||||
|
||||
await new Role(partialAdminRole).save();
|
||||
await initializeRoles();
|
||||
|
||||
const adminRole = await getRoleByName(SystemRoles.ADMIN);
|
||||
expect(adminRole.permissions[PermissionTypes.PROMPTS]).toEqual(
|
||||
partialAdminRole.permissions[PermissionTypes.PROMPTS],
|
||||
);
|
||||
expect(adminRole.permissions[PermissionTypes.AGENTS]).toBeDefined();
|
||||
expect(adminRole.permissions[PermissionTypes.AGENTS].CREATE).toBeDefined();
|
||||
expect(adminRole.permissions[PermissionTypes.AGENTS].USE).toBeDefined();
|
||||
expect(adminRole.permissions[PermissionTypes.AGENTS].SHARE).toBeDefined();
|
||||
});
|
||||
|
||||
it('should include MULTI_CONVO permissions when creating default roles', async () => {
|
||||
await initializeRoles();
|
||||
|
||||
const adminRole = await getRoleByName(SystemRoles.ADMIN);
|
||||
const userRole = await getRoleByName(SystemRoles.USER);
|
||||
|
||||
expect(adminRole.permissions[PermissionTypes.MULTI_CONVO]).toBeDefined();
|
||||
expect(userRole.permissions[PermissionTypes.MULTI_CONVO]).toBeDefined();
|
||||
expect(adminRole.permissions[PermissionTypes.MULTI_CONVO].USE).toBe(
|
||||
roleDefaults[SystemRoles.ADMIN].permissions[PermissionTypes.MULTI_CONVO].USE,
|
||||
);
|
||||
expect(userRole.permissions[PermissionTypes.MULTI_CONVO].USE).toBe(
|
||||
roleDefaults[SystemRoles.USER].permissions[PermissionTypes.MULTI_CONVO].USE,
|
||||
);
|
||||
});
|
||||
|
||||
it('should add MULTI_CONVO permissions to existing roles without them', async () => {
|
||||
const partialUserRole = {
|
||||
name: SystemRoles.USER,
|
||||
permissions: {
|
||||
[PermissionTypes.PROMPTS]:
|
||||
roleDefaults[SystemRoles.USER].permissions[PermissionTypes.PROMPTS],
|
||||
[PermissionTypes.BOOKMARKS]:
|
||||
roleDefaults[SystemRoles.USER].permissions[PermissionTypes.BOOKMARKS],
|
||||
},
|
||||
};
|
||||
|
||||
await new Role(partialUserRole).save();
|
||||
await initializeRoles();
|
||||
|
||||
const userRole = await getRoleByName(SystemRoles.USER);
|
||||
expect(userRole.permissions[PermissionTypes.MULTI_CONVO]).toBeDefined();
|
||||
expect(userRole.permissions[PermissionTypes.MULTI_CONVO].USE).toBeDefined();
|
||||
});
|
||||
});
|
||||
96
api/models/ToolCall.js
Normal file
96
api/models/ToolCall.js
Normal file
@ -0,0 +1,96 @@
|
||||
const { ToolCall } = require('~/db/models');
|
||||
|
||||
/**
|
||||
* Create a new tool call
|
||||
* @param {IToolCallData} toolCallData - The tool call data
|
||||
* @returns {Promise<IToolCallData>} The created tool call document
|
||||
*/
|
||||
async function createToolCall(toolCallData) {
|
||||
try {
|
||||
return await ToolCall.create(toolCallData);
|
||||
} catch (error) {
|
||||
throw new Error(`Error creating tool call: ${error.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get a tool call by ID
|
||||
* @param {string} id - The tool call document ID
|
||||
* @returns {Promise<IToolCallData|null>} The tool call document or null if not found
|
||||
*/
|
||||
async function getToolCallById(id) {
|
||||
try {
|
||||
return await ToolCall.findById(id).lean();
|
||||
} catch (error) {
|
||||
throw new Error(`Error fetching tool call: ${error.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get tool calls by message ID and user
|
||||
* @param {string} messageId - The message ID
|
||||
* @param {string} userId - The user's ObjectId
|
||||
* @returns {Promise<Array>} Array of tool call documents
|
||||
*/
|
||||
async function getToolCallsByMessage(messageId, userId) {
|
||||
try {
|
||||
return await ToolCall.find({ messageId, user: userId }).lean();
|
||||
} catch (error) {
|
||||
throw new Error(`Error fetching tool calls: ${error.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get tool calls by conversation ID and user
|
||||
* @param {string} conversationId - The conversation ID
|
||||
* @param {string} userId - The user's ObjectId
|
||||
* @returns {Promise<IToolCallData[]>} Array of tool call documents
|
||||
*/
|
||||
async function getToolCallsByConvo(conversationId, userId) {
|
||||
try {
|
||||
return await ToolCall.find({ conversationId, user: userId }).lean();
|
||||
} catch (error) {
|
||||
throw new Error(`Error fetching tool calls: ${error.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Update a tool call
|
||||
* @param {string} id - The tool call document ID
|
||||
* @param {Partial<IToolCallData>} updateData - The data to update
|
||||
* @returns {Promise<IToolCallData|null>} The updated tool call document or null if not found
|
||||
*/
|
||||
async function updateToolCall(id, updateData) {
|
||||
try {
|
||||
return await ToolCall.findByIdAndUpdate(id, updateData, { new: true }).lean();
|
||||
} catch (error) {
|
||||
throw new Error(`Error updating tool call: ${error.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Delete a tool call
|
||||
* @param {string} userId - The related user's ObjectId
|
||||
* @param {string} [conversationId] - The tool call conversation ID
|
||||
* @returns {Promise<{ ok?: number; n?: number; deletedCount?: number }>} The result of the delete operation
|
||||
*/
|
||||
async function deleteToolCalls(userId, conversationId) {
|
||||
try {
|
||||
const query = { user: userId };
|
||||
if (conversationId) {
|
||||
query.conversationId = conversationId;
|
||||
}
|
||||
return await ToolCall.deleteMany(query);
|
||||
} catch (error) {
|
||||
throw new Error(`Error deleting tool call: ${error.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
createToolCall,
|
||||
updateToolCall,
|
||||
deleteToolCalls,
|
||||
getToolCallById,
|
||||
getToolCallsByConvo,
|
||||
getToolCallsByMessage,
|
||||
};
|
||||
346
api/models/Transaction.js
Normal file
346
api/models/Transaction.js
Normal file
@ -0,0 +1,346 @@
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { getMultiplier, getCacheMultiplier } = require('./tx');
|
||||
const { Transaction, Balance } = require('~/db/models');
|
||||
|
||||
const cancelRate = 1.15;
|
||||
|
||||
/**
|
||||
* Updates a user's token balance based on a transaction using optimistic concurrency control
|
||||
* without schema changes. Compatible with DocumentDB.
|
||||
* @async
|
||||
* @function
|
||||
* @param {Object} params - The function parameters.
|
||||
* @param {string|mongoose.Types.ObjectId} params.user - The user ID.
|
||||
* @param {number} params.incrementValue - The value to increment the balance by (can be negative).
|
||||
* @param {import('mongoose').UpdateQuery<import('@librechat/data-schemas').IBalance>['$set']} [params.setValues] - Optional additional fields to set.
|
||||
* @returns {Promise<Object>} Returns the updated balance document (lean).
|
||||
* @throws {Error} Throws an error if the update fails after multiple retries.
|
||||
*/
|
||||
const updateBalance = async ({ user, incrementValue, setValues }) => {
|
||||
let maxRetries = 10; // Number of times to retry on conflict
|
||||
let delay = 50; // Initial retry delay in ms
|
||||
let lastError = null;
|
||||
|
||||
for (let attempt = 1; attempt <= maxRetries; attempt++) {
|
||||
let currentBalanceDoc;
|
||||
try {
|
||||
// 1. Read the current document state
|
||||
currentBalanceDoc = await Balance.findOne({ user }).lean();
|
||||
const currentCredits = currentBalanceDoc ? currentBalanceDoc.tokenCredits : 0;
|
||||
|
||||
// 2. Calculate the desired new state
|
||||
const potentialNewCredits = currentCredits + incrementValue;
|
||||
const newCredits = Math.max(0, potentialNewCredits); // Ensure balance doesn't go below zero
|
||||
|
||||
// 3. Prepare the update payload
|
||||
const updatePayload = {
|
||||
$set: {
|
||||
tokenCredits: newCredits,
|
||||
...(setValues || {}), // Merge other values to set
|
||||
},
|
||||
};
|
||||
|
||||
// 4. Attempt the conditional update or upsert
|
||||
let updatedBalance = null;
|
||||
if (currentBalanceDoc) {
|
||||
// --- Document Exists: Perform Conditional Update ---
|
||||
// Try to update only if the tokenCredits match the value we read (currentCredits)
|
||||
updatedBalance = await Balance.findOneAndUpdate(
|
||||
{
|
||||
user: user,
|
||||
tokenCredits: currentCredits, // Optimistic lock: condition based on the read value
|
||||
},
|
||||
updatePayload,
|
||||
{
|
||||
new: true, // Return the modified document
|
||||
// lean: true, // .lean() is applied after query execution in Mongoose >= 6
|
||||
},
|
||||
).lean(); // Use lean() for plain JS object
|
||||
|
||||
if (updatedBalance) {
|
||||
// Success! The update was applied based on the expected current state.
|
||||
return updatedBalance;
|
||||
}
|
||||
// If updatedBalance is null, it means tokenCredits changed between read and write (conflict).
|
||||
lastError = new Error(`Concurrency conflict for user ${user} on attempt ${attempt}.`);
|
||||
// Proceed to retry logic below.
|
||||
} else {
|
||||
// --- Document Does Not Exist: Perform Conditional Upsert ---
|
||||
// Try to insert the document, but only if it still doesn't exist.
|
||||
// Using tokenCredits: {$exists: false} helps prevent race conditions where
|
||||
// another process creates the doc between our findOne and findOneAndUpdate.
|
||||
try {
|
||||
updatedBalance = await Balance.findOneAndUpdate(
|
||||
{
|
||||
user: user,
|
||||
// Attempt to match only if the document doesn't exist OR was just created
|
||||
// without tokenCredits (less likely but possible). A simple { user } filter
|
||||
// might also work, relying on the retry for conflicts.
|
||||
// Let's use a simpler filter and rely on retry for races.
|
||||
// tokenCredits: { $exists: false } // This condition might be too strict if doc exists with 0 credits
|
||||
},
|
||||
updatePayload,
|
||||
{
|
||||
upsert: true, // Create if doesn't exist
|
||||
new: true, // Return the created/updated document
|
||||
// setDefaultsOnInsert: true, // Ensure schema defaults are applied on insert
|
||||
// lean: true,
|
||||
},
|
||||
).lean();
|
||||
|
||||
if (updatedBalance) {
|
||||
// Upsert succeeded (likely created the document)
|
||||
return updatedBalance;
|
||||
}
|
||||
// If null, potentially a rare race condition during upsert. Retry should handle it.
|
||||
lastError = new Error(
|
||||
`Upsert race condition suspected for user ${user} on attempt ${attempt}.`,
|
||||
);
|
||||
} catch (error) {
|
||||
if (error.code === 11000) {
|
||||
// E11000 duplicate key error on index
|
||||
// This means another process created the document *just* before our upsert.
|
||||
// It's a concurrency conflict during creation. We should retry.
|
||||
lastError = error; // Store the error
|
||||
// Proceed to retry logic below.
|
||||
} else {
|
||||
// Different error, rethrow
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
} // End if/else (document exists?)
|
||||
} catch (error) {
|
||||
// Catch errors from findOne or unexpected findOneAndUpdate errors
|
||||
logger.error(`[updateBalance] Error during attempt ${attempt} for user ${user}:`, error);
|
||||
lastError = error; // Store the error
|
||||
// Consider stopping retries for non-transient errors, but for now, we retry.
|
||||
}
|
||||
|
||||
// If we reached here, it means the update failed (conflict or error), wait and retry
|
||||
if (attempt < maxRetries) {
|
||||
const jitter = Math.random() * delay * 0.5; // Add jitter to delay
|
||||
await new Promise((resolve) => setTimeout(resolve, delay + jitter));
|
||||
delay = Math.min(delay * 2, 2000); // Exponential backoff with cap
|
||||
}
|
||||
} // End for loop (retries)
|
||||
|
||||
// If loop finishes without success, throw the last encountered error or a generic one
|
||||
logger.error(
|
||||
`[updateBalance] Failed to update balance for user ${user} after ${maxRetries} attempts.`,
|
||||
);
|
||||
throw (
|
||||
lastError ||
|
||||
new Error(
|
||||
`Failed to update balance for user ${user} after maximum retries due to persistent conflicts.`,
|
||||
)
|
||||
);
|
||||
};
|
||||
|
||||
/** Method to calculate and set the tokenValue for a transaction */
|
||||
function calculateTokenValue(txn) {
|
||||
if (!txn.valueKey || !txn.tokenType) {
|
||||
txn.tokenValue = txn.rawAmount;
|
||||
}
|
||||
const { valueKey, tokenType, model, endpointTokenConfig } = txn;
|
||||
const multiplier = Math.abs(getMultiplier({ valueKey, tokenType, model, endpointTokenConfig }));
|
||||
txn.rate = multiplier;
|
||||
txn.tokenValue = txn.rawAmount * multiplier;
|
||||
if (txn.context && txn.tokenType === 'completion' && txn.context === 'incomplete') {
|
||||
txn.tokenValue = Math.ceil(txn.tokenValue * cancelRate);
|
||||
txn.rate *= cancelRate;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* New static method to create an auto-refill transaction that does NOT trigger a balance update.
|
||||
* @param {object} txData - Transaction data.
|
||||
* @param {string} txData.user - The user ID.
|
||||
* @param {string} txData.tokenType - The type of token.
|
||||
* @param {string} txData.context - The context of the transaction.
|
||||
* @param {number} txData.rawAmount - The raw amount of tokens.
|
||||
* @returns {Promise<object>} - The created transaction.
|
||||
*/
|
||||
async function createAutoRefillTransaction(txData) {
|
||||
if (txData.rawAmount != null && isNaN(txData.rawAmount)) {
|
||||
return;
|
||||
}
|
||||
const transaction = new Transaction(txData);
|
||||
transaction.endpointTokenConfig = txData.endpointTokenConfig;
|
||||
calculateTokenValue(transaction);
|
||||
await transaction.save();
|
||||
|
||||
const balanceResponse = await updateBalance({
|
||||
user: transaction.user,
|
||||
incrementValue: txData.rawAmount,
|
||||
setValues: { lastRefill: new Date() },
|
||||
});
|
||||
const result = {
|
||||
rate: transaction.rate,
|
||||
user: transaction.user.toString(),
|
||||
balance: balanceResponse.tokenCredits,
|
||||
};
|
||||
logger.debug('[Balance.check] Auto-refill performed', result);
|
||||
result.transaction = transaction;
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* Static method to create a transaction and update the balance
|
||||
* @param {txData} _txData - Transaction data.
|
||||
*/
|
||||
async function createTransaction(_txData) {
|
||||
const { balance, transactions, ...txData } = _txData;
|
||||
if (txData.rawAmount != null && isNaN(txData.rawAmount)) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (transactions?.enabled === false) {
|
||||
return;
|
||||
}
|
||||
|
||||
const transaction = new Transaction(txData);
|
||||
transaction.endpointTokenConfig = txData.endpointTokenConfig;
|
||||
calculateTokenValue(transaction);
|
||||
|
||||
await transaction.save();
|
||||
if (!balance?.enabled) {
|
||||
return;
|
||||
}
|
||||
|
||||
let incrementValue = transaction.tokenValue;
|
||||
const balanceResponse = await updateBalance({
|
||||
user: transaction.user,
|
||||
incrementValue,
|
||||
});
|
||||
|
||||
return {
|
||||
rate: transaction.rate,
|
||||
user: transaction.user.toString(),
|
||||
balance: balanceResponse.tokenCredits,
|
||||
[transaction.tokenType]: incrementValue,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Static method to create a structured transaction and update the balance
|
||||
* @param {txData} _txData - Transaction data.
|
||||
*/
|
||||
async function createStructuredTransaction(_txData) {
|
||||
const { balance, transactions, ...txData } = _txData;
|
||||
if (transactions?.enabled === false) {
|
||||
return;
|
||||
}
|
||||
|
||||
const transaction = new Transaction({
|
||||
...txData,
|
||||
endpointTokenConfig: txData.endpointTokenConfig,
|
||||
});
|
||||
|
||||
calculateStructuredTokenValue(transaction);
|
||||
|
||||
await transaction.save();
|
||||
|
||||
if (!balance?.enabled) {
|
||||
return;
|
||||
}
|
||||
|
||||
let incrementValue = transaction.tokenValue;
|
||||
|
||||
const balanceResponse = await updateBalance({
|
||||
user: transaction.user,
|
||||
incrementValue,
|
||||
});
|
||||
|
||||
return {
|
||||
rate: transaction.rate,
|
||||
user: transaction.user.toString(),
|
||||
balance: balanceResponse.tokenCredits,
|
||||
[transaction.tokenType]: incrementValue,
|
||||
};
|
||||
}
|
||||
|
||||
/** Method to calculate token value for structured tokens */
|
||||
function calculateStructuredTokenValue(txn) {
|
||||
if (!txn.tokenType) {
|
||||
txn.tokenValue = txn.rawAmount;
|
||||
return;
|
||||
}
|
||||
|
||||
const { model, endpointTokenConfig } = txn;
|
||||
|
||||
if (txn.tokenType === 'prompt') {
|
||||
const inputMultiplier = getMultiplier({ tokenType: 'prompt', model, endpointTokenConfig });
|
||||
const writeMultiplier =
|
||||
getCacheMultiplier({ cacheType: 'write', model, endpointTokenConfig }) ?? inputMultiplier;
|
||||
const readMultiplier =
|
||||
getCacheMultiplier({ cacheType: 'read', model, endpointTokenConfig }) ?? inputMultiplier;
|
||||
|
||||
txn.rateDetail = {
|
||||
input: inputMultiplier,
|
||||
write: writeMultiplier,
|
||||
read: readMultiplier,
|
||||
};
|
||||
|
||||
const totalPromptTokens =
|
||||
Math.abs(txn.inputTokens || 0) +
|
||||
Math.abs(txn.writeTokens || 0) +
|
||||
Math.abs(txn.readTokens || 0);
|
||||
|
||||
if (totalPromptTokens > 0) {
|
||||
txn.rate =
|
||||
(Math.abs(inputMultiplier * (txn.inputTokens || 0)) +
|
||||
Math.abs(writeMultiplier * (txn.writeTokens || 0)) +
|
||||
Math.abs(readMultiplier * (txn.readTokens || 0))) /
|
||||
totalPromptTokens;
|
||||
} else {
|
||||
txn.rate = Math.abs(inputMultiplier); // Default to input rate if no tokens
|
||||
}
|
||||
|
||||
txn.tokenValue = -(
|
||||
Math.abs(txn.inputTokens || 0) * inputMultiplier +
|
||||
Math.abs(txn.writeTokens || 0) * writeMultiplier +
|
||||
Math.abs(txn.readTokens || 0) * readMultiplier
|
||||
);
|
||||
|
||||
txn.rawAmount = -totalPromptTokens;
|
||||
} else if (txn.tokenType === 'completion') {
|
||||
const multiplier = getMultiplier({ tokenType: txn.tokenType, model, endpointTokenConfig });
|
||||
txn.rate = Math.abs(multiplier);
|
||||
txn.tokenValue = -Math.abs(txn.rawAmount) * multiplier;
|
||||
txn.rawAmount = -Math.abs(txn.rawAmount);
|
||||
}
|
||||
|
||||
if (txn.context && txn.tokenType === 'completion' && txn.context === 'incomplete') {
|
||||
txn.tokenValue = Math.ceil(txn.tokenValue * cancelRate);
|
||||
txn.rate *= cancelRate;
|
||||
if (txn.rateDetail) {
|
||||
txn.rateDetail = Object.fromEntries(
|
||||
Object.entries(txn.rateDetail).map(([k, v]) => [k, v * cancelRate]),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Queries and retrieves transactions based on a given filter.
|
||||
* @async
|
||||
* @function getTransactions
|
||||
* @param {Object} filter - MongoDB filter object to apply when querying transactions.
|
||||
* @returns {Promise<Array>} A promise that resolves to an array of matched transactions.
|
||||
* @throws {Error} Throws an error if querying the database fails.
|
||||
*/
|
||||
async function getTransactions(filter) {
|
||||
try {
|
||||
return await Transaction.find(filter).lean();
|
||||
} catch (error) {
|
||||
logger.error('Error querying transactions:', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
getTransactions,
|
||||
createTransaction,
|
||||
createAutoRefillTransaction,
|
||||
createStructuredTransaction,
|
||||
};
|
||||
566
api/models/Transaction.spec.js
Normal file
566
api/models/Transaction.spec.js
Normal file
@ -0,0 +1,566 @@
|
||||
const mongoose = require('mongoose');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
const { spendTokens, spendStructuredTokens } = require('./spendTokens');
|
||||
const { getMultiplier, getCacheMultiplier } = require('./tx');
|
||||
const { createTransaction, createStructuredTransaction } = require('./Transaction');
|
||||
const { Balance, Transaction } = require('~/db/models');
|
||||
|
||||
let mongoServer;
|
||||
beforeAll(async () => {
|
||||
mongoServer = await MongoMemoryServer.create();
|
||||
const mongoUri = mongoServer.getUri();
|
||||
await mongoose.connect(mongoUri);
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await mongoose.disconnect();
|
||||
await mongoServer.stop();
|
||||
});
|
||||
|
||||
beforeEach(async () => {
|
||||
await mongoose.connection.dropDatabase();
|
||||
});
|
||||
|
||||
describe('Regular Token Spending Tests', () => {
|
||||
test('Balance should decrease when spending tokens with spendTokens', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000; // $10.00
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: 100,
|
||||
completionTokens: 50,
|
||||
};
|
||||
|
||||
// Act
|
||||
await spendTokens(txData, tokenUsage);
|
||||
|
||||
// Assert
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
|
||||
const completionMultiplier = getMultiplier({ model, tokenType: 'completion' });
|
||||
const expectedTotalCost = 100 * promptMultiplier + 50 * completionMultiplier;
|
||||
const expectedBalance = initialBalance - expectedTotalCost;
|
||||
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(expectedBalance, 0);
|
||||
});
|
||||
|
||||
test('spendTokens should handle zero completion tokens', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: 100,
|
||||
completionTokens: 0,
|
||||
};
|
||||
|
||||
// Act
|
||||
await spendTokens(txData, tokenUsage);
|
||||
|
||||
// Assert
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
|
||||
const expectedCost = 100 * promptMultiplier;
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
|
||||
test('spendTokens should handle undefined token counts', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
const tokenUsage = {};
|
||||
|
||||
// Act
|
||||
const result = await spendTokens(txData, tokenUsage);
|
||||
|
||||
// Assert: No transaction should be created
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
test('spendTokens should handle only prompt tokens', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
const tokenUsage = { promptTokens: 100 };
|
||||
|
||||
// Act
|
||||
await spendTokens(txData, tokenUsage);
|
||||
|
||||
// Assert
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
|
||||
const expectedCost = 100 * promptMultiplier;
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
|
||||
test('spendTokens should not update balance when balance feature is disabled', async () => {
|
||||
// Arrange: Balance config is now passed directly in txData
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
balance: { enabled: false },
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: 100,
|
||||
completionTokens: 50,
|
||||
};
|
||||
|
||||
// Act
|
||||
await spendTokens(txData, tokenUsage);
|
||||
|
||||
// Assert: Balance should remain unchanged.
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(updatedBalance.tokenCredits).toBe(initialBalance);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Structured Token Spending Tests', () => {
|
||||
test('Balance should decrease and rawAmount should be set when spending a large number of structured tokens', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 17613154.55; // $17.61
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-3-5-sonnet';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'c23a18da-706c-470a-ac28-ec87ed065199',
|
||||
model,
|
||||
context: 'message',
|
||||
endpointTokenConfig: null,
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: {
|
||||
input: 11,
|
||||
write: 140522,
|
||||
read: 0,
|
||||
},
|
||||
completionTokens: 5,
|
||||
};
|
||||
|
||||
const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
|
||||
const completionMultiplier = getMultiplier({ model, tokenType: 'completion' });
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' });
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' });
|
||||
|
||||
// Act
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
// Calculate expected costs.
|
||||
const expectedPromptCost =
|
||||
tokenUsage.promptTokens.input * promptMultiplier +
|
||||
tokenUsage.promptTokens.write * writeMultiplier +
|
||||
tokenUsage.promptTokens.read * readMultiplier;
|
||||
const expectedCompletionCost = tokenUsage.completionTokens * completionMultiplier;
|
||||
const expectedTotalCost = expectedPromptCost + expectedCompletionCost;
|
||||
const expectedBalance = initialBalance - expectedTotalCost;
|
||||
|
||||
// Assert
|
||||
expect(result.completion.balance).toBeLessThan(initialBalance);
|
||||
const allowedDifference = 100;
|
||||
expect(Math.abs(result.completion.balance - expectedBalance)).toBeLessThan(allowedDifference);
|
||||
const balanceDecrease = initialBalance - result.completion.balance;
|
||||
expect(balanceDecrease).toBeCloseTo(expectedTotalCost, 0);
|
||||
|
||||
const expectedPromptTokenValue = -expectedPromptCost;
|
||||
const expectedCompletionTokenValue = -expectedCompletionCost;
|
||||
expect(result.prompt.prompt).toBeCloseTo(expectedPromptTokenValue, 1);
|
||||
expect(result.completion.completion).toBe(expectedCompletionTokenValue);
|
||||
});
|
||||
|
||||
test('should handle zero completion tokens in structured spending', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 17613154.55;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-3-5-sonnet';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-convo',
|
||||
model,
|
||||
context: 'message',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: {
|
||||
input: 10,
|
||||
write: 100,
|
||||
read: 5,
|
||||
},
|
||||
completionTokens: 0,
|
||||
};
|
||||
|
||||
// Act
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
// Assert
|
||||
expect(result.prompt).toBeDefined();
|
||||
expect(result.completion).toBeUndefined();
|
||||
expect(result.prompt.prompt).toBeLessThan(0);
|
||||
});
|
||||
|
||||
test('should handle only prompt tokens in structured spending', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 17613154.55;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-3-5-sonnet';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-convo',
|
||||
model,
|
||||
context: 'message',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: {
|
||||
input: 10,
|
||||
write: 100,
|
||||
read: 5,
|
||||
},
|
||||
};
|
||||
|
||||
// Act
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
// Assert
|
||||
expect(result.prompt).toBeDefined();
|
||||
expect(result.completion).toBeUndefined();
|
||||
expect(result.prompt.prompt).toBeLessThan(0);
|
||||
});
|
||||
|
||||
test('should handle undefined token counts in structured spending', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 17613154.55;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-3-5-sonnet';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-convo',
|
||||
model,
|
||||
context: 'message',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
const tokenUsage = {};
|
||||
|
||||
// Act
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
// Assert
|
||||
expect(result).toEqual({
|
||||
prompt: undefined,
|
||||
completion: undefined,
|
||||
});
|
||||
});
|
||||
|
||||
test('should handle incomplete context for completion tokens', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 17613154.55;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-3-5-sonnet';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-convo',
|
||||
model,
|
||||
context: 'incomplete',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: {
|
||||
input: 10,
|
||||
write: 100,
|
||||
read: 5,
|
||||
},
|
||||
completionTokens: 50,
|
||||
};
|
||||
|
||||
// Act
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
// Assert:
|
||||
// (Assuming a multiplier for completion of 15 and a cancel rate of 1.15 as noted in the original test.)
|
||||
expect(result.completion.completion).toBeCloseTo(-50 * 15 * 1.15, 0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('NaN Handling Tests', () => {
|
||||
test('should skip transaction creation when rawAmount is NaN', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
rawAmount: NaN,
|
||||
tokenType: 'prompt',
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
// Act
|
||||
const result = await createTransaction(txData);
|
||||
|
||||
// Assert: No transaction should be created and balance remains unchanged.
|
||||
expect(result).toBeUndefined();
|
||||
const balance = await Balance.findOne({ user: userId });
|
||||
expect(balance.tokenCredits).toBe(initialBalance);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Transactions Config Tests', () => {
|
||||
test('createTransaction should not save when transactions.enabled is false', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
rawAmount: -100,
|
||||
tokenType: 'prompt',
|
||||
transactions: { enabled: false },
|
||||
};
|
||||
|
||||
// Act
|
||||
const result = await createTransaction(txData);
|
||||
|
||||
// Assert: No transaction should be created
|
||||
expect(result).toBeUndefined();
|
||||
const transactions = await Transaction.find({ user: userId });
|
||||
expect(transactions).toHaveLength(0);
|
||||
const balance = await Balance.findOne({ user: userId });
|
||||
expect(balance.tokenCredits).toBe(initialBalance);
|
||||
});
|
||||
|
||||
test('createTransaction should save when transactions.enabled is true', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
rawAmount: -100,
|
||||
tokenType: 'prompt',
|
||||
transactions: { enabled: true },
|
||||
balance: { enabled: true },
|
||||
};
|
||||
|
||||
// Act
|
||||
const result = await createTransaction(txData);
|
||||
|
||||
// Assert: Transaction should be created
|
||||
expect(result).toBeDefined();
|
||||
expect(result.balance).toBeLessThan(initialBalance);
|
||||
const transactions = await Transaction.find({ user: userId });
|
||||
expect(transactions).toHaveLength(1);
|
||||
expect(transactions[0].rawAmount).toBe(-100);
|
||||
});
|
||||
|
||||
test('createTransaction should save when balance.enabled is true even if transactions config is missing', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
rawAmount: -100,
|
||||
tokenType: 'prompt',
|
||||
balance: { enabled: true },
|
||||
// No transactions config provided
|
||||
};
|
||||
|
||||
// Act
|
||||
const result = await createTransaction(txData);
|
||||
|
||||
// Assert: Transaction should be created (backward compatibility)
|
||||
expect(result).toBeDefined();
|
||||
expect(result.balance).toBeLessThan(initialBalance);
|
||||
const transactions = await Transaction.find({ user: userId });
|
||||
expect(transactions).toHaveLength(1);
|
||||
});
|
||||
|
||||
test('createTransaction should save transaction but not update balance when balance is disabled but transactions enabled', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
rawAmount: -100,
|
||||
tokenType: 'prompt',
|
||||
transactions: { enabled: true },
|
||||
balance: { enabled: false },
|
||||
};
|
||||
|
||||
// Act
|
||||
const result = await createTransaction(txData);
|
||||
|
||||
// Assert: Transaction should be created but balance unchanged
|
||||
expect(result).toBeUndefined();
|
||||
const transactions = await Transaction.find({ user: userId });
|
||||
expect(transactions).toHaveLength(1);
|
||||
expect(transactions[0].rawAmount).toBe(-100);
|
||||
const balance = await Balance.findOne({ user: userId });
|
||||
expect(balance.tokenCredits).toBe(initialBalance);
|
||||
});
|
||||
|
||||
test('createStructuredTransaction should not save when transactions.enabled is false', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-3-5-sonnet';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'message',
|
||||
tokenType: 'prompt',
|
||||
inputTokens: -10,
|
||||
writeTokens: -100,
|
||||
readTokens: -5,
|
||||
transactions: { enabled: false },
|
||||
};
|
||||
|
||||
// Act
|
||||
const result = await createStructuredTransaction(txData);
|
||||
|
||||
// Assert: No transaction should be created
|
||||
expect(result).toBeUndefined();
|
||||
const transactions = await Transaction.find({ user: userId });
|
||||
expect(transactions).toHaveLength(0);
|
||||
const balance = await Balance.findOne({ user: userId });
|
||||
expect(balance.tokenCredits).toBe(initialBalance);
|
||||
});
|
||||
|
||||
test('createStructuredTransaction should save transaction but not update balance when balance is disabled but transactions enabled', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-3-5-sonnet';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'message',
|
||||
tokenType: 'prompt',
|
||||
inputTokens: -10,
|
||||
writeTokens: -100,
|
||||
readTokens: -5,
|
||||
transactions: { enabled: true },
|
||||
balance: { enabled: false },
|
||||
};
|
||||
|
||||
// Act
|
||||
const result = await createStructuredTransaction(txData);
|
||||
|
||||
// Assert: Transaction should be created but balance unchanged
|
||||
expect(result).toBeUndefined();
|
||||
const transactions = await Transaction.find({ user: userId });
|
||||
expect(transactions).toHaveLength(1);
|
||||
expect(transactions[0].inputTokens).toBe(-10);
|
||||
expect(transactions[0].writeTokens).toBe(-100);
|
||||
expect(transactions[0].readTokens).toBe(-5);
|
||||
const balance = await Balance.findOne({ user: userId });
|
||||
expect(balance.tokenCredits).toBe(initialBalance);
|
||||
});
|
||||
});
|
||||
156
api/models/balanceMethods.js
Normal file
156
api/models/balanceMethods.js
Normal file
@ -0,0 +1,156 @@
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { ViolationTypes } = require('librechat-data-provider');
|
||||
const { createAutoRefillTransaction } = require('./Transaction');
|
||||
const { logViolation } = require('~/cache');
|
||||
const { getMultiplier } = require('./tx');
|
||||
const { Balance } = require('~/db/models');
|
||||
|
||||
function isInvalidDate(date) {
|
||||
return isNaN(date);
|
||||
}
|
||||
|
||||
/**
|
||||
* Simple check method that calculates token cost and returns balance info.
|
||||
* The auto-refill logic has been moved to balanceMethods.js to prevent circular dependencies.
|
||||
*/
|
||||
const checkBalanceRecord = async function ({
|
||||
user,
|
||||
model,
|
||||
endpoint,
|
||||
valueKey,
|
||||
tokenType,
|
||||
amount,
|
||||
endpointTokenConfig,
|
||||
}) {
|
||||
const multiplier = getMultiplier({ valueKey, tokenType, model, endpoint, endpointTokenConfig });
|
||||
const tokenCost = amount * multiplier;
|
||||
|
||||
// Retrieve the balance record
|
||||
let record = await Balance.findOne({ user }).lean();
|
||||
if (!record) {
|
||||
logger.debug('[Balance.check] No balance record found for user', { user });
|
||||
return {
|
||||
canSpend: false,
|
||||
balance: 0,
|
||||
tokenCost,
|
||||
};
|
||||
}
|
||||
let balance = record.tokenCredits;
|
||||
|
||||
logger.debug('[Balance.check] Initial state', {
|
||||
user,
|
||||
model,
|
||||
endpoint,
|
||||
valueKey,
|
||||
tokenType,
|
||||
amount,
|
||||
balance,
|
||||
multiplier,
|
||||
endpointTokenConfig: !!endpointTokenConfig,
|
||||
});
|
||||
|
||||
// Only perform auto-refill if spending would bring the balance to 0 or below
|
||||
if (balance - tokenCost <= 0 && record.autoRefillEnabled && record.refillAmount > 0) {
|
||||
const lastRefillDate = new Date(record.lastRefill);
|
||||
const now = new Date();
|
||||
if (
|
||||
isInvalidDate(lastRefillDate) ||
|
||||
now >=
|
||||
addIntervalToDate(lastRefillDate, record.refillIntervalValue, record.refillIntervalUnit)
|
||||
) {
|
||||
try {
|
||||
/** @type {{ rate: number, user: string, balance: number, transaction: import('@librechat/data-schemas').ITransaction}} */
|
||||
const result = await createAutoRefillTransaction({
|
||||
user: user,
|
||||
tokenType: 'credits',
|
||||
context: 'autoRefill',
|
||||
rawAmount: record.refillAmount,
|
||||
});
|
||||
balance = result.balance;
|
||||
} catch (error) {
|
||||
logger.error('[Balance.check] Failed to record transaction for auto-refill', error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.debug('[Balance.check] Token cost', { tokenCost });
|
||||
return { canSpend: balance >= tokenCost, balance, tokenCost };
|
||||
};
|
||||
|
||||
/**
|
||||
* Adds a time interval to a given date.
|
||||
* @param {Date} date - The starting date.
|
||||
* @param {number} value - The numeric value of the interval.
|
||||
* @param {'seconds'|'minutes'|'hours'|'days'|'weeks'|'months'} unit - The unit of time.
|
||||
* @returns {Date} A new Date representing the starting date plus the interval.
|
||||
*/
|
||||
const addIntervalToDate = (date, value, unit) => {
|
||||
const result = new Date(date);
|
||||
switch (unit) {
|
||||
case 'seconds':
|
||||
result.setSeconds(result.getSeconds() + value);
|
||||
break;
|
||||
case 'minutes':
|
||||
result.setMinutes(result.getMinutes() + value);
|
||||
break;
|
||||
case 'hours':
|
||||
result.setHours(result.getHours() + value);
|
||||
break;
|
||||
case 'days':
|
||||
result.setDate(result.getDate() + value);
|
||||
break;
|
||||
case 'weeks':
|
||||
result.setDate(result.getDate() + value * 7);
|
||||
break;
|
||||
case 'months':
|
||||
result.setMonth(result.getMonth() + value);
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
return result;
|
||||
};
|
||||
|
||||
/**
|
||||
* Checks the balance for a user and determines if they can spend a certain amount.
|
||||
* If the user cannot spend the amount, it logs a violation and denies the request.
|
||||
*
|
||||
* @async
|
||||
* @function
|
||||
* @param {Object} params - The function parameters.
|
||||
* @param {ServerRequest} params.req - The Express request object.
|
||||
* @param {Express.Response} params.res - The Express response object.
|
||||
* @param {Object} params.txData - The transaction data.
|
||||
* @param {string} params.txData.user - The user ID or identifier.
|
||||
* @param {('prompt' | 'completion')} params.txData.tokenType - The type of token.
|
||||
* @param {number} params.txData.amount - The amount of tokens.
|
||||
* @param {string} params.txData.model - The model name or identifier.
|
||||
* @param {string} [params.txData.endpointTokenConfig] - The token configuration for the endpoint.
|
||||
* @returns {Promise<boolean>} Throws error if the user cannot spend the amount.
|
||||
* @throws {Error} Throws an error if there's an issue with the balance check.
|
||||
*/
|
||||
const checkBalance = async ({ req, res, txData }) => {
|
||||
const { canSpend, balance, tokenCost } = await checkBalanceRecord(txData);
|
||||
if (canSpend) {
|
||||
return true;
|
||||
}
|
||||
|
||||
const type = ViolationTypes.TOKEN_BALANCE;
|
||||
const errorMessage = {
|
||||
type,
|
||||
balance,
|
||||
tokenCost,
|
||||
promptTokens: txData.amount,
|
||||
};
|
||||
|
||||
if (txData.generations && txData.generations.length > 0) {
|
||||
errorMessage.generations = txData.generations;
|
||||
}
|
||||
|
||||
await logViolation(req, res, type, errorMessage, 0);
|
||||
throw new Error(JSON.stringify(errorMessage));
|
||||
};
|
||||
|
||||
module.exports = {
|
||||
checkBalance,
|
||||
};
|
||||
275
api/models/convoStructure.spec.js
Normal file
275
api/models/convoStructure.spec.js
Normal file
@ -0,0 +1,275 @@
|
||||
const mongoose = require('mongoose');
|
||||
const { buildTree } = require('librechat-data-provider');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
const { getMessages, bulkSaveMessages } = require('./Message');
|
||||
const { Message } = require('~/db/models');
|
||||
|
||||
let mongod;
|
||||
beforeAll(async () => {
|
||||
mongod = await MongoMemoryServer.create();
|
||||
const uri = mongod.getUri();
|
||||
await mongoose.connect(uri);
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await mongoose.disconnect();
|
||||
await mongod.stop();
|
||||
});
|
||||
|
||||
beforeEach(async () => {
|
||||
await Message.deleteMany({});
|
||||
});
|
||||
|
||||
describe('Conversation Structure Tests', () => {
|
||||
test('Conversation folding/corrupting with inconsistent timestamps', async () => {
|
||||
const userId = 'testUser';
|
||||
const conversationId = 'testConversation';
|
||||
|
||||
// Create messages with inconsistent timestamps
|
||||
const messages = [
|
||||
{
|
||||
messageId: 'message0',
|
||||
parentMessageId: null,
|
||||
text: 'Message 0',
|
||||
createdAt: new Date('2023-01-01T00:00:00Z'),
|
||||
},
|
||||
{
|
||||
messageId: 'message1',
|
||||
parentMessageId: 'message0',
|
||||
text: 'Message 1',
|
||||
createdAt: new Date('2023-01-01T00:02:00Z'),
|
||||
},
|
||||
{
|
||||
messageId: 'message2',
|
||||
parentMessageId: 'message1',
|
||||
text: 'Message 2',
|
||||
createdAt: new Date('2023-01-01T00:01:00Z'),
|
||||
}, // Note: Earlier than its parent
|
||||
{
|
||||
messageId: 'message3',
|
||||
parentMessageId: 'message1',
|
||||
text: 'Message 3',
|
||||
createdAt: new Date('2023-01-01T00:03:00Z'),
|
||||
},
|
||||
{
|
||||
messageId: 'message4',
|
||||
parentMessageId: 'message2',
|
||||
text: 'Message 4',
|
||||
createdAt: new Date('2023-01-01T00:04:00Z'),
|
||||
},
|
||||
];
|
||||
|
||||
// Add common properties to all messages
|
||||
messages.forEach((msg) => {
|
||||
msg.conversationId = conversationId;
|
||||
msg.user = userId;
|
||||
msg.isCreatedByUser = false;
|
||||
msg.error = false;
|
||||
msg.unfinished = false;
|
||||
});
|
||||
|
||||
// Save messages with overrideTimestamp omitted (default is false)
|
||||
await bulkSaveMessages(messages, true);
|
||||
|
||||
// Retrieve messages (this will sort by createdAt)
|
||||
const retrievedMessages = await getMessages({ conversationId, user: userId });
|
||||
|
||||
// Build tree
|
||||
const tree = buildTree({ messages: retrievedMessages });
|
||||
|
||||
// Check if the tree is incorrect (folded/corrupted)
|
||||
expect(tree.length).toBeGreaterThan(1); // Should have multiple root messages, indicating corruption
|
||||
});
|
||||
|
||||
test('Fix: Conversation structure maintained with more than 16 messages', async () => {
|
||||
const userId = 'testUser';
|
||||
const conversationId = 'testConversation';
|
||||
|
||||
// Create more than 16 messages
|
||||
const messages = Array.from({ length: 20 }, (_, i) => ({
|
||||
messageId: `message${i}`,
|
||||
parentMessageId: i === 0 ? null : `message${i - 1}`,
|
||||
conversationId,
|
||||
user: userId,
|
||||
text: `Message ${i}`,
|
||||
createdAt: new Date(Date.now() + (i % 2 === 0 ? i * 500000 : -i * 500000)),
|
||||
}));
|
||||
|
||||
// Save messages with new timestamps being generated (message objects ignored)
|
||||
await bulkSaveMessages(messages);
|
||||
|
||||
// Retrieve messages (this will sort by createdAt, but it shouldn't matter now)
|
||||
const retrievedMessages = await getMessages({ conversationId, user: userId });
|
||||
|
||||
// Build tree
|
||||
const tree = buildTree({ messages: retrievedMessages });
|
||||
|
||||
// Check if the tree is correct
|
||||
expect(tree.length).toBe(1); // Should have only one root message
|
||||
let currentNode = tree[0];
|
||||
for (let i = 1; i < 20; i++) {
|
||||
expect(currentNode.children.length).toBe(1);
|
||||
currentNode = currentNode.children[0];
|
||||
expect(currentNode.text).toBe(`Message ${i}`);
|
||||
}
|
||||
expect(currentNode.children.length).toBe(0); // Last message should have no children
|
||||
});
|
||||
|
||||
test('Simulate MongoDB ordering issue with more than 16 messages and close timestamps', async () => {
|
||||
const userId = 'testUser';
|
||||
const conversationId = 'testConversation';
|
||||
|
||||
// Create more than 16 messages with very close timestamps
|
||||
const messages = Array.from({ length: 20 }, (_, i) => ({
|
||||
messageId: `message${i}`,
|
||||
parentMessageId: i === 0 ? null : `message${i - 1}`,
|
||||
conversationId,
|
||||
user: userId,
|
||||
text: `Message ${i}`,
|
||||
createdAt: new Date(Date.now() + (i % 2 === 0 ? i * 1 : -i * 1)),
|
||||
}));
|
||||
|
||||
// Add common properties to all messages
|
||||
messages.forEach((msg) => {
|
||||
msg.isCreatedByUser = false;
|
||||
msg.error = false;
|
||||
msg.unfinished = false;
|
||||
});
|
||||
|
||||
await bulkSaveMessages(messages, true);
|
||||
const retrievedMessages = await getMessages({ conversationId, user: userId });
|
||||
const tree = buildTree({ messages: retrievedMessages });
|
||||
expect(tree.length).toBeGreaterThan(1);
|
||||
});
|
||||
|
||||
test('Fix: Preserve order with more than 16 messages by maintaining original timestamps', async () => {
|
||||
const userId = 'testUser';
|
||||
const conversationId = 'testConversation';
|
||||
|
||||
// Create more than 16 messages with distinct timestamps
|
||||
const messages = Array.from({ length: 20 }, (_, i) => ({
|
||||
messageId: `message${i}`,
|
||||
parentMessageId: i === 0 ? null : `message${i - 1}`,
|
||||
conversationId,
|
||||
user: userId,
|
||||
text: `Message ${i}`,
|
||||
createdAt: new Date(Date.now() + i * 1000), // Ensure each message has a distinct timestamp
|
||||
}));
|
||||
|
||||
// Add common properties to all messages
|
||||
messages.forEach((msg) => {
|
||||
msg.isCreatedByUser = false;
|
||||
msg.error = false;
|
||||
msg.unfinished = false;
|
||||
});
|
||||
|
||||
// Save messages with overriding timestamps (preserve original timestamps)
|
||||
await bulkSaveMessages(messages, true);
|
||||
|
||||
// Retrieve messages (this will sort by createdAt)
|
||||
const retrievedMessages = await getMessages({ conversationId, user: userId });
|
||||
|
||||
// Build tree
|
||||
const tree = buildTree({ messages: retrievedMessages });
|
||||
|
||||
// Check if the tree is correct
|
||||
expect(tree.length).toBe(1); // Should have only one root message
|
||||
let currentNode = tree[0];
|
||||
for (let i = 1; i < 20; i++) {
|
||||
expect(currentNode.children.length).toBe(1);
|
||||
currentNode = currentNode.children[0];
|
||||
expect(currentNode.text).toBe(`Message ${i}`);
|
||||
}
|
||||
expect(currentNode.children.length).toBe(0); // Last message should have no children
|
||||
});
|
||||
|
||||
test('Random order dates between parent and children messages', async () => {
|
||||
const userId = 'testUser';
|
||||
const conversationId = 'testConversation';
|
||||
|
||||
// Create messages with deliberately out-of-order timestamps but sequential creation
|
||||
const messages = [
|
||||
{
|
||||
messageId: 'parent',
|
||||
parentMessageId: null,
|
||||
text: 'Parent Message',
|
||||
createdAt: new Date('2023-01-01T00:00:00Z'), // Make parent earliest
|
||||
},
|
||||
{
|
||||
messageId: 'child1',
|
||||
parentMessageId: 'parent',
|
||||
text: 'Child Message 1',
|
||||
createdAt: new Date('2023-01-01T00:01:00Z'),
|
||||
},
|
||||
{
|
||||
messageId: 'child2',
|
||||
parentMessageId: 'parent',
|
||||
text: 'Child Message 2',
|
||||
createdAt: new Date('2023-01-01T00:02:00Z'),
|
||||
},
|
||||
{
|
||||
messageId: 'grandchild1',
|
||||
parentMessageId: 'child1',
|
||||
text: 'Grandchild Message 1',
|
||||
createdAt: new Date('2023-01-01T00:03:00Z'),
|
||||
},
|
||||
];
|
||||
|
||||
// Add common properties to all messages
|
||||
messages.forEach((msg) => {
|
||||
msg.conversationId = conversationId;
|
||||
msg.user = userId;
|
||||
msg.isCreatedByUser = false;
|
||||
msg.error = false;
|
||||
msg.unfinished = false;
|
||||
});
|
||||
|
||||
// Save messages with overrideTimestamp set to true
|
||||
await bulkSaveMessages(messages, true);
|
||||
|
||||
// Retrieve messages
|
||||
const retrievedMessages = await getMessages({ conversationId, user: userId });
|
||||
|
||||
// Debug log to see what's being returned
|
||||
console.log(
|
||||
'Retrieved Messages:',
|
||||
retrievedMessages.map((msg) => ({
|
||||
messageId: msg.messageId,
|
||||
parentMessageId: msg.parentMessageId,
|
||||
createdAt: msg.createdAt,
|
||||
})),
|
||||
);
|
||||
|
||||
// Build tree
|
||||
const tree = buildTree({ messages: retrievedMessages });
|
||||
|
||||
// Debug log to see the tree structure
|
||||
console.log(
|
||||
'Tree structure:',
|
||||
tree.map((root) => ({
|
||||
messageId: root.messageId,
|
||||
children: root.children.map((child) => ({
|
||||
messageId: child.messageId,
|
||||
children: child.children.map((grandchild) => ({
|
||||
messageId: grandchild.messageId,
|
||||
})),
|
||||
})),
|
||||
})),
|
||||
);
|
||||
|
||||
// Verify the structure before making assertions
|
||||
expect(retrievedMessages.length).toBe(4); // Should have all 4 messages
|
||||
|
||||
// Check if messages are properly linked
|
||||
const parentMsg = retrievedMessages.find((msg) => msg.messageId === 'parent');
|
||||
expect(parentMsg.parentMessageId).toBeNull(); // Parent should have null parentMessageId
|
||||
|
||||
const childMsg1 = retrievedMessages.find((msg) => msg.messageId === 'child1');
|
||||
expect(childMsg1.parentMessageId).toBe('parent');
|
||||
|
||||
// Then check tree structure
|
||||
expect(tree.length).toBe(1); // Should have only one root message
|
||||
expect(tree[0].messageId).toBe('parent');
|
||||
expect(tree[0].children.length).toBe(2); // Should have two children
|
||||
});
|
||||
});
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue
Block a user