Gemini Embedding Support in PraisonAI

This document describes how to use Google's Gemini embedding models with PraisonAI's memory system.

Overview

PraisonAI now supports configurable embedding models through LiteLLM, including Google's Gemini embedding models. This allows you to use models like text-embedding-004 or gemini-embedding-exp-03-07 for semantic search and memory storage.

Configuration

Default Configuration (OpenAI)

By default, PraisonAI uses OpenAI's text-embedding-3-small model:

from praisonaiagents import PraisonAIAgents

agents = PraisonAIAgents(
    agents=[...],
    tasks=[...],
    memory=True  # Uses default OpenAI embeddings
)

Using Gemini Embeddings

To use Gemini embeddings, configure the embedder parameter:

from praisonaiagents import PraisonAIAgents

agents = PraisonAIAgents(
    agents=[...],
    tasks=[...],
    memory=True,
    embedder={
        "provider": "gemini",
        "config": {
            "model": "text-embedding-004"  # or "gemini-embedding-exp-03-07"
        }
    }
)

Prerequisites

API Key: Set your Google API key:
```
export GOOGLE_API_KEY='your-api-key'
```
Dependencies: Ensure you have the memory dependencies installed:
```
pip install "praisonaiagents[memory]"
```

Supported Models

Gemini Embedding Models

text-embedding-004 - Stable embedding model
gemini-embedding-exp-03-07 - Experimental model with advanced features

Task Types (Optional)

Gemini models support task-specific optimizations:

SEMANTIC_SIMILARITY
CLASSIFICATION
CLUSTERING
RETRIEVAL_DOCUMENT
RETRIEVAL_QUERY
QUESTION_ANSWERING
FACT_VERIFICATION

Example with task type:

embedder={
    "provider": "gemini",
    "config": {
        "model": "text-embedding-004",
        "task_type": "RETRIEVAL_DOCUMENT"
    }
}

Complete Example

import os
from praisonaiagents import Agent, Task, PraisonAIAgents

# Ensure Google API key is set
os.environ["GOOGLE_API_KEY"] = "your-api-key"

# Create agents
researcher = Agent(
    name="Researcher",
    role="Information Researcher",
    goal="Store information with semantic understanding",
    llm="gpt-4o-mini"
)

# Create task
task = Task(
    description="Research and store information about AI advancements",
    expected_output="Confirmation of stored information",
    agent=researcher
)

# Initialize with Gemini embeddings
agents = PraisonAIAgents(
    agents=[researcher],
    tasks=[task],
    memory=True,
    embedder={
        "provider": "gemini",
        "config": {
            "model": "text-embedding-004"
        }
    }
)

# Run the agents
agents.start()

Backward Compatibility

The implementation maintains full backward compatibility:

No embedder config: Uses default OpenAI text-embedding-3-small
Existing configs: Continue to work without modification
Other providers: Supports any embedding model available through LiteLLM

Memory Configuration with Mem0

When using Mem0 as the memory provider, the embedder configuration is passed through:

memory_config = {
    "provider": "mem0",
    "config": {
        "api_key": "your-mem0-key",
        "embedder": {
            "provider": "gemini",
            "config": {
                "model": "text-embedding-004"
            }
        }
    }
}

Troubleshooting

API Key Error: Ensure GOOGLE_API_KEY is set in your environment
Model Not Found: Verify the model name is correct
LiteLLM Error: Ensure litellm is installed and up to date
Embedding Size Mismatch: When switching models, you may need to clear the vector database as different models produce different embedding dimensions

3.9 KiB Raw Permalink Blame History