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:
@@ -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',
|
||||
}
|
||||
});
|
||||
}}
|
||||
|
||||
@@ -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")
|
||||
|
||||
+12
-15
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user