support ollama
This commit is contained in:
parent
03b3cf2b4e
commit
0e63e57e6c
6 changed files with 47 additions and 22 deletions
|
|
@ -1,8 +1,11 @@
|
||||||
from .client import EmbedderClient
|
from .client import EmbedderClient
|
||||||
from .openai import OpenAIEmbedder, OpenAIEmbedderConfig
|
from .openai import OpenAIEmbedder, OpenAIEmbedderConfig
|
||||||
|
from .ollama import OllamaEmbedder, OllamaEmbedderConfig
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
'EmbedderClient',
|
'EmbedderClient',
|
||||||
'OpenAIEmbedder',
|
'OpenAIEmbedder',
|
||||||
'OpenAIEmbedderConfig',
|
'OpenAIEmbedderConfig',
|
||||||
|
'OllamaEmbedder',
|
||||||
|
'OllamaEmbedderConfig',
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -14,6 +14,17 @@ MODEL_NAME=gpt-4.1-mini
|
||||||
# Optional: Only needed for non-standard OpenAI endpoints
|
# Optional: Only needed for non-standard OpenAI endpoints
|
||||||
# OPENAI_BASE_URL=https://api.openai.com/v1
|
# OPENAI_BASE_URL=https://api.openai.com/v1
|
||||||
|
|
||||||
|
# Embedder Configuration
|
||||||
|
# Provider is auto-detected based on configuration:
|
||||||
|
# - Azure: if AZURE_OPENAI_EMBEDDING_ENDPOINT is set
|
||||||
|
# - Ollama: if USE_OLLAMA_FOR_EMBEDDER is set to true
|
||||||
|
# - OpenAI: default (no additional config needed)
|
||||||
|
# USE_OLLAMA_FOR_EMBEDDER=true # Set this to true to use Ollama
|
||||||
|
# OLLAMA_EMBEDDER_API_KEY=ollama # Ollama API key (optional, defaults to 'ollama')
|
||||||
|
# OLLAMA_EMBEDDER_BASE_URL=http://localhost:11434 # Ollama base URL (when using Ollama)
|
||||||
|
# OLLAMA_EMBEDDER_MODEL_NAME=nomic-embed-text # Ollama embedding model to use
|
||||||
|
# OLLAMA_EMBEDDER_DIMENSION=768 # Ollama embedding dimension (model-specific)
|
||||||
|
|
||||||
# Optional: Group ID for namespacing graph data
|
# Optional: Group ID for namespacing graph data
|
||||||
# GROUP_ID=my_project
|
# GROUP_ID=my_project
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -18,10 +18,11 @@ OPENAI_BASE_URL=https://generativelanguage.googleapis.com/v1beta
|
||||||
# Embedder Configuration
|
# Embedder Configuration
|
||||||
# Optional: Separate API key and URL for embedder (falls back to OPENAI_API_KEY and OPENAI_BASE_URL if not set)
|
# Optional: Separate API key and URL for embedder (falls back to OPENAI_API_KEY and OPENAI_BASE_URL if not set)
|
||||||
# Note: OpenRouter does not support embeddings API, using Ollama as free alternative
|
# Note: OpenRouter does not support embeddings API, using Ollama as free alternative
|
||||||
EMBEDDER_API_KEY=ollama
|
USE_OLLAMA_FOR_EMBEDDER=true
|
||||||
EMBEDDER_BASE_URL=http://localhost:11434
|
OLLAMA_EMBEDDER_API_KEY=ollama
|
||||||
EMBEDDER_MODEL_NAME=nomic-embed-text
|
OLLAMA_EMBEDDER_BASE_URL=http://localhost:11434
|
||||||
EMBEDDER_DIMENSION=768
|
OLLAMA_EMBEDDER_MODEL_NAME=nomic-embed-text
|
||||||
|
OLLAMA_EMBEDDER_DIMENSION=768
|
||||||
|
|
||||||
# Optional: Group ID for namespacing graph data
|
# Optional: Group ID for namespacing graph data
|
||||||
# GROUP_ID=my_project
|
# GROUP_ID=my_project
|
||||||
|
|
|
||||||
|
|
@ -17,10 +17,11 @@ OPENAI_BASE_URL=https://openrouter.ai/api/v1
|
||||||
# Embedder Configuration
|
# Embedder Configuration
|
||||||
# Optional: Separate API key and URL for embedder (falls back to OPENAI_API_KEY and OPENAI_BASE_URL if not set)
|
# Optional: Separate API key and URL for embedder (falls back to OPENAI_API_KEY and OPENAI_BASE_URL if not set)
|
||||||
# Note: OpenRouter does not support embeddings API, using Ollama as free alternative
|
# Note: OpenRouter does not support embeddings API, using Ollama as free alternative
|
||||||
EMBEDDER_API_KEY=ollama
|
USE_OLLAMA_FOR_EMBEDDER=true
|
||||||
EMBEDDER_BASE_URL=http://localhost:11434
|
OLLAMA_EMBEDDER_API_KEY=ollama
|
||||||
EMBEDDER_MODEL_NAME=nomic-embed-text
|
OLLAMA_EMBEDDER_BASE_URL=http://localhost:11434
|
||||||
EMBEDDER_DIMENSION=768
|
OLLAMA_EMBEDDER_MODEL_NAME=nomic-embed-text
|
||||||
|
OLLAMA_EMBEDDER_DIMENSION=768
|
||||||
|
|
||||||
# Optional: Group ID for namespacing graph data
|
# Optional: Group ID for namespacing graph data
|
||||||
# GROUP_ID=my_project
|
# GROUP_ID=my_project
|
||||||
|
|
|
||||||
|
|
@ -92,6 +92,11 @@ The server uses the following environment variables:
|
||||||
- `MODEL_NAME`: OpenAI model name to use for LLM operations.
|
- `MODEL_NAME`: OpenAI model name to use for LLM operations.
|
||||||
- `SMALL_MODEL_NAME`: OpenAI model name to use for smaller LLM operations.
|
- `SMALL_MODEL_NAME`: OpenAI model name to use for smaller LLM operations.
|
||||||
- `LLM_TEMPERATURE`: Temperature for LLM responses (0.0-2.0).
|
- `LLM_TEMPERATURE`: Temperature for LLM responses (0.0-2.0).
|
||||||
|
- `USE_OLLAMA_FOR_EMBEDDER`: Set to `true` to use Ollama for embeddings (auto-detects Ollama provider)
|
||||||
|
- `OLLAMA_EMBEDDER_API_KEY`: Ollama API key (optional, defaults to 'ollama')
|
||||||
|
- `OLLAMA_EMBEDDER_BASE_URL`: Ollama base URL for embedder API (when using Ollama)
|
||||||
|
- `OLLAMA_EMBEDDER_MODEL_NAME`: Ollama embedding model name
|
||||||
|
- `OLLAMA_EMBEDDER_DIMENSION`: Ollama embedding dimension
|
||||||
- `AZURE_OPENAI_ENDPOINT`: Optional Azure OpenAI LLM endpoint URL
|
- `AZURE_OPENAI_ENDPOINT`: Optional Azure OpenAI LLM endpoint URL
|
||||||
- `AZURE_OPENAI_DEPLOYMENT_NAME`: Optional Azure OpenAI LLM deployment name
|
- `AZURE_OPENAI_DEPLOYMENT_NAME`: Optional Azure OpenAI LLM deployment name
|
||||||
- `AZURE_OPENAI_API_VERSION`: Optional Azure OpenAI LLM API version
|
- `AZURE_OPENAI_API_VERSION`: Optional Azure OpenAI LLM API version
|
||||||
|
|
|
||||||
|
|
@ -355,11 +355,11 @@ class GraphitiEmbedderConfig(BaseModel):
|
||||||
|
|
||||||
model: str = DEFAULT_EMBEDDER_MODEL
|
model: str = DEFAULT_EMBEDDER_MODEL
|
||||||
api_key: str | None = None
|
api_key: str | None = None
|
||||||
provider: str = "openai" # "openai", "ollama", or "azure"
|
|
||||||
azure_openai_endpoint: str | None = None
|
azure_openai_endpoint: str | None = None
|
||||||
azure_openai_deployment_name: str | None = None
|
azure_openai_deployment_name: str | None = None
|
||||||
azure_openai_api_version: str | None = None
|
azure_openai_api_version: str | None = None
|
||||||
azure_openai_use_managed_identity: bool = False
|
azure_openai_use_managed_identity: bool = False
|
||||||
|
use_ollama_for_embedder: bool = False
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_env(cls) -> 'GraphitiEmbedderConfig':
|
def from_env(cls) -> 'GraphitiEmbedderConfig':
|
||||||
|
|
@ -370,14 +370,11 @@ class GraphitiEmbedderConfig(BaseModel):
|
||||||
model = model_env if model_env.strip() else DEFAULT_EMBEDDER_MODEL
|
model = model_env if model_env.strip() else DEFAULT_EMBEDDER_MODEL
|
||||||
|
|
||||||
# Get embedder-specific API key and base URL, fallback to general OpenAI settings
|
# Get embedder-specific API key and base URL, fallback to general OpenAI settings
|
||||||
api_key = os.environ.get('EMBEDDER_API_KEY') or os.environ.get('OPENAI_API_KEY')
|
|
||||||
|
|
||||||
# Detect provider based on configuration
|
# Detect provider based on configuration (similar to Azure pattern)
|
||||||
provider = "openai" # default
|
use_ollama_for_embedder = (
|
||||||
if api_key and api_key.lower() == "ollama":
|
os.environ.get('USE_OLLAMA_FOR_EMBEDDER', 'false').lower() == 'true'
|
||||||
provider = "ollama"
|
)
|
||||||
|
|
||||||
logger.info(f'GraphitiEmbedderConfig provider: {provider}')
|
|
||||||
|
|
||||||
azure_openai_endpoint = os.environ.get('AZURE_OPENAI_EMBEDDING_ENDPOINT', None)
|
azure_openai_endpoint = os.environ.get('AZURE_OPENAI_EMBEDDING_ENDPOINT', None)
|
||||||
azure_openai_api_version = os.environ.get('AZURE_OPENAI_EMBEDDING_API_VERSION', None)
|
azure_openai_api_version = os.environ.get('AZURE_OPENAI_EMBEDDING_API_VERSION', None)
|
||||||
|
|
@ -415,12 +412,19 @@ class GraphitiEmbedderConfig(BaseModel):
|
||||||
api_key=api_key,
|
api_key=api_key,
|
||||||
azure_openai_api_version=azure_openai_api_version,
|
azure_openai_api_version=azure_openai_api_version,
|
||||||
azure_openai_deployment_name=azure_openai_deployment_name,
|
azure_openai_deployment_name=azure_openai_deployment_name,
|
||||||
|
use_ollama_for_embedder=False,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
|
if use_ollama_for_embedder:
|
||||||
|
api_key_env = os.environ.get("OLLAMA_EMBEDDER_API_KEY")
|
||||||
|
api_key = api_key_env if api_key_env else 'ollama'
|
||||||
|
logger.info(f'ollama api_key: {api_key}')
|
||||||
|
else:
|
||||||
|
api_key = os.environ.get("OPENAI_API_KEY")
|
||||||
return cls(
|
return cls(
|
||||||
model=model,
|
model=model,
|
||||||
api_key=os.environ.get('OPENAI_API_KEY'),
|
api_key=api_key,
|
||||||
provider=provider,
|
use_ollama_for_embedder=use_ollama_for_embedder,
|
||||||
)
|
)
|
||||||
|
|
||||||
def create_client(self) -> EmbedderClient | None:
|
def create_client(self) -> EmbedderClient | None:
|
||||||
|
|
@ -452,16 +456,16 @@ class GraphitiEmbedderConfig(BaseModel):
|
||||||
else:
|
else:
|
||||||
logger.error('OPENAI_API_KEY must be set when using Azure OpenAI API')
|
logger.error('OPENAI_API_KEY must be set when using Azure OpenAI API')
|
||||||
return None
|
return None
|
||||||
elif self.provider == "ollama":
|
elif self.use_ollama_for_embedder:
|
||||||
|
|
||||||
base_url_env = os.environ.get('EMBEDDER_BASE_URL')
|
base_url_env = os.environ.get('OLLAMA_EMBEDDER_BASE_URL')
|
||||||
base_url = base_url_env if base_url_env else 'http://localhost:11434'
|
base_url = base_url_env if base_url_env else 'http://localhost:11434'
|
||||||
|
|
||||||
model_env = os.environ.get('EMBEDDER_MODEL_NAME')
|
model_env = os.environ.get('OLLAMA_EMBEDDER_MODEL_NAME')
|
||||||
model = model_env if model_env else 'nomic-embed-text'
|
model = model_env if model_env else 'nomic-embed-text'
|
||||||
|
|
||||||
# Get embedding dimension from environment
|
# Get embedding dimension from environment
|
||||||
embedding_dim_env = os.environ.get('EMBEDDER_DIMENSION')
|
embedding_dim_env = os.environ.get('OLLAMA_EMBEDDER_DIMENSION')
|
||||||
embedding_dim = int(embedding_dim_env) if embedding_dim_env else 768
|
embedding_dim = int(embedding_dim_env) if embedding_dim_env else 768
|
||||||
|
|
||||||
logger.info(f'ollama model: {model}')
|
logger.info(f'ollama model: {model}')
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue