From d979e9078f2878aec90a5f6aaddce727253845bf Mon Sep 17 00:00:00 2001 From: yangdx Date: Thu, 24 Jul 2025 12:11:35 +0800 Subject: [PATCH 1/4] feat: Integrate Jina embeddings API support - Implemented Jina embedding function - Add new EMBEDDING_BINDING type of jina for LightRAG Server - Add env var sample --- env.example | 9 +++- lightrag/api/lightrag_server.py | 11 +++- lightrag/llm/jina.py | 90 ++++++++++++++++++++++++++++----- 3 files changed, 96 insertions(+), 14 deletions(-) diff --git a/env.example b/env.example index 42751abc..d1d8fee4 100644 --- a/env.example +++ b/env.example @@ -123,7 +123,7 @@ LLM_BINDING_API_KEY=your_api_key #################################################################################### ### Embedding Configuration (Should not be changed after the first file processed) #################################################################################### -### Embedding Binding type: openai, ollama, lollms, azure_openai +### Embedding Binding type: openai, ollama, lollms, azure_openai, jina EMBEDDING_BINDING=ollama EMBEDDING_MODEL=bge-m3:latest EMBEDDING_DIM=1024 @@ -139,6 +139,13 @@ EMBEDDING_BINDING_HOST=http://localhost:11434 # AZURE_EMBEDDING_ENDPOINT=your_endpoint # AZURE_EMBEDDING_API_KEY=your_api_key +### Jina AI Embedding +EMBEDDING_BINDING=jina +EMBEDDING_BINDING_HOST=https://api.jina.ai/v1/embeddings +EMBEDDING_MODEL=jina-embeddings-v4 +EMBEDDING_DIM=2048 +EMBEDDING_BINDING_API_KEY=your_api_key + ############################ ### Data storage selection ############################ diff --git a/lightrag/api/lightrag_server.py b/lightrag/api/lightrag_server.py index 17bbaaed..5bb13813 100644 --- a/lightrag/api/lightrag_server.py +++ b/lightrag/api/lightrag_server.py @@ -89,7 +89,7 @@ def create_app(args): ]: raise Exception("llm binding not supported") - if args.embedding_binding not in ["lollms", "ollama", "openai", "azure_openai"]: + if args.embedding_binding not in ["lollms", "ollama", "openai", "azure_openai", "jina"]: raise Exception("embedding binding not supported") # Set default hosts if not provided @@ -213,6 +213,8 @@ def create_app(args): if args.llm_binding_host == "openai-ollama" or args.embedding_binding == "ollama": from lightrag.llm.openai import openai_complete_if_cache from lightrag.llm.ollama import ollama_embed + if args.embedding_binding == "jina": + from lightrag.llm.jina import jina_embed async def openai_alike_model_complete( prompt, @@ -284,6 +286,13 @@ def create_app(args): api_key=args.embedding_binding_api_key, ) if args.embedding_binding == "azure_openai" + else jina_embed( + texts, + dimensions=args.embedding_dim, + base_url=args.embedding_binding_host, + api_key=args.embedding_binding_api_key, + ) + if args.embedding_binding == "jina" else openai_embed( texts, model=args.embedding_model, diff --git a/lightrag/llm/jina.py b/lightrag/llm/jina.py index c3484b39..720877f5 100644 --- a/lightrag/llm/jina.py +++ b/lightrag/llm/jina.py @@ -2,45 +2,111 @@ import os import pipmaster as pm # Pipmaster for dynamic library install # install specific modules -if not pm.is_installed("lmdeploy"): - pm.install("lmdeploy") +if not pm.is_installed("aiohttp"): + pm.install("aiohttp") if not pm.is_installed("tenacity"): pm.install("tenacity") - import numpy as np import aiohttp +from tenacity import ( + retry, + stop_after_attempt, + wait_exponential, + retry_if_exception_type, +) +from lightrag.utils import wrap_embedding_func_with_attrs, logger async def fetch_data(url, headers, data): async with aiohttp.ClientSession() as session: async with session.post(url, headers=headers, json=data) as response: + if response.status != 200: + error_text = await response.text() + logger.error(f"Jina API error {response.status}: {error_text}") + raise aiohttp.ClientResponseError( + request_info=response.request_info, + history=response.history, + status=response.status, + message=f"Jina API error: {error_text}" + ) response_json = await response.json() data_list = response_json.get("data", []) return data_list +@wrap_embedding_func_with_attrs(embedding_dim=2048, max_token_size=8192) +@retry( + stop=stop_after_attempt(3), + wait=wait_exponential(multiplier=1, min=4, max=60), + retry=( + retry_if_exception_type(aiohttp.ClientError) + | retry_if_exception_type(aiohttp.ClientResponseError) + ), +) async def jina_embed( texts: list[str], - dimensions: int = 1024, + dimensions: int = 2048, late_chunking: bool = False, base_url: str = None, api_key: str = None, ) -> np.ndarray: + """Generate embeddings for a list of texts using Jina AI's API. + + Args: + texts: List of texts to embed. + dimensions: The embedding dimensions (default: 2048 for jina-embeddings-v4). + late_chunking: Whether to use late chunking. + base_url: Optional base URL for the Jina API. + api_key: Optional Jina API key. If None, uses the JINA_API_KEY environment variable. + + Returns: + A numpy array of embeddings, one per input text. + + Raises: + aiohttp.ClientError: If there is a connection error with the Jina API. + aiohttp.ClientResponseError: If the Jina API returns an error response. + """ if api_key: os.environ["JINA_API_KEY"] = api_key - url = "https://api.jina.ai/v1/embeddings" if not base_url else base_url + + if "JINA_API_KEY" not in os.environ: + raise ValueError("JINA_API_KEY environment variable is required") + + url = base_url or "https://api.jina.ai/v1/embeddings" headers = { "Content-Type": "application/json", "Authorization": f"Bearer {os.environ['JINA_API_KEY']}", } data = { - "model": "jina-embeddings-v3", - "normalized": True, - "embedding_type": "float", - "dimensions": f"{dimensions}", - "late_chunking": late_chunking, + "model": "jina-embeddings-v4", + "task": "text-matching", + "dimensions": dimensions, "input": texts, } - data_list = await fetch_data(url, headers, data) - return np.array([dp["embedding"] for dp in data_list]) + + # Only add optional parameters if they have non-default values + if late_chunking: + data["late_chunking"] = late_chunking + + logger.debug(f"Jina embedding request: {len(texts)} texts, dimensions: {dimensions}") + + try: + data_list = await fetch_data(url, headers, data) + + if not data_list: + logger.error("Jina API returned empty data list") + raise ValueError("Jina API returned empty data list") + + if len(data_list) != len(texts): + logger.error(f"Jina API returned {len(data_list)} embeddings for {len(texts)} texts") + raise ValueError(f"Jina API returned {len(data_list)} embeddings for {len(texts)} texts") + + embeddings = np.array([dp["embedding"] for dp in data_list]) + logger.debug(f"Jina embeddings generated: shape {embeddings.shape}") + + return embeddings + + except Exception as e: + logger.error(f"Jina embedding error: {e}") + raise From 2767212ba0dc12fa1d9dd91c227e910820f29079 Mon Sep 17 00:00:00 2001 From: yangdx Date: Thu, 24 Jul 2025 12:25:50 +0800 Subject: [PATCH 2/4] Fix linting --- env.example | 2 +- lightrag/api/lightrag_server.py | 8 +++++++- lightrag/llm/jina.py | 34 +++++++++++++++++++-------------- 3 files changed, 28 insertions(+), 16 deletions(-) diff --git a/env.example b/env.example index d1d8fee4..c2f3c3a9 100644 --- a/env.example +++ b/env.example @@ -139,7 +139,7 @@ EMBEDDING_BINDING_HOST=http://localhost:11434 # AZURE_EMBEDDING_ENDPOINT=your_endpoint # AZURE_EMBEDDING_API_KEY=your_api_key -### Jina AI Embedding +### Jina AI Embedding EMBEDDING_BINDING=jina EMBEDDING_BINDING_HOST=https://api.jina.ai/v1/embeddings EMBEDDING_MODEL=jina-embeddings-v4 diff --git a/lightrag/api/lightrag_server.py b/lightrag/api/lightrag_server.py index 5bb13813..ebffe9a3 100644 --- a/lightrag/api/lightrag_server.py +++ b/lightrag/api/lightrag_server.py @@ -89,7 +89,13 @@ def create_app(args): ]: raise Exception("llm binding not supported") - if args.embedding_binding not in ["lollms", "ollama", "openai", "azure_openai", "jina"]: + if args.embedding_binding not in [ + "lollms", + "ollama", + "openai", + "azure_openai", + "jina", + ]: raise Exception("embedding binding not supported") # Set default hosts if not provided diff --git a/lightrag/llm/jina.py b/lightrag/llm/jina.py index 720877f5..6a1e95d2 100644 --- a/lightrag/llm/jina.py +++ b/lightrag/llm/jina.py @@ -28,7 +28,7 @@ async def fetch_data(url, headers, data): request_info=response.request_info, history=response.history, status=response.status, - message=f"Jina API error: {error_text}" + message=f"Jina API error: {error_text}", ) response_json = await response.json() data_list = response_json.get("data", []) @@ -69,10 +69,10 @@ async def jina_embed( """ if api_key: os.environ["JINA_API_KEY"] = api_key - + if "JINA_API_KEY" not in os.environ: raise ValueError("JINA_API_KEY environment variable is required") - + url = base_url or "https://api.jina.ai/v1/embeddings" headers = { "Content-Type": "application/json", @@ -84,29 +84,35 @@ async def jina_embed( "dimensions": dimensions, "input": texts, } - + # Only add optional parameters if they have non-default values if late_chunking: data["late_chunking"] = late_chunking - - logger.debug(f"Jina embedding request: {len(texts)} texts, dimensions: {dimensions}") - + + logger.debug( + f"Jina embedding request: {len(texts)} texts, dimensions: {dimensions}" + ) + try: data_list = await fetch_data(url, headers, data) - + if not data_list: logger.error("Jina API returned empty data list") raise ValueError("Jina API returned empty data list") - + if len(data_list) != len(texts): - logger.error(f"Jina API returned {len(data_list)} embeddings for {len(texts)} texts") - raise ValueError(f"Jina API returned {len(data_list)} embeddings for {len(texts)} texts") - + logger.error( + f"Jina API returned {len(data_list)} embeddings for {len(texts)} texts" + ) + raise ValueError( + f"Jina API returned {len(data_list)} embeddings for {len(texts)} texts" + ) + embeddings = np.array([dp["embedding"] for dp in data_list]) logger.debug(f"Jina embeddings generated: shape {embeddings.shape}") - + return embeddings - + except Exception as e: logger.error(f"Jina embedding error: {e}") raise From d8e7b770994cfd8fecb0dcdb34c6472e00673fc7 Mon Sep 17 00:00:00 2001 From: yangdx Date: Thu, 24 Jul 2025 04:00:43 +0800 Subject: [PATCH 3/4] Bump api version to 0188 --- lightrag/api/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lightrag/api/__init__.py b/lightrag/api/__init__.py index df6c6909..744d1e58 100644 --- a/lightrag/api/__init__.py +++ b/lightrag/api/__init__.py @@ -1 +1 @@ -__api_version__ = "0187" +__api_version__ = "0188" From 5574a308565f5967472fc8810c2419630073dd72 Mon Sep 17 00:00:00 2001 From: yangdx Date: Thu, 24 Jul 2025 12:45:13 +0800 Subject: [PATCH 4/4] fix(postgres): handle ssl_mode="allow" in _create_ssl_context Add "allow" to the list of recognized SSL modes in PostgreSQL connection helper. Previously, ssl_mode="allow" would fall through to "Unknown SSL mode" warning. Now it's properly handled alongside "require" and "prefer" modes. --- lightrag/kg/postgres_impl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lightrag/kg/postgres_impl.py b/lightrag/kg/postgres_impl.py index 9ac0f96f..0e49a67f 100644 --- a/lightrag/kg/postgres_impl.py +++ b/lightrag/kg/postgres_impl.py @@ -80,7 +80,7 @@ class PostgreSQLDB: if ssl_mode in ["disable", "allow", "prefer", "require"]: if ssl_mode == "disable": return None - elif ssl_mode in ["require", "prefer"]: + elif ssl_mode in ["require", "prefer", "allow"]: # Return None for simple SSL requirement, handled in initdb return None