fix(graphiti): address CodeRabbit review comments

- Apply ruff formatting to Python files
- Fix ENV_GET handler to populate graphitiProviderConfig from .env
- Add Google AI to get_available_providers() function
- Fix type assertions to include google/groq/huggingface providers
- Fix asyncio deprecation: use get_running_loop() instead of get_event_loop()
This commit is contained in:
adryserage
2025-12-19 07:56:40 -05:00
parent 1a38a06e6e
commit 679b8cd948
6 changed files with 70 additions and 32 deletions
@@ -317,6 +317,45 @@ ${existingVars['GRAPHITI_DATABASE'] ? `GRAPHITI_DATABASE=${existingVars['GRAPHIT
config.enableFancyUi = false;
}
// Populate graphitiProviderConfig from .env file
const llmProvider = vars['GRAPHITI_LLM_PROVIDER'];
const embeddingProvider = vars['GRAPHITI_EMBEDDER_PROVIDER'];
if (llmProvider || embeddingProvider || vars['ANTHROPIC_API_KEY'] || vars['AZURE_OPENAI_API_KEY'] ||
vars['VOYAGE_API_KEY'] || vars['GOOGLE_API_KEY'] || vars['OLLAMA_BASE_URL']) {
config.graphitiProviderConfig = {
llmProvider: (llmProvider as 'openai' | 'anthropic' | 'azure_openai' | 'ollama' | 'google' | 'groq') || 'openai',
embeddingProvider: (embeddingProvider as 'openai' | 'voyage' | 'azure_openai' | 'ollama' | 'google' | 'huggingface') || 'openai',
// OpenAI
openaiApiKey: vars['OPENAI_API_KEY'],
openaiModel: vars['OPENAI_MODEL'],
openaiEmbeddingModel: vars['OPENAI_EMBEDDING_MODEL'],
// Anthropic
anthropicApiKey: vars['ANTHROPIC_API_KEY'],
anthropicModel: vars['GRAPHITI_ANTHROPIC_MODEL'],
// Azure OpenAI
azureOpenaiApiKey: vars['AZURE_OPENAI_API_KEY'],
azureOpenaiBaseUrl: vars['AZURE_OPENAI_BASE_URL'],
azureOpenaiLlmDeployment: vars['AZURE_OPENAI_LLM_DEPLOYMENT'],
azureOpenaiEmbeddingDeployment: vars['AZURE_OPENAI_EMBEDDING_DEPLOYMENT'],
// Voyage
voyageApiKey: vars['VOYAGE_API_KEY'],
voyageEmbeddingModel: vars['VOYAGE_EMBEDDING_MODEL'],
// Google
googleApiKey: vars['GOOGLE_API_KEY'],
googleLlmModel: vars['GOOGLE_LLM_MODEL'],
googleEmbeddingModel: vars['GOOGLE_EMBEDDING_MODEL'],
// Ollama
ollamaBaseUrl: vars['OLLAMA_BASE_URL'],
ollamaLlmModel: vars['OLLAMA_LLM_MODEL'],
ollamaEmbeddingModel: vars['OLLAMA_EMBEDDING_MODEL'],
ollamaEmbeddingDim: vars['OLLAMA_EMBEDDING_DIM'] ? parseInt(vars['OLLAMA_EMBEDDING_DIM'], 10) : undefined,
// FalkorDB
falkorDbHost: vars['GRAPHITI_FALKORDB_HOST'],
falkorDbPort: vars['GRAPHITI_FALKORDB_PORT'] ? parseInt(vars['GRAPHITI_FALKORDB_PORT'], 10) : undefined,
falkorDbPassword: vars['GRAPHITI_FALKORDB_PASSWORD'],
};
}
return { success: true, data: config };
}
);
@@ -146,7 +146,7 @@ export function MemoryBackendSection({
onValueChange={(value) => onUpdateConfig({
graphitiProviderConfig: {
...envConfig.graphitiProviderConfig,
llmProvider: value as 'openai' | 'anthropic' | 'azure_openai' | 'ollama',
llmProvider: value as 'openai' | 'anthropic' | 'azure_openai' | 'ollama' | 'google' | 'groq',
embeddingProvider: envConfig.graphitiProviderConfig?.embeddingProvider || 'openai',
}
})}
@@ -176,7 +176,7 @@ export function MemoryBackendSection({
graphitiProviderConfig: {
...envConfig.graphitiProviderConfig,
llmProvider: envConfig.graphitiProviderConfig?.llmProvider || 'openai',
embeddingProvider: value as 'openai' | 'voyage' | 'azure_openai' | 'ollama',
embeddingProvider: value as 'openai' | 'voyage' | 'azure_openai' | 'ollama' | 'google' | 'huggingface',
}
})}
>
@@ -164,7 +164,7 @@ export function SecuritySettings({
updateEnvConfig({
graphitiProviderConfig: {
...currentConfig,
llmProvider: value as 'openai' | 'anthropic' | 'azure_openai' | 'ollama',
llmProvider: value as 'openai' | 'anthropic' | 'azure_openai' | 'ollama' | 'google' | 'groq',
}
});
}}
@@ -198,7 +198,7 @@ export function SecuritySettings({
updateEnvConfig({
graphitiProviderConfig: {
...currentConfig,
embeddingProvider: value as 'openai' | 'voyage' | 'azure_openai' | 'ollama',
embeddingProvider: value as 'openai' | 'voyage' | 'azure_openai' | 'ollama' | 'google' | 'huggingface',
}
});
}}
+8 -1
View File
@@ -199,7 +199,9 @@ class GraphitiConfig:
# 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")
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)
@@ -542,6 +544,11 @@ def get_available_providers() -> dict:
if config.voyage_api_key:
available_embedder.append("voyage")
# Check Google AI
if config.google_api_key:
available_llm.append("google")
available_embedder.append("google")
# Check Ollama
if config.ollama_llm_model:
available_llm.append("ollama")
@@ -75,17 +75,17 @@ class GoogleEmbedder:
text = str(input_data)
# Run the synchronous API call in a thread pool
loop = asyncio.get_event_loop()
loop = asyncio.get_running_loop()
result = await loop.run_in_executor(
None,
lambda: self._genai.embed_content(
model=f"models/{self.model}",
content=text,
task_type="retrieval_document"
)
task_type="retrieval_document",
),
)
return result['embedding']
return result["embedding"]
async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
"""
@@ -100,29 +100,29 @@ class GoogleEmbedder:
import asyncio
# Google's API supports batch embedding
loop = asyncio.get_event_loop()
loop = asyncio.get_running_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]
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"
)
task_type="retrieval_document",
),
)
# Handle single vs batch response
if isinstance(result['embedding'][0], list):
all_embeddings.extend(result['embedding'])
if isinstance(result["embedding"][0], list):
all_embeddings.extend(result["embedding"])
else:
all_embeddings.append(result['embedding'])
all_embeddings.append(result["embedding"])
return all_embeddings
@@ -146,7 +146,4 @@ def create_google_embedder(config: "GraphitiConfig") -> Any:
model = config.google_embedding_model or DEFAULT_GOOGLE_EMBEDDING_MODEL
return GoogleEmbedder(
api_key=config.google_api_key,
model=model
)
return GoogleEmbedder(api_key=config.google_api_key, model=model)
@@ -89,14 +89,13 @@ class GoogleLLMClient:
# Create model with system instruction if provided
if system_instruction:
model = self._genai.GenerativeModel(
self.model,
system_instruction=system_instruction
self.model, system_instruction=system_instruction
)
else:
model = self._model
# Generate response
loop = asyncio.get_event_loop()
loop = asyncio.get_running_loop()
if response_model:
# For structured output, use JSON mode
@@ -107,13 +106,13 @@ class GoogleLLMClient:
response = await loop.run_in_executor(
None,
lambda: model.generate_content(
google_messages,
generation_config=generation_config
)
google_messages, generation_config=generation_config
),
)
# Parse JSON response into the model
import json
try:
data = json.loads(response.text)
return response_model(**data)
@@ -122,8 +121,7 @@ class GoogleLLMClient:
return response.text
else:
response = await loop.run_in_executor(
None,
lambda: model.generate_content(google_messages)
None, lambda: model.generate_content(google_messages)
)
return response.text
@@ -169,7 +167,4 @@ def create_google_llm_client(config: "GraphitiConfig") -> Any:
model = config.google_llm_model or DEFAULT_GOOGLE_LLM_MODEL
return GoogleLLMClient(
api_key=config.google_api_key,
model=model
)
return GoogleLLMClient(api_key=config.google_api_key, model=model)