This commit implements a clean architecture where environment variables are only read, never set, ensuring compatibility with other ComfyUI packages. The changes: - Remove all code that sets global environment variables - Clean up config.yaml to remove API key placeholders - Add documentation for safe environment handling - Add comprehensive tests to verify non-polluting behavior - Ensure backward compatibility for other packages The commit includes detailed docs explaining the clean architecture and safety guarantees for package compatibility.
227 lines
9.7 KiB
Markdown
Executable File
227 lines
9.7 KiB
Markdown
Executable File
# ComfyUI_LiteLLM
|
|
|
|
ComfyUI_LiteLLM is an addon for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) that provides nodes for interacting with various language models using the [LiteLLM](https://github.com/BerriAI/litellm) library directly within the ComfyUI interface. This addon serves as a bridge between ComfyUI and LiteLLM, allowing you to leverage any model supported by LiteLLM in your ComfyUI workflows.
|
|
|
|
## Features
|
|
|
|
This addon provides ComfyUI nodes that interface with LiteLLM's functionality, including:
|
|
|
|
- `LiteLLMModelProvider`: Selects LiteLLM-supported models for use in other nodes
|
|
- `LiteLLMCompletion`: Generates completions using any LiteLLM-supported model
|
|
- `LiteLLMCompletionPrePend`: Similar to LiteLLMCompletion but with a pre-pended prompt
|
|
- `LiteLLMCompletionListOfPrompts`: Generates completions for a list of prompts
|
|
- `CreateReflectionFilter` and `FirstCodeBlockReflectionFilter`: Creates reflection filters for use with `LiteLLMCompletionWithReflectionFilter`
|
|
- `LiteLLMCompletionWithReflectionFilter`: Generates completions with a reflection filter
|
|
- `LiteLLMCompletionProvider`: Provides a completion function for use in other nodes
|
|
- `LiteLLMMessage`: Creates LiteLLM messages
|
|
- `ShowLastMessage` and `ShowMessages`: Display message content in the UI
|
|
- `ListToMessages` and `MessagesToList`: Convert between list and message formats
|
|
- `AppendMessages`: Combine multiple message sets
|
|
- `MessagesToText` and `TextToMessages`: Convert between text and message formats
|
|
- `LLLMInterpretUrl`: Fetch and process content from URLs
|
|
|
|
### **🚀 LightRAG Integration**
|
|
Complete graph-based Retrieval-Augmented Generation with any LiteLLM model:
|
|
- **DocumentProcessor**: Process documents using LightRAG with any LiteLLM-supported model
|
|
- **QueryNode**: Query LightRAG knowledge graphs with multiple retrieval modes (naive, local, global, hybrid)
|
|
- **AgentMemoryProvider**: Create memory functions for LiteLLM agents
|
|
- **LinuxMemoryDirectory**: High-performance memory directory for faster processing
|
|
- **Local Embeddings**: Uses high-quality Stella 1.5B model locally (no OpenAI dependency)
|
|
- **Automatic Compatibility**: LiteLLM completion providers work seamlessly with LightRAG's async interface
|
|
|
|
## Supported Models
|
|
|
|
This addon supports all models that LiteLLM can interface with. This includes, but is not limited to:
|
|
|
|
- OpenAI models
|
|
- Anthropic models (Claude)
|
|
- Google models (PaLM, Gemini)
|
|
- Azure OpenAI
|
|
- AWS Bedrock models
|
|
- Cohere models
|
|
- AI21 models
|
|
- And many more
|
|
|
|
For the most up-to-date list of supported models, please refer to the [LiteLLM documentation](https://docs.litellm.ai/docs/).
|
|
|
|
## 🧠 LightRAG Integration - Graph-Based RAG with Any LLM
|
|
|
|
This addon features a **production-ready integration** between [LightRAG](https://github.com/HKUDS/LightRAG) (Graph-based Retrieval-Augmented Generation) and LiteLLM completion providers. Use any LiteLLM-supported model with LightRAG's powerful graph-based knowledge indexing and retrieval.
|
|
|
|
### 🎯 Key Features
|
|
|
|
- **🔄 Universal Model Support**: Use any LiteLLM model (OpenAI, Anthropic, Google, Kluster, etc.) with LightRAG
|
|
- **🏠 Local Embeddings**: High-quality Stella 1.5B model runs locally (no OpenAI API calls)
|
|
- **📊 Graph-Based Knowledge**: Build intelligent knowledge graphs from documents with entity and relationship extraction
|
|
- **🚀 Incremental Processing**: Add documents progressively while maintaining graph relationships
|
|
- **🔍 Multiple Query Modes**: naive, local, global, and hybrid retrieval strategies
|
|
- **💾 Persistent Storage**: Working directories maintain state across sessions
|
|
- **🧠 Agent Memory**: Seamless integration with LiteLLM agents for contextual memory
|
|
- **⚡ High Performance**: Linux memory directory support for faster processing
|
|
|
|
### 🔧 How It Works
|
|
|
|
The integration includes an **automatic compatibility wrapper** that bridges the gap between:
|
|
- **LiteLLM**: Synchronous `completion_function(prompt)`
|
|
- **LightRAG**: Async `llm_func(prompt, system_prompt, history_messages, **kwargs)`
|
|
|
|
The wrapper intelligently combines system prompts, conversation history, and user prompts into the format expected by your chosen LLM.
|
|
|
|
### 📋 Available Nodes
|
|
|
|
| Node | Purpose | Key Features |
|
|
|------|---------|--------------|
|
|
| **DocumentProcessor** | Process documents into knowledge graphs | Chunking, entity extraction, relationship mapping |
|
|
| **QueryNode** | Query the knowledge graph | 4 retrieval modes, contextual responses |
|
|
| **AgentMemoryProvider** | Memory for LiteLLM agents | Query refinement, context retrieval |
|
|
| **LinuxMemoryDirectory** | High-speed storage | `/dev/shm` for faster processing |
|
|
|
|
### 🛠️ Example Workflows
|
|
|
|
**Basic Document Processing:**
|
|
```
|
|
LiteLLMModelProvider → LiteLLMCompletionProvider → DocumentProcessor
|
|
```
|
|
|
|
**Complete RAG Pipeline:**
|
|
```
|
|
LiteLLMModelProvider → LiteLLMCompletionProvider → DocumentProcessor → QueryNode
|
|
```
|
|
|
|
**Agent with Memory:**
|
|
```
|
|
DocumentProcessor → AgentMemoryProvider → [Your Agent Workflow]
|
|
```
|
|
|
|
### 🧪 Thoroughly Tested
|
|
|
|
- ✅ **Real API Integration**: Tested with actual Kluster API calls
|
|
- ✅ **Multi-Document Processing**: Handles incremental document addition
|
|
- ✅ **Cross-Document Queries**: Finds relationships across multiple documents
|
|
- ✅ **State Persistence**: Working directories maintain state across sessions
|
|
- ✅ **Error Handling**: Robust error recovery and informative messages
|
|
- ✅ **No OpenAI Dependencies**: Local embeddings eliminate authentication errors
|
|
|
|
### 🔧 Local Embeddings Setup
|
|
|
|
**Requirements:**
|
|
```bash
|
|
pip install sentence-transformers
|
|
```
|
|
|
|
**DocumentProcessor Configuration:**
|
|
- `embedding_provider`: "local" (default - uses Stella 1.5B model)
|
|
- `embedding_model`: "NovaSearch/stella_en_1.5B_v5" (default)
|
|
- `embedding_dimension`: 1024 (recommended, options: 256-8192)
|
|
- `override_callable`: **REQUIRED** - Connect your LiteLLM completion provider
|
|
|
|
**Benefits:**
|
|
- ✅ High-quality embeddings (state-of-the-art MTEB scores)
|
|
- ✅ Complete privacy - no data sent to external APIs
|
|
- ✅ Fast local inference after initial model download (~3GB)
|
|
- ✅ No API costs for embeddings
|
|
- ✅ Works offline after setup
|
|
|
|
## Installation
|
|
|
|
### Manual Installation
|
|
|
|
1. Clone this repository into your `custom_nodes` directory in your ComfyUI installation.
|
|
2. Install the required dependencies by running `pip install -r requirements.txt`.
|
|
3. Restart ComfyUI.
|
|
|
|
### Installation with ComfyUI-Manager
|
|
|
|
1. Add this addon's repository URL to your ComfyUI-Manager configuration.
|
|
2. Install the addon through ComfyUI-Manager.
|
|
3. Restart ComfyUI.
|
|
|
|
## Configuration
|
|
|
|
### Environment Variables (Recommended for Production)
|
|
|
|
For production use in ComfyUI, set your API keys as system environment variables:
|
|
|
|
```bash
|
|
export OPENAI_API_KEY="your-openai-api-key"
|
|
export ANTHROPIC_API_KEY="your-anthropic-api-key"
|
|
export KLUSTER_API_KEY="your-kluster-api-key"
|
|
# ... other API keys as needed
|
|
```
|
|
|
|
### Configuration Files
|
|
|
|
The addon uses a clean architecture that separates configuration from secrets:
|
|
|
|
- **`config.yaml`**: Contains non-sensitive settings like base URLs and paths. Never put API keys here.
|
|
- **`.env`**: Used only for testing and development. Not used in production ComfyUI.
|
|
|
|
### Environment Safety
|
|
|
|
This addon follows best practices for environment management:
|
|
- ✅ **Never pollutes global environment variables** - won't interfere with other ComfyUI custom nodes
|
|
- ✅ **Reads environment variables safely** - uses system environment first, falls back to `.env` for testing
|
|
- ✅ **No API keys in config files** - keeps secrets separate from configuration
|
|
- ✅ **Compatible with other node packs** - won't override API keys that other packages expect
|
|
|
|
### For Developers and Testing
|
|
|
|
If you're developing or testing, you can create a `.env` file in the addon directory:
|
|
|
|
```bash
|
|
# ComfyUI_LiteLLM/.env
|
|
OPENAI_API_KEY=your-test-key
|
|
KLUSTER_API_KEY=your-test-key
|
|
# etc.
|
|
```
|
|
|
|
This `.env` file is only used when system environment variables aren't available, making it perfect for development without affecting production deployments.
|
|
|
|
## Usage
|
|
|
|
After installation, the LiteLLM nodes will be available in the ComfyUI node editor under the "ETK/LLM/LiteLLM" category. You can use these nodes to integrate any LiteLLM-supported language model into your ComfyUI workflows.
|
|
|
|
### Basic LLM Usage:
|
|
1. Set up your API keys as system environment variables (see Configuration section above)
|
|
2. Select the desired model in the `LiteLLMModelProvider` node
|
|
3. Connect the nodes in your workflow
|
|
|
|
### LightRAG Document Processing:
|
|
1. **Install sentence-transformers**: `pip install sentence-transformers`
|
|
2. **Create workflow**: `LiteLLMModelProvider` → `LiteLLMCompletionProvider` → `DocumentProcessor`
|
|
3. **Connect completion provider** to DocumentProcessor's `override_callable` input (REQUIRED)
|
|
4. **Set embedding provider** to "local" for high-quality local embeddings
|
|
5. **First run** downloads Stella model (~3GB), then cached for future use
|
|
|
|
The addon will automatically use your environment variables for authentication without interfering with other ComfyUI custom nodes that may also use the same API services.
|
|
|
|
## Contributing
|
|
|
|
Contributions are welcome! If you encounter any issues or have suggestions for improvements, please open an issue or submit a pull request.
|
|
|
|
### Running Tests
|
|
|
|
The addon includes comprehensive tests that can be run with:
|
|
|
|
```bash
|
|
cd ComfyUI_LiteLLM
|
|
python3 -m pytest tests/ -v
|
|
```
|
|
|
|
Tests use the `.env` file for API keys when available, or skip tests requiring real API calls if keys aren't configured.
|
|
|
|
## License
|
|
|
|
This addon is released under the [MIT License](LICENSE).
|
|
|
|
## Dependencies
|
|
|
|
- pyyaml
|
|
- addict
|
|
- litellm
|
|
- boto3
|
|
- lightrag-hku (for LightRAG graph-based RAG integration)
|
|
- sentence-transformers (for local embeddings with Stella 1.5B model)
|
|
|
|
For the full list of dependencies, please refer to the `requirements.txt` file.
|