feat(graphiti): add Google AI as LLM and embedding provider
Add full Google AI (Gemini) support for Graphiti memory system: Backend: - Add google-generativeai dependency to requirements.txt - Create GoogleEmbedder class with text-embedding-004 default model - Create GoogleLLMClient class with gemini-2.0-flash default model - Add GOOGLE to LLMProvider and EmbedderProvider enums - Add google_api_key, google_llm_model, google_embedding_model config - Update factory to create Google LLM client and embedder - Add validation for Google provider configuration Frontend: - Add 'google' to GraphitiLLMProvider and GraphitiEmbeddingProvider types - Add Google AI option to LLM provider dropdown in Setup Wizard - Add Google AI option to embedding provider dropdown - Add Google API key input field with link to Google AI Studio - Update MemoryBackendSection and SecuritySettings components - Update env-handlers to save GOOGLE_API_KEY, GOOGLE_LLM_MODEL, and GOOGLE_EMBEDDING_MODEL to .env files This allows users to use Google's Gemini models for both LLM operations (graph extraction, search, reasoning) and embeddings in Graphiti memory.
This commit is contained in:
adryserage
@@ -81,6 +81,21 @@ auto-claude/.venv/bin/pytest tests/ -m "not slow"
|
||||
python auto-claude/validate_spec.py --spec-dir auto-claude/specs/001-feature --checkpoint all
|
||||
```
|
||||
|
||||
### Releases
|
||||
```bash
|
||||
# Automated version bump and release (recommended)
|
||||
node scripts/bump-version.js patch # 2.5.5 -> 2.5.6
|
||||
node scripts/bump-version.js minor # 2.5.5 -> 2.6.0
|
||||
node scripts/bump-version.js major # 2.5.5 -> 3.0.0
|
||||
node scripts/bump-version.js 2.6.0 # Set specific version
|
||||
|
||||
# Then push to trigger GitHub release workflows
|
||||
git push origin main
|
||||
git push origin v2.6.0
|
||||
```
|
||||
|
||||
See [RELEASE.md](RELEASE.md) for detailed release process documentation.
|
||||
|
||||
## Architecture
|
||||
|
||||
### Core Pipeline
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "auto-claude-ui",
|
||||
"version": "2.5.0",
|
||||
"version": "2.5.5",
|
||||
"description": "Desktop UI for Auto Claude autonomous coding framework",
|
||||
"main": "./out/main/index.js",
|
||||
"author": "Auto Claude Team",
|
||||
|
||||
@@ -65,6 +65,41 @@ export function registerEnvHandlers(
|
||||
if (config.graphitiEnabled !== undefined) {
|
||||
existingVars['GRAPHITI_ENABLED'] = config.graphitiEnabled ? 'true' : 'false';
|
||||
}
|
||||
// Graphiti Provider Configuration
|
||||
if (config.graphitiProviderConfig) {
|
||||
const pc = config.graphitiProviderConfig;
|
||||
if (pc.llmProvider) existingVars['GRAPHITI_LLM_PROVIDER'] = pc.llmProvider;
|
||||
if (pc.embeddingProvider) existingVars['GRAPHITI_EMBEDDER_PROVIDER'] = pc.embeddingProvider;
|
||||
// OpenAI
|
||||
if (pc.openaiApiKey) existingVars['OPENAI_API_KEY'] = pc.openaiApiKey;
|
||||
if (pc.openaiModel) existingVars['OPENAI_MODEL'] = pc.openaiModel;
|
||||
if (pc.openaiEmbeddingModel) existingVars['OPENAI_EMBEDDING_MODEL'] = pc.openaiEmbeddingModel;
|
||||
// Anthropic
|
||||
if (pc.anthropicApiKey) existingVars['ANTHROPIC_API_KEY'] = pc.anthropicApiKey;
|
||||
if (pc.anthropicModel) existingVars['GRAPHITI_ANTHROPIC_MODEL'] = pc.anthropicModel;
|
||||
// Azure OpenAI
|
||||
if (pc.azureOpenaiApiKey) existingVars['AZURE_OPENAI_API_KEY'] = pc.azureOpenaiApiKey;
|
||||
if (pc.azureOpenaiBaseUrl) existingVars['AZURE_OPENAI_BASE_URL'] = pc.azureOpenaiBaseUrl;
|
||||
if (pc.azureOpenaiLlmDeployment) existingVars['AZURE_OPENAI_LLM_DEPLOYMENT'] = pc.azureOpenaiLlmDeployment;
|
||||
if (pc.azureOpenaiEmbeddingDeployment) existingVars['AZURE_OPENAI_EMBEDDING_DEPLOYMENT'] = pc.azureOpenaiEmbeddingDeployment;
|
||||
// Voyage
|
||||
if (pc.voyageApiKey) existingVars['VOYAGE_API_KEY'] = pc.voyageApiKey;
|
||||
if (pc.voyageEmbeddingModel) existingVars['VOYAGE_EMBEDDING_MODEL'] = pc.voyageEmbeddingModel;
|
||||
// Google
|
||||
if (pc.googleApiKey) existingVars['GOOGLE_API_KEY'] = pc.googleApiKey;
|
||||
if (pc.googleLlmModel) existingVars['GOOGLE_LLM_MODEL'] = pc.googleLlmModel;
|
||||
if (pc.googleEmbeddingModel) existingVars['GOOGLE_EMBEDDING_MODEL'] = pc.googleEmbeddingModel;
|
||||
// Ollama
|
||||
if (pc.ollamaBaseUrl) existingVars['OLLAMA_BASE_URL'] = pc.ollamaBaseUrl;
|
||||
if (pc.ollamaLlmModel) existingVars['OLLAMA_LLM_MODEL'] = pc.ollamaLlmModel;
|
||||
if (pc.ollamaEmbeddingModel) existingVars['OLLAMA_EMBEDDING_MODEL'] = pc.ollamaEmbeddingModel;
|
||||
if (pc.ollamaEmbeddingDim) existingVars['OLLAMA_EMBEDDING_DIM'] = String(pc.ollamaEmbeddingDim);
|
||||
// FalkorDB
|
||||
if (pc.falkorDbHost) existingVars['GRAPHITI_FALKORDB_HOST'] = pc.falkorDbHost;
|
||||
if (pc.falkorDbPort) existingVars['GRAPHITI_FALKORDB_PORT'] = String(pc.falkorDbPort);
|
||||
if (pc.falkorDbPassword) existingVars['GRAPHITI_FALKORDB_PASSWORD'] = pc.falkorDbPassword;
|
||||
}
|
||||
// Legacy fields (still supported)
|
||||
if (config.openaiApiKey !== undefined) {
|
||||
existingVars['OPENAI_API_KEY'] = config.openaiApiKey;
|
||||
}
|
||||
@@ -116,9 +151,45 @@ ${existingVars['ENABLE_FANCY_UI'] !== undefined ? `ENABLE_FANCY_UI=${existingVar
|
||||
|
||||
# =============================================================================
|
||||
# GRAPHITI MEMORY INTEGRATION (OPTIONAL)
|
||||
# Multi-provider support: OpenAI, Anthropic, Azure OpenAI, Ollama, Voyage
|
||||
# =============================================================================
|
||||
${existingVars['GRAPHITI_ENABLED'] ? `GRAPHITI_ENABLED=${existingVars['GRAPHITI_ENABLED']}` : '# GRAPHITI_ENABLED=false'}
|
||||
|
||||
# Provider Selection
|
||||
${existingVars['GRAPHITI_LLM_PROVIDER'] ? `GRAPHITI_LLM_PROVIDER=${existingVars['GRAPHITI_LLM_PROVIDER']}` : '# GRAPHITI_LLM_PROVIDER=openai'}
|
||||
${existingVars['GRAPHITI_EMBEDDER_PROVIDER'] ? `GRAPHITI_EMBEDDER_PROVIDER=${existingVars['GRAPHITI_EMBEDDER_PROVIDER']}` : '# GRAPHITI_EMBEDDER_PROVIDER=openai'}
|
||||
|
||||
# OpenAI Settings
|
||||
${existingVars['OPENAI_API_KEY'] ? `OPENAI_API_KEY=${existingVars['OPENAI_API_KEY']}` : '# OPENAI_API_KEY='}
|
||||
${existingVars['OPENAI_MODEL'] ? `OPENAI_MODEL=${existingVars['OPENAI_MODEL']}` : '# OPENAI_MODEL=gpt-4o-mini'}
|
||||
${existingVars['OPENAI_EMBEDDING_MODEL'] ? `OPENAI_EMBEDDING_MODEL=${existingVars['OPENAI_EMBEDDING_MODEL']}` : '# OPENAI_EMBEDDING_MODEL=text-embedding-3-small'}
|
||||
|
||||
# Anthropic Settings (LLM only - use with Voyage or OpenAI for embeddings)
|
||||
${existingVars['ANTHROPIC_API_KEY'] ? `ANTHROPIC_API_KEY=${existingVars['ANTHROPIC_API_KEY']}` : '# ANTHROPIC_API_KEY='}
|
||||
${existingVars['GRAPHITI_ANTHROPIC_MODEL'] ? `GRAPHITI_ANTHROPIC_MODEL=${existingVars['GRAPHITI_ANTHROPIC_MODEL']}` : '# GRAPHITI_ANTHROPIC_MODEL=claude-sonnet-4-5-latest'}
|
||||
|
||||
# Azure OpenAI Settings
|
||||
${existingVars['AZURE_OPENAI_API_KEY'] ? `AZURE_OPENAI_API_KEY=${existingVars['AZURE_OPENAI_API_KEY']}` : '# AZURE_OPENAI_API_KEY='}
|
||||
${existingVars['AZURE_OPENAI_BASE_URL'] ? `AZURE_OPENAI_BASE_URL=${existingVars['AZURE_OPENAI_BASE_URL']}` : '# AZURE_OPENAI_BASE_URL='}
|
||||
${existingVars['AZURE_OPENAI_LLM_DEPLOYMENT'] ? `AZURE_OPENAI_LLM_DEPLOYMENT=${existingVars['AZURE_OPENAI_LLM_DEPLOYMENT']}` : '# AZURE_OPENAI_LLM_DEPLOYMENT='}
|
||||
${existingVars['AZURE_OPENAI_EMBEDDING_DEPLOYMENT'] ? `AZURE_OPENAI_EMBEDDING_DEPLOYMENT=${existingVars['AZURE_OPENAI_EMBEDDING_DEPLOYMENT']}` : '# AZURE_OPENAI_EMBEDDING_DEPLOYMENT='}
|
||||
|
||||
# Voyage AI Settings (Embeddings only - great with Anthropic)
|
||||
${existingVars['VOYAGE_API_KEY'] ? `VOYAGE_API_KEY=${existingVars['VOYAGE_API_KEY']}` : '# VOYAGE_API_KEY='}
|
||||
${existingVars['VOYAGE_EMBEDDING_MODEL'] ? `VOYAGE_EMBEDDING_MODEL=${existingVars['VOYAGE_EMBEDDING_MODEL']}` : '# VOYAGE_EMBEDDING_MODEL=voyage-3'}
|
||||
|
||||
# Google AI Settings (LLM and Embeddings - Gemini)
|
||||
${existingVars['GOOGLE_API_KEY'] ? `GOOGLE_API_KEY=${existingVars['GOOGLE_API_KEY']}` : '# GOOGLE_API_KEY='}
|
||||
${existingVars['GOOGLE_LLM_MODEL'] ? `GOOGLE_LLM_MODEL=${existingVars['GOOGLE_LLM_MODEL']}` : '# GOOGLE_LLM_MODEL=gemini-2.0-flash'}
|
||||
${existingVars['GOOGLE_EMBEDDING_MODEL'] ? `GOOGLE_EMBEDDING_MODEL=${existingVars['GOOGLE_EMBEDDING_MODEL']}` : '# GOOGLE_EMBEDDING_MODEL=text-embedding-004'}
|
||||
|
||||
# Ollama Settings (Local - free)
|
||||
${existingVars['OLLAMA_BASE_URL'] ? `OLLAMA_BASE_URL=${existingVars['OLLAMA_BASE_URL']}` : '# OLLAMA_BASE_URL=http://localhost:11434'}
|
||||
${existingVars['OLLAMA_LLM_MODEL'] ? `OLLAMA_LLM_MODEL=${existingVars['OLLAMA_LLM_MODEL']}` : '# OLLAMA_LLM_MODEL='}
|
||||
${existingVars['OLLAMA_EMBEDDING_MODEL'] ? `OLLAMA_EMBEDDING_MODEL=${existingVars['OLLAMA_EMBEDDING_MODEL']}` : '# OLLAMA_EMBEDDING_MODEL='}
|
||||
${existingVars['OLLAMA_EMBEDDING_DIM'] ? `OLLAMA_EMBEDDING_DIM=${existingVars['OLLAMA_EMBEDDING_DIM']}` : '# OLLAMA_EMBEDDING_DIM=768'}
|
||||
|
||||
# FalkorDB Connection
|
||||
${existingVars['GRAPHITI_FALKORDB_HOST'] ? `GRAPHITI_FALKORDB_HOST=${existingVars['GRAPHITI_FALKORDB_HOST']}` : '# GRAPHITI_FALKORDB_HOST=localhost'}
|
||||
${existingVars['GRAPHITI_FALKORDB_PORT'] ? `GRAPHITI_FALKORDB_PORT=${existingVars['GRAPHITI_FALKORDB_PORT']}` : '# GRAPHITI_FALKORDB_PORT=6380'}
|
||||
${existingVars['GRAPHITI_FALKORDB_PASSWORD'] ? `GRAPHITI_FALKORDB_PASSWORD=${existingVars['GRAPHITI_FALKORDB_PASSWORD']}` : '# GRAPHITI_FALKORDB_PASSWORD='}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -146,7 +146,7 @@ export function MemoryBackendSection({
|
||||
onValueChange={(value) => onUpdateConfig({
|
||||
graphitiProviderConfig: {
|
||||
...envConfig.graphitiProviderConfig,
|
||||
llmProvider: value as 'openai' | 'anthropic' | 'google' | 'groq',
|
||||
llmProvider: value as 'openai' | 'anthropic' | 'azure_openai' | 'ollama',
|
||||
embeddingProvider: envConfig.graphitiProviderConfig?.embeddingProvider || 'openai',
|
||||
}
|
||||
})}
|
||||
@@ -155,10 +155,11 @@ export function MemoryBackendSection({
|
||||
<SelectValue placeholder="Select LLM provider" />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="openai">OpenAI (GPT-5-mini)</SelectItem>
|
||||
<SelectItem value="openai">OpenAI (GPT-4o-mini)</SelectItem>
|
||||
<SelectItem value="anthropic">Anthropic (Claude)</SelectItem>
|
||||
<SelectItem value="google">Google (Gemini)</SelectItem>
|
||||
<SelectItem value="groq">Groq (Llama)</SelectItem>
|
||||
<SelectItem value="google">Google AI (Gemini)</SelectItem>
|
||||
<SelectItem value="azure_openai">Azure OpenAI</SelectItem>
|
||||
<SelectItem value="ollama">Ollama (Local)</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
@@ -175,7 +176,7 @@ export function MemoryBackendSection({
|
||||
graphitiProviderConfig: {
|
||||
...envConfig.graphitiProviderConfig,
|
||||
llmProvider: envConfig.graphitiProviderConfig?.llmProvider || 'openai',
|
||||
embeddingProvider: value as 'openai' | 'voyage' | 'google' | 'huggingface',
|
||||
embeddingProvider: value as 'openai' | 'voyage' | 'azure_openai' | 'ollama',
|
||||
}
|
||||
})}
|
||||
>
|
||||
@@ -185,8 +186,9 @@ export function MemoryBackendSection({
|
||||
<SelectContent>
|
||||
<SelectItem value="openai">OpenAI</SelectItem>
|
||||
<SelectItem value="voyage">Voyage AI</SelectItem>
|
||||
<SelectItem value="google">Google</SelectItem>
|
||||
<SelectItem value="huggingface">HuggingFace (Local)</SelectItem>
|
||||
<SelectItem value="google">Google AI</SelectItem>
|
||||
<SelectItem value="azure_openai">Azure OpenAI</SelectItem>
|
||||
<SelectItem value="ollama">Ollama (Local)</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
|
||||
@@ -164,7 +164,7 @@ export function SecuritySettings({
|
||||
updateEnvConfig({
|
||||
graphitiProviderConfig: {
|
||||
...currentConfig,
|
||||
llmProvider: value as 'openai' | 'anthropic' | 'google' | 'groq',
|
||||
llmProvider: value as 'openai' | 'anthropic' | 'azure_openai' | 'ollama',
|
||||
}
|
||||
});
|
||||
}}
|
||||
@@ -173,10 +173,11 @@ export function SecuritySettings({
|
||||
<SelectValue placeholder="Select LLM provider" />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="openai">OpenAI (GPT-5-mini)</SelectItem>
|
||||
<SelectItem value="openai">OpenAI (GPT-4o-mini)</SelectItem>
|
||||
<SelectItem value="anthropic">Anthropic (Claude)</SelectItem>
|
||||
<SelectItem value="google">Google (Gemini)</SelectItem>
|
||||
<SelectItem value="groq">Groq (Llama)</SelectItem>
|
||||
<SelectItem value="google">Google AI (Gemini)</SelectItem>
|
||||
<SelectItem value="azure_openai">Azure OpenAI</SelectItem>
|
||||
<SelectItem value="ollama">Ollama (Local)</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
@@ -197,7 +198,7 @@ export function SecuritySettings({
|
||||
updateEnvConfig({
|
||||
graphitiProviderConfig: {
|
||||
...currentConfig,
|
||||
embeddingProvider: value as 'openai' | 'voyage' | 'google' | 'huggingface',
|
||||
embeddingProvider: value as 'openai' | 'voyage' | 'azure_openai' | 'ollama',
|
||||
}
|
||||
});
|
||||
}}
|
||||
@@ -208,8 +209,9 @@ export function SecuritySettings({
|
||||
<SelectContent>
|
||||
<SelectItem value="openai">OpenAI</SelectItem>
|
||||
<SelectItem value="voyage">Voyage AI</SelectItem>
|
||||
<SelectItem value="google">Google</SelectItem>
|
||||
<SelectItem value="huggingface">HuggingFace (Local)</SelectItem>
|
||||
<SelectItem value="google">Google AI</SelectItem>
|
||||
<SelectItem value="azure_openai">Azure OpenAI</SelectItem>
|
||||
<SelectItem value="ollama">Ollama (Local)</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
|
||||
@@ -186,31 +186,61 @@ export interface GraphitiConnectionTestResult {
|
||||
}
|
||||
|
||||
// Graphiti Provider Types (Memory System V2)
|
||||
export type GraphitiProviderType = 'openai' | 'anthropic' | 'google' | 'groq' | 'ollama';
|
||||
export type GraphitiEmbeddingProvider = 'openai' | 'voyage' | 'google' | 'huggingface' | 'ollama';
|
||||
// LLM Providers: OpenAI, Anthropic, Azure OpenAI, Ollama (local), Google, Groq
|
||||
export type GraphitiLLMProvider = 'openai' | 'anthropic' | 'azure_openai' | 'ollama' | 'google' | 'groq';
|
||||
// Embedding Providers: OpenAI, Voyage AI, Azure OpenAI, Ollama (local), Google, HuggingFace
|
||||
export type GraphitiEmbeddingProvider = 'openai' | 'voyage' | 'azure_openai' | 'ollama' | 'google' | 'huggingface';
|
||||
|
||||
// Legacy type alias for backward compatibility
|
||||
export type GraphitiProviderType = GraphitiLLMProvider;
|
||||
|
||||
export interface GraphitiProviderConfig {
|
||||
// LLM Provider
|
||||
llmProvider: GraphitiProviderType;
|
||||
llmProvider: GraphitiLLMProvider;
|
||||
llmModel?: string; // Model name, uses provider default if not specified
|
||||
|
||||
// Embedding Provider
|
||||
embeddingProvider: GraphitiEmbeddingProvider;
|
||||
embeddingModel?: string; // Embedding model, uses provider default if not specified
|
||||
|
||||
// Provider-specific API keys (stored securely)
|
||||
// OpenAI settings
|
||||
openaiApiKey?: string;
|
||||
anthropicApiKey?: string;
|
||||
googleApiKey?: string;
|
||||
groqApiKey?: string;
|
||||
voyageApiKey?: string;
|
||||
openaiModel?: string;
|
||||
openaiEmbeddingModel?: string;
|
||||
|
||||
// Ollama-specific config (local LLM, no API key required)
|
||||
// Anthropic settings (LLM only - needs separate embedder)
|
||||
anthropicApiKey?: string;
|
||||
anthropicModel?: string;
|
||||
|
||||
// Azure OpenAI settings
|
||||
azureOpenaiApiKey?: string;
|
||||
azureOpenaiBaseUrl?: string;
|
||||
azureOpenaiLlmDeployment?: string;
|
||||
azureOpenaiEmbeddingDeployment?: string;
|
||||
|
||||
// Voyage AI settings (embeddings only - commonly used with Anthropic)
|
||||
voyageApiKey?: string;
|
||||
voyageEmbeddingModel?: string;
|
||||
|
||||
// Google AI settings (LLM and embeddings)
|
||||
googleApiKey?: string;
|
||||
googleLlmModel?: string;
|
||||
googleEmbeddingModel?: string;
|
||||
|
||||
// Ollama settings (local LLM, no API key required)
|
||||
ollamaBaseUrl?: string; // Default: http://localhost:11434
|
||||
ollamaLlmModel?: string;
|
||||
ollamaEmbeddingModel?: string;
|
||||
ollamaEmbeddingDim?: number;
|
||||
|
||||
// Groq settings
|
||||
groqApiKey?: string;
|
||||
groqModel?: string;
|
||||
|
||||
// HuggingFace settings (embeddings only)
|
||||
huggingfaceApiKey?: string;
|
||||
huggingfaceEmbeddingModel?: string;
|
||||
|
||||
// FalkorDB connection (required for all providers)
|
||||
falkorDbHost?: string;
|
||||
falkorDbPort?: number;
|
||||
|
||||
@@ -82,6 +82,7 @@ class LLMProvider(str, Enum):
|
||||
ANTHROPIC = "anthropic"
|
||||
AZURE_OPENAI = "azure_openai"
|
||||
OLLAMA = "ollama"
|
||||
GOOGLE = "google"
|
||||
|
||||
|
||||
class EmbedderProvider(str, Enum):
|
||||
@@ -91,6 +92,7 @@ class EmbedderProvider(str, Enum):
|
||||
VOYAGE = "voyage"
|
||||
AZURE_OPENAI = "azure_openai"
|
||||
OLLAMA = "ollama"
|
||||
GOOGLE = "google"
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -128,6 +130,11 @@ class GraphitiConfig:
|
||||
voyage_api_key: str = ""
|
||||
voyage_embedding_model: str = "voyage-3"
|
||||
|
||||
# Google AI settings (LLM and embeddings)
|
||||
google_api_key: str = ""
|
||||
google_llm_model: str = "gemini-2.0-flash"
|
||||
google_embedding_model: str = "text-embedding-004"
|
||||
|
||||
# Ollama settings (local)
|
||||
ollama_base_url: str = DEFAULT_OLLAMA_BASE_URL
|
||||
ollama_llm_model: str = ""
|
||||
@@ -189,6 +196,11 @@ class GraphitiConfig:
|
||||
voyage_api_key = os.environ.get("VOYAGE_API_KEY", "")
|
||||
voyage_embedding_model = os.environ.get("VOYAGE_EMBEDDING_MODEL", "voyage-3")
|
||||
|
||||
# Google AI settings
|
||||
google_api_key = os.environ.get("GOOGLE_API_KEY", "")
|
||||
google_llm_model = os.environ.get("GOOGLE_LLM_MODEL", "gemini-2.0-flash")
|
||||
google_embedding_model = os.environ.get("GOOGLE_EMBEDDING_MODEL", "text-embedding-004")
|
||||
|
||||
# Ollama settings
|
||||
ollama_base_url = os.environ.get("OLLAMA_BASE_URL", DEFAULT_OLLAMA_BASE_URL)
|
||||
ollama_llm_model = os.environ.get("OLLAMA_LLM_MODEL", "")
|
||||
@@ -220,6 +232,9 @@ class GraphitiConfig:
|
||||
azure_openai_embedding_deployment=azure_openai_embedding_deployment,
|
||||
voyage_api_key=voyage_api_key,
|
||||
voyage_embedding_model=voyage_embedding_model,
|
||||
google_api_key=google_api_key,
|
||||
google_llm_model=google_llm_model,
|
||||
google_embedding_model=google_embedding_model,
|
||||
ollama_base_url=ollama_base_url,
|
||||
ollama_llm_model=ollama_llm_model,
|
||||
ollama_embedding_model=ollama_embedding_model,
|
||||
@@ -262,6 +277,8 @@ class GraphitiConfig:
|
||||
)
|
||||
elif self.llm_provider == "ollama":
|
||||
return bool(self.ollama_llm_model)
|
||||
elif self.llm_provider == "google":
|
||||
return bool(self.google_api_key)
|
||||
return False
|
||||
|
||||
def _validate_embedder_provider(self) -> bool:
|
||||
@@ -278,6 +295,8 @@ class GraphitiConfig:
|
||||
)
|
||||
elif self.embedder_provider == "ollama":
|
||||
return bool(self.ollama_embedding_model and self.ollama_embedding_dim)
|
||||
elif self.embedder_provider == "google":
|
||||
return bool(self.google_api_key)
|
||||
return False
|
||||
|
||||
def get_validation_errors(self) -> list[str]:
|
||||
@@ -309,6 +328,9 @@ class GraphitiConfig:
|
||||
elif self.llm_provider == "ollama":
|
||||
if not self.ollama_llm_model:
|
||||
errors.append("Ollama LLM provider requires OLLAMA_LLM_MODEL")
|
||||
elif self.llm_provider == "google":
|
||||
if not self.google_api_key:
|
||||
errors.append("Google LLM provider requires GOOGLE_API_KEY")
|
||||
else:
|
||||
errors.append(f"Unknown LLM provider: {self.llm_provider}")
|
||||
|
||||
@@ -339,6 +361,9 @@ class GraphitiConfig:
|
||||
)
|
||||
if not self.ollama_embedding_dim:
|
||||
errors.append("Ollama embedder provider requires OLLAMA_EMBEDDING_DIM")
|
||||
elif self.embedder_provider == "google":
|
||||
if not self.google_api_key:
|
||||
errors.append("Google embedder provider requires GOOGLE_API_KEY")
|
||||
else:
|
||||
errors.append(f"Unknown embedder provider: {self.embedder_provider}")
|
||||
|
||||
|
||||
@@ -11,6 +11,7 @@ if TYPE_CHECKING:
|
||||
from graphiti_config import GraphitiConfig
|
||||
|
||||
from .azure_openai_embedder import create_azure_openai_embedder
|
||||
from .google_embedder import create_google_embedder
|
||||
from .ollama_embedder import create_ollama_embedder
|
||||
from .openai_embedder import create_openai_embedder
|
||||
from .voyage_embedder import create_voyage_embedder
|
||||
@@ -20,4 +21,5 @@ __all__ = [
|
||||
"create_voyage_embedder",
|
||||
"create_azure_openai_embedder",
|
||||
"create_ollama_embedder",
|
||||
"create_google_embedder",
|
||||
]
|
||||
|
||||
+152
@@ -0,0 +1,152 @@
|
||||
"""
|
||||
Google AI Embedder Provider
|
||||
===========================
|
||||
|
||||
Google Gemini embedder implementation for Graphiti.
|
||||
Uses the google-generativeai SDK for text embeddings.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from graphiti_config import GraphitiConfig
|
||||
|
||||
from ..exceptions import ProviderError, ProviderNotInstalled
|
||||
|
||||
|
||||
# Default embedding model for Google
|
||||
DEFAULT_GOOGLE_EMBEDDING_MODEL = "text-embedding-004"
|
||||
|
||||
|
||||
class GoogleEmbedder:
|
||||
"""
|
||||
Google AI Embedder using the Gemini API.
|
||||
|
||||
Implements the EmbedderClient interface expected by graphiti-core.
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: str, model: str = DEFAULT_GOOGLE_EMBEDDING_MODEL):
|
||||
"""
|
||||
Initialize the Google embedder.
|
||||
|
||||
Args:
|
||||
api_key: Google AI API key
|
||||
model: Embedding model name (default: text-embedding-004)
|
||||
"""
|
||||
try:
|
||||
import google.generativeai as genai
|
||||
except ImportError as e:
|
||||
raise ProviderNotInstalled(
|
||||
f"Google embedder requires google-generativeai. "
|
||||
f"Install with: pip install google-generativeai\n"
|
||||
f"Error: {e}"
|
||||
)
|
||||
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
|
||||
# Configure the Google AI client
|
||||
genai.configure(api_key=api_key)
|
||||
self._genai = genai
|
||||
|
||||
async def create(self, input_data: str | list[str]) -> list[float]:
|
||||
"""
|
||||
Create embeddings for the input data.
|
||||
|
||||
Args:
|
||||
input_data: Text string or list of strings to embed
|
||||
|
||||
Returns:
|
||||
List of floats representing the embedding vector
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# Handle single string input
|
||||
if isinstance(input_data, str):
|
||||
text = input_data
|
||||
elif isinstance(input_data, list) and len(input_data) > 0:
|
||||
# Join list items if it's a list of strings
|
||||
if isinstance(input_data[0], str):
|
||||
text = " ".join(input_data)
|
||||
else:
|
||||
# It might be token IDs, convert to string
|
||||
text = str(input_data)
|
||||
else:
|
||||
text = str(input_data)
|
||||
|
||||
# Run the synchronous API call in a thread pool
|
||||
loop = asyncio.get_event_loop()
|
||||
result = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self._genai.embed_content(
|
||||
model=f"models/{self.model}",
|
||||
content=text,
|
||||
task_type="retrieval_document"
|
||||
)
|
||||
)
|
||||
|
||||
return result['embedding']
|
||||
|
||||
async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
|
||||
"""
|
||||
Create embeddings for a batch of inputs.
|
||||
|
||||
Args:
|
||||
input_data_list: List of text strings to embed
|
||||
|
||||
Returns:
|
||||
List of embedding vectors
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# Google's API supports batch embedding
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
# Process in batches to avoid rate limits
|
||||
batch_size = 100
|
||||
all_embeddings = []
|
||||
|
||||
for i in range(0, len(input_data_list), batch_size):
|
||||
batch = input_data_list[i:i + batch_size]
|
||||
|
||||
result = await loop.run_in_executor(
|
||||
None,
|
||||
lambda b=batch: self._genai.embed_content(
|
||||
model=f"models/{self.model}",
|
||||
content=b,
|
||||
task_type="retrieval_document"
|
||||
)
|
||||
)
|
||||
|
||||
# Handle single vs batch response
|
||||
if isinstance(result['embedding'][0], list):
|
||||
all_embeddings.extend(result['embedding'])
|
||||
else:
|
||||
all_embeddings.append(result['embedding'])
|
||||
|
||||
return all_embeddings
|
||||
|
||||
|
||||
def create_google_embedder(config: "GraphitiConfig") -> Any:
|
||||
"""
|
||||
Create Google AI embedder.
|
||||
|
||||
Args:
|
||||
config: GraphitiConfig with Google settings
|
||||
|
||||
Returns:
|
||||
Google embedder instance
|
||||
|
||||
Raises:
|
||||
ProviderNotInstalled: If google-generativeai is not installed
|
||||
ProviderError: If API key is missing
|
||||
"""
|
||||
if not config.google_api_key:
|
||||
raise ProviderError("Google embedder requires GOOGLE_API_KEY")
|
||||
|
||||
model = config.google_embedding_model or DEFAULT_GOOGLE_EMBEDDING_MODEL
|
||||
|
||||
return GoogleEmbedder(
|
||||
api_key=config.google_api_key,
|
||||
model=model
|
||||
)
|
||||
@@ -13,6 +13,7 @@ if TYPE_CHECKING:
|
||||
|
||||
from .embedder_providers import (
|
||||
create_azure_openai_embedder,
|
||||
create_google_embedder,
|
||||
create_ollama_embedder,
|
||||
create_openai_embedder,
|
||||
create_voyage_embedder,
|
||||
@@ -21,6 +22,7 @@ from .exceptions import ProviderError
|
||||
from .llm_providers import (
|
||||
create_anthropic_llm_client,
|
||||
create_azure_openai_llm_client,
|
||||
create_google_llm_client,
|
||||
create_ollama_llm_client,
|
||||
create_openai_llm_client,
|
||||
)
|
||||
@@ -54,6 +56,8 @@ def create_llm_client(config: "GraphitiConfig") -> Any:
|
||||
return create_azure_openai_llm_client(config)
|
||||
elif provider == "ollama":
|
||||
return create_ollama_llm_client(config)
|
||||
elif provider == "google":
|
||||
return create_google_llm_client(config)
|
||||
else:
|
||||
raise ProviderError(f"Unknown LLM provider: {provider}")
|
||||
|
||||
@@ -84,5 +88,7 @@ def create_embedder(config: "GraphitiConfig") -> Any:
|
||||
return create_azure_openai_embedder(config)
|
||||
elif provider == "ollama":
|
||||
return create_ollama_embedder(config)
|
||||
elif provider == "google":
|
||||
return create_google_embedder(config)
|
||||
else:
|
||||
raise ProviderError(f"Unknown embedder provider: {provider}")
|
||||
|
||||
@@ -12,6 +12,7 @@ if TYPE_CHECKING:
|
||||
|
||||
from .anthropic_llm import create_anthropic_llm_client
|
||||
from .azure_openai_llm import create_azure_openai_llm_client
|
||||
from .google_llm import create_google_llm_client
|
||||
from .ollama_llm import create_ollama_llm_client
|
||||
from .openai_llm import create_openai_llm_client
|
||||
|
||||
@@ -20,4 +21,5 @@ __all__ = [
|
||||
"create_anthropic_llm_client",
|
||||
"create_azure_openai_llm_client",
|
||||
"create_ollama_llm_client",
|
||||
"create_google_llm_client",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,175 @@
|
||||
"""
|
||||
Google AI LLM Provider
|
||||
======================
|
||||
|
||||
Google Gemini LLM client implementation for Graphiti.
|
||||
Uses the google-generativeai SDK.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from graphiti_config import GraphitiConfig
|
||||
|
||||
from ..exceptions import ProviderError, ProviderNotInstalled
|
||||
|
||||
|
||||
# Default model for Google LLM
|
||||
DEFAULT_GOOGLE_LLM_MODEL = "gemini-2.0-flash"
|
||||
|
||||
|
||||
class GoogleLLMClient:
|
||||
"""
|
||||
Google AI LLM Client using the Gemini API.
|
||||
|
||||
Implements the LLMClient interface expected by graphiti-core.
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: str, model: str = DEFAULT_GOOGLE_LLM_MODEL):
|
||||
"""
|
||||
Initialize the Google LLM client.
|
||||
|
||||
Args:
|
||||
api_key: Google AI API key
|
||||
model: Model name (default: gemini-2.0-flash)
|
||||
"""
|
||||
try:
|
||||
import google.generativeai as genai
|
||||
except ImportError as e:
|
||||
raise ProviderNotInstalled(
|
||||
f"Google LLM requires google-generativeai. "
|
||||
f"Install with: pip install google-generativeai\n"
|
||||
f"Error: {e}"
|
||||
)
|
||||
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
|
||||
# Configure the Google AI client
|
||||
genai.configure(api_key=api_key)
|
||||
self._genai = genai
|
||||
self._model = genai.GenerativeModel(model)
|
||||
|
||||
async def generate_response(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
response_model: Any = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""
|
||||
Generate a response from the LLM.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts with 'role' and 'content'
|
||||
response_model: Optional Pydantic model for structured output
|
||||
**kwargs: Additional arguments
|
||||
|
||||
Returns:
|
||||
Generated response (string or structured object)
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# Convert messages to Google format
|
||||
# Google uses 'user' and 'model' roles
|
||||
google_messages = []
|
||||
system_instruction = None
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
# Google handles system messages as system_instruction
|
||||
system_instruction = content
|
||||
elif role == "assistant":
|
||||
google_messages.append({"role": "model", "parts": [content]})
|
||||
else:
|
||||
google_messages.append({"role": "user", "parts": [content]})
|
||||
|
||||
# Create model with system instruction if provided
|
||||
if system_instruction:
|
||||
model = self._genai.GenerativeModel(
|
||||
self.model,
|
||||
system_instruction=system_instruction
|
||||
)
|
||||
else:
|
||||
model = self._model
|
||||
|
||||
# Generate response
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
if response_model:
|
||||
# For structured output, use JSON mode
|
||||
generation_config = self._genai.GenerationConfig(
|
||||
response_mime_type="application/json"
|
||||
)
|
||||
|
||||
response = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: model.generate_content(
|
||||
google_messages,
|
||||
generation_config=generation_config
|
||||
)
|
||||
)
|
||||
|
||||
# Parse JSON response into the model
|
||||
import json
|
||||
try:
|
||||
data = json.loads(response.text)
|
||||
return response_model(**data)
|
||||
except (json.JSONDecodeError, Exception):
|
||||
# If parsing fails, return raw text
|
||||
return response.text
|
||||
else:
|
||||
response = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: model.generate_content(google_messages)
|
||||
)
|
||||
|
||||
return response.text
|
||||
|
||||
async def generate_response_with_tools(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[Any],
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""
|
||||
Generate a response with tool calling support.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts
|
||||
tools: List of tool definitions
|
||||
**kwargs: Additional arguments
|
||||
|
||||
Returns:
|
||||
Generated response with potential tool calls
|
||||
"""
|
||||
# For now, fall back to regular generation
|
||||
# Tool calling can be added later if needed
|
||||
return await self.generate_response(messages, **kwargs)
|
||||
|
||||
|
||||
def create_google_llm_client(config: "GraphitiConfig") -> Any:
|
||||
"""
|
||||
Create Google AI LLM client.
|
||||
|
||||
Args:
|
||||
config: GraphitiConfig with Google settings
|
||||
|
||||
Returns:
|
||||
Google LLM client instance
|
||||
|
||||
Raises:
|
||||
ProviderNotInstalled: If google-generativeai is not installed
|
||||
ProviderError: If API key is missing
|
||||
"""
|
||||
if not config.google_api_key:
|
||||
raise ProviderError("Google LLM provider requires GOOGLE_API_KEY")
|
||||
|
||||
model = config.google_llm_model or DEFAULT_GOOGLE_LLM_MODEL
|
||||
|
||||
return GoogleLLMClient(
|
||||
api_key=config.google_api_key,
|
||||
model=model
|
||||
)
|
||||
@@ -4,3 +4,6 @@ python-dotenv>=1.0.0
|
||||
|
||||
# Memory Integration (highly recommended) but can be disabled by commenting out the line below
|
||||
graphiti-core[falkordb]>=0.5.0
|
||||
|
||||
# Google AI embeddings (optional - for Gemini embeddings)
|
||||
google-generativeai>=0.8.0
|
||||
|
||||
Reference in New Issue
Block a user