From 03b3cf2b4e02d20b7fd2fa50ed1de429d6900591 Mon Sep 17 00:00:00 2001 From: kavenGw Date: Sun, 20 Jul 2025 10:05:43 +0800 Subject: [PATCH 1/2] support ollama --- graphiti_core/embedder/ollama.py | 131 ++++++++++++++++++++++ mcp_server/.env.example.gemini_ollama | 44 ++++++++ mcp_server/.env.example.openrouter_ollama | 43 +++++++ mcp_server/graphiti_mcp_server.py | 36 ++++++ 4 files changed, 254 insertions(+) create mode 100644 graphiti_core/embedder/ollama.py create mode 100644 mcp_server/.env.example.gemini_ollama create mode 100644 mcp_server/.env.example.openrouter_ollama diff --git a/graphiti_core/embedder/ollama.py b/graphiti_core/embedder/ollama.py new file mode 100644 index 00000000..83a67902 --- /dev/null +++ b/graphiti_core/embedder/ollama.py @@ -0,0 +1,131 @@ +""" +Copyright 2024, Zep Software, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import logging +from collections.abc import Iterable +from typing import Any + +import httpx +from pydantic import Field + +from .client import EmbedderClient, EmbedderConfig + +logger = logging.getLogger(__name__) + +DEFAULT_EMBEDDING_MODEL = 'nomic-embed-text' +DEFAULT_BASE_URL = 'http://localhost:11434' + + +class OllamaEmbedderConfig(EmbedderConfig): + embedding_model: str = Field(default=DEFAULT_EMBEDDING_MODEL) + base_url: str = Field(default=DEFAULT_BASE_URL) + + +class OllamaEmbedder(EmbedderClient): + """ + Ollama Embedder Client + + Uses Ollama's native API endpoint for embeddings. + """ + + def __init__(self, config: OllamaEmbedderConfig | None = None): + if config is None: + config = OllamaEmbedderConfig() + self.config = config + self.base_url = config.base_url.rstrip('/') + self.embed_url = f"{self.base_url}/api/embed" + + async def create( + self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]] + ) -> list[float]: + """ + Create embeddings for the given input data using Ollama's embedding model. + + Args: + input_data: The input data to create embeddings for. Can be a string, list of strings, + or an iterable of integers or iterables of integers. + + Returns: + A list of floats representing the embedding vector. + """ + # Convert input to string if needed + if isinstance(input_data, str): + text_input = input_data + elif isinstance(input_data, list) and len(input_data) > 0: + if isinstance(input_data[0], str): + # For list of strings, take the first one for single embedding + text_input = input_data[0] + else: + # Convert other types to string + text_input = str(input_data[0]) + else: + text_input = str(input_data) + + payload = { + "model": self.config.embedding_model, + "input": text_input + } + + try: + async with httpx.AsyncClient() as client: + response = await client.post( + self.embed_url, + json=payload, + headers={"Content-Type": "application/json"}, + timeout=30.0 + ) + + if response.status_code != 200: + error_text = response.text + raise Exception(f"Ollama API error {response.status_code}: {error_text}") + + result = response.json() + + if "embeddings" not in result: + raise Exception(f"No embeddings in response: {result}") + + embeddings = result["embeddings"] + if not embeddings or len(embeddings) == 0: + raise Exception("Empty embeddings returned") + + # Return the first embedding, truncated to the configured dimension + embedding = embeddings[0] + return embedding[: self.config.embedding_dim] + + except httpx.HTTPStatusError as e: + logger.error(f"HTTP error creating Ollama embedding: {e.response.status_code} - {e.response.text}") + raise Exception(f"Ollama API error {e.response.status_code}: {e.response.text}") + except Exception as e: + logger.error(f"Error creating Ollama embedding: {e}") + raise + + async def create_batch(self, input_data_list: list[str]) -> list[list[float]]: + """ + Create batch embeddings using Ollama's embedding model. + + Note: Ollama doesn't support batch embeddings natively, so we process them sequentially. + """ + embeddings = [] + + for text in input_data_list: + try: + embedding = await self.create(text) + embeddings.append(embedding) + except Exception as e: + logger.error(f"Error creating embedding for text '{text[:50]}...': {e}") + raise + + return embeddings diff --git a/mcp_server/.env.example.gemini_ollama b/mcp_server/.env.example.gemini_ollama new file mode 100644 index 00000000..70c13935 --- /dev/null +++ b/mcp_server/.env.example.gemini_ollama @@ -0,0 +1,44 @@ +# Graphiti MCP Server Environment Configuration + +# Neo4j Database Configuration +# These settings are used to connect to your Neo4j database +NEO4J_URI=bolt://localhost:7687 +NEO4J_USER=neo4j +NEO4J_PASSWORD=demodemo + +# OpenAI API Configuration +# Required for LLM operations +OPENAI_API_KEY=your_gemini_api_key_here +MODEL_NAME=gemini-2.5-flash +SMALL_MODEL_NAME=gemini-2.5-flash + +# Optional: Only needed for non-standard OpenAI endpoints +OPENAI_BASE_URL=https://generativelanguage.googleapis.com/v1beta + +# Embedder Configuration +# 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 +EMBEDDER_API_KEY=ollama +EMBEDDER_BASE_URL=http://localhost:11434 +EMBEDDER_MODEL_NAME=nomic-embed-text +EMBEDDER_DIMENSION=768 + +# Optional: Group ID for namespacing graph data +# GROUP_ID=my_project + +# Optional: Path configuration for Docker +# PATH=/root/.local/bin:${PATH} + +# Optional: Memory settings for Neo4j (used in Docker Compose) +# NEO4J_server_memory_heap_initial__size=512m +# NEO4J_server_memory_heap_max__size=1G +# NEO4J_server_memory_pagecache_size=512m + +# Azure OpenAI configuration +# Optional: Only needed for Azure OpenAI endpoints +# AZURE_OPENAI_ENDPOINT=your_azure_openai_endpoint_here +# AZURE_OPENAI_API_VERSION=2025-01-01-preview +# AZURE_OPENAI_DEPLOYMENT_NAME=gpt-4o-gpt-4o-mini-deployment +# AZURE_OPENAI_EMBEDDING_API_VERSION=2023-05-15 +# AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME=text-embedding-3-large-deployment +# AZURE_OPENAI_USE_MANAGED_IDENTITY=false diff --git a/mcp_server/.env.example.openrouter_ollama b/mcp_server/.env.example.openrouter_ollama new file mode 100644 index 00000000..14181c3b --- /dev/null +++ b/mcp_server/.env.example.openrouter_ollama @@ -0,0 +1,43 @@ +# Graphiti MCP Server Environment Configuration + +# Neo4j Database Configuration +# These settings are used to connect to your Neo4j database +NEO4J_URI=bolt://localhost:7687 +NEO4J_USER=neo4j +NEO4J_PASSWORD=demodemo + +# OpenAI API Configuration +# Required for LLM operations +OPENAI_API_KEY=your_open_router_api_key_here +MODEL_NAME=gpt-4.1-mini + +# Optional: Only needed for non-standard OpenAI endpoints +OPENAI_BASE_URL=https://openrouter.ai/api/v1 + +# Embedder Configuration +# 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 +EMBEDDER_API_KEY=ollama +EMBEDDER_BASE_URL=http://localhost:11434 +EMBEDDER_MODEL_NAME=nomic-embed-text +EMBEDDER_DIMENSION=768 + +# Optional: Group ID for namespacing graph data +# GROUP_ID=my_project + +# Optional: Path configuration for Docker +# PATH=/root/.local/bin:${PATH} + +# Optional: Memory settings for Neo4j (used in Docker Compose) +# NEO4J_server_memory_heap_initial__size=512m +# NEO4J_server_memory_heap_max__size=1G +# NEO4J_server_memory_pagecache_size=512m + +# Azure OpenAI configuration +# Optional: Only needed for Azure OpenAI endpoints +# AZURE_OPENAI_ENDPOINT=your_azure_openai_endpoint_here +# AZURE_OPENAI_API_VERSION=2025-01-01-preview +# AZURE_OPENAI_DEPLOYMENT_NAME=gpt-4o-gpt-4o-mini-deployment +# AZURE_OPENAI_EMBEDDING_API_VERSION=2023-05-15 +# AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME=text-embedding-3-large-deployment +# AZURE_OPENAI_USE_MANAGED_IDENTITY=false diff --git a/mcp_server/graphiti_mcp_server.py b/mcp_server/graphiti_mcp_server.py index 9b382074..e8ef4da1 100644 --- a/mcp_server/graphiti_mcp_server.py +++ b/mcp_server/graphiti_mcp_server.py @@ -23,6 +23,7 @@ from graphiti_core.edges import EntityEdge from graphiti_core.embedder.azure_openai import AzureOpenAIEmbedderClient from graphiti_core.embedder.client import EmbedderClient from graphiti_core.embedder.openai import OpenAIEmbedder, OpenAIEmbedderConfig +from graphiti_core.embedder.ollama import OllamaEmbedder, OllamaEmbedderConfig from graphiti_core.llm_client import LLMClient from graphiti_core.llm_client.azure_openai_client import AzureOpenAILLMClient from graphiti_core.llm_client.config import LLMConfig @@ -354,6 +355,7 @@ class GraphitiEmbedderConfig(BaseModel): model: str = DEFAULT_EMBEDDER_MODEL api_key: str | None = None + provider: str = "openai" # "openai", "ollama", or "azure" azure_openai_endpoint: str | None = None azure_openai_deployment_name: str | None = None azure_openai_api_version: str | None = None @@ -367,6 +369,16 @@ class GraphitiEmbedderConfig(BaseModel): model_env = os.environ.get('EMBEDDER_MODEL_NAME', '') model = model_env if model_env.strip() else DEFAULT_EMBEDDER_MODEL + # 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 + provider = "openai" # default + if api_key and api_key.lower() == "ollama": + provider = "ollama" + + logger.info(f'GraphitiEmbedderConfig provider: {provider}') + 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_deployment_name = os.environ.get( @@ -408,6 +420,7 @@ class GraphitiEmbedderConfig(BaseModel): return cls( model=model, api_key=os.environ.get('OPENAI_API_KEY'), + provider=provider, ) def create_client(self) -> EmbedderClient | None: @@ -439,6 +452,29 @@ class GraphitiEmbedderConfig(BaseModel): else: logger.error('OPENAI_API_KEY must be set when using Azure OpenAI API') return None + elif self.provider == "ollama": + + base_url_env = os.environ.get('EMBEDDER_BASE_URL') + base_url = base_url_env if base_url_env else 'http://localhost:11434' + + model_env = os.environ.get('EMBEDDER_MODEL_NAME') + model = model_env if model_env else 'nomic-embed-text' + + # Get embedding dimension from environment + embedding_dim_env = os.environ.get('EMBEDDER_DIMENSION') + embedding_dim = int(embedding_dim_env) if embedding_dim_env else 768 + + logger.info(f'ollama model: {model}') + logger.info(f'ollama base_url: {base_url}') + logger.info(f'ollama embedding_dim: {embedding_dim}') + + # Ollama API setup + ollama_config = OllamaEmbedderConfig( + embedding_model=model, + base_url=base_url, + embedding_dim=embedding_dim # nomic-embed-text default + ) + return OllamaEmbedder(config=ollama_config) else: # OpenAI API setup if not self.api_key: From 0e63e57e6cf18d9adedf336616c3bf34162da6b6 Mon Sep 17 00:00:00 2001 From: kavenGw Date: Sun, 20 Jul 2025 11:26:16 +0800 Subject: [PATCH 2/2] support ollama --- graphiti_core/embedder/__init__.py | 3 +++ mcp_server/.env.example | 11 ++++++++ mcp_server/.env.example.gemini_ollama | 9 ++++--- mcp_server/.env.example.openrouter_ollama | 9 ++++--- mcp_server/README.md | 5 ++++ mcp_server/graphiti_mcp_server.py | 32 +++++++++++++---------- 6 files changed, 47 insertions(+), 22 deletions(-) diff --git a/graphiti_core/embedder/__init__.py b/graphiti_core/embedder/__init__.py index aea15619..417aa598 100644 --- a/graphiti_core/embedder/__init__.py +++ b/graphiti_core/embedder/__init__.py @@ -1,8 +1,11 @@ from .client import EmbedderClient from .openai import OpenAIEmbedder, OpenAIEmbedderConfig +from .ollama import OllamaEmbedder, OllamaEmbedderConfig __all__ = [ 'EmbedderClient', 'OpenAIEmbedder', 'OpenAIEmbedderConfig', + 'OllamaEmbedder', + 'OllamaEmbedderConfig', ] diff --git a/mcp_server/.env.example b/mcp_server/.env.example index 1e70ee56..795b5a48 100644 --- a/mcp_server/.env.example +++ b/mcp_server/.env.example @@ -14,6 +14,17 @@ MODEL_NAME=gpt-4.1-mini # Optional: Only needed for non-standard OpenAI endpoints # 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 # GROUP_ID=my_project diff --git a/mcp_server/.env.example.gemini_ollama b/mcp_server/.env.example.gemini_ollama index 70c13935..9b68c355 100644 --- a/mcp_server/.env.example.gemini_ollama +++ b/mcp_server/.env.example.gemini_ollama @@ -18,10 +18,11 @@ OPENAI_BASE_URL=https://generativelanguage.googleapis.com/v1beta # Embedder Configuration # 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 -EMBEDDER_API_KEY=ollama -EMBEDDER_BASE_URL=http://localhost:11434 -EMBEDDER_MODEL_NAME=nomic-embed-text -EMBEDDER_DIMENSION=768 +USE_OLLAMA_FOR_EMBEDDER=true +OLLAMA_EMBEDDER_API_KEY=ollama +OLLAMA_EMBEDDER_BASE_URL=http://localhost:11434 +OLLAMA_EMBEDDER_MODEL_NAME=nomic-embed-text +OLLAMA_EMBEDDER_DIMENSION=768 # Optional: Group ID for namespacing graph data # GROUP_ID=my_project diff --git a/mcp_server/.env.example.openrouter_ollama b/mcp_server/.env.example.openrouter_ollama index 14181c3b..0939f2dc 100644 --- a/mcp_server/.env.example.openrouter_ollama +++ b/mcp_server/.env.example.openrouter_ollama @@ -17,10 +17,11 @@ OPENAI_BASE_URL=https://openrouter.ai/api/v1 # Embedder Configuration # 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 -EMBEDDER_API_KEY=ollama -EMBEDDER_BASE_URL=http://localhost:11434 -EMBEDDER_MODEL_NAME=nomic-embed-text -EMBEDDER_DIMENSION=768 +USE_OLLAMA_FOR_EMBEDDER=true +OLLAMA_EMBEDDER_API_KEY=ollama +OLLAMA_EMBEDDER_BASE_URL=http://localhost:11434 +OLLAMA_EMBEDDER_MODEL_NAME=nomic-embed-text +OLLAMA_EMBEDDER_DIMENSION=768 # Optional: Group ID for namespacing graph data # GROUP_ID=my_project diff --git a/mcp_server/README.md b/mcp_server/README.md index d957feb8..ee52928b 100644 --- a/mcp_server/README.md +++ b/mcp_server/README.md @@ -92,6 +92,11 @@ The server uses the following environment variables: - `MODEL_NAME`: OpenAI model name to use for LLM operations. - `SMALL_MODEL_NAME`: OpenAI model name to use for smaller LLM operations. - `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_DEPLOYMENT_NAME`: Optional Azure OpenAI LLM deployment name - `AZURE_OPENAI_API_VERSION`: Optional Azure OpenAI LLM API version diff --git a/mcp_server/graphiti_mcp_server.py b/mcp_server/graphiti_mcp_server.py index e8ef4da1..b0901b6d 100644 --- a/mcp_server/graphiti_mcp_server.py +++ b/mcp_server/graphiti_mcp_server.py @@ -355,11 +355,11 @@ class GraphitiEmbedderConfig(BaseModel): model: str = DEFAULT_EMBEDDER_MODEL api_key: str | None = None - provider: str = "openai" # "openai", "ollama", or "azure" azure_openai_endpoint: str | None = None azure_openai_deployment_name: str | None = None azure_openai_api_version: str | None = None azure_openai_use_managed_identity: bool = False + use_ollama_for_embedder: bool = False @classmethod def from_env(cls) -> 'GraphitiEmbedderConfig': @@ -370,14 +370,11 @@ class GraphitiEmbedderConfig(BaseModel): model = model_env if model_env.strip() else DEFAULT_EMBEDDER_MODEL # 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 - provider = "openai" # default - if api_key and api_key.lower() == "ollama": - provider = "ollama" - - logger.info(f'GraphitiEmbedderConfig provider: {provider}') + # Detect provider based on configuration (similar to Azure pattern) + use_ollama_for_embedder = ( + os.environ.get('USE_OLLAMA_FOR_EMBEDDER', 'false').lower() == 'true' + ) azure_openai_endpoint = os.environ.get('AZURE_OPENAI_EMBEDDING_ENDPOINT', 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, azure_openai_api_version=azure_openai_api_version, azure_openai_deployment_name=azure_openai_deployment_name, + use_ollama_for_embedder=False, ) 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( model=model, - api_key=os.environ.get('OPENAI_API_KEY'), - provider=provider, + api_key=api_key, + use_ollama_for_embedder=use_ollama_for_embedder, ) def create_client(self) -> EmbedderClient | None: @@ -452,16 +456,16 @@ class GraphitiEmbedderConfig(BaseModel): else: logger.error('OPENAI_API_KEY must be set when using Azure OpenAI API') 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' - 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' # 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 logger.info(f'ollama model: {model}')