From bae801bbdbaf00148128b0388e90f157d784550d Mon Sep 17 00:00:00 2001 From: mmavka Date: Thu, 21 May 2026 17:37:55 +0200 Subject: [PATCH] Restore lightrag-hku 1.4.x compatibility and drop hardcoded 1536-dim embeddings Three small fixes that together let knowledge-mcp work against the latest lightrag-hku (1.4.16) plus arbitrary embedding dimensions. 1. Drop llm_model_max_token_size from the LightRAG constructor call in rag.py. The kwarg was removed upstream in 1.4.x, so KB creation was failing with TypeError before this. 2. Filter ids, model_func, chunk_top_k, and enable_rerank out of QueryParam kwargs before construction. Those fields were removed from QueryParam upstream but still ship in the default per-KB config template, so every query was crashing. 3. Replace the openai_embed wrapper in openai_embedding_func with a direct AsyncOpenAI client call. openai_embed is decorated with wrap_embedding_func_with_attrs(embedding_dim=1536), which silently pinned the embedding dimension to OpenAI text-embedding-3-* and broke any non-1536-dim endpoint (Ollama bge-m3, mxbai-embed-large, etc.). Tested locally against lightrag-hku 1.4.16, OpenRouter as the LLM endpoint, and Ollama bge-m3 (1024-dim) for embeddings. KB create, add_text, and all five query MCP tools work end-to-end after the patches; before, each issue blocks at a different stage. --- knowledge_mcp/openai_func.py | 20 +++++++++++--------- knowledge_mcp/rag.py | 10 +++++++--- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/knowledge_mcp/openai_func.py b/knowledge_mcp/openai_func.py index aaf1c82..3d77f9d 100644 --- a/knowledge_mcp/openai_func.py +++ b/knowledge_mcp/openai_func.py @@ -77,15 +77,17 @@ async def vision_model_func(prompt, system_prompt=None, history_messages=[], ima ) async def openai_embedding_func(texts: list[str]) -> np.ndarray: - embedding_model = Config.get_instance().lightrag.embedding.model_name - - # Call the OpenAI embed function with all the parameters - return await openai_embed( - texts=texts, - model=embedding_model, - api_key=Config.get_instance().lightrag.embedding.api_key, - base_url=Config.get_instance().lightrag.embedding.api_base, - ) + # Call the OpenAI-compatible embeddings endpoint directly instead of going + # through lightrag.llm.openai.openai_embed, which is decorated with + # @wrap_embedding_func_with_attrs(embedding_dim=1536) and therefore pins + # the dimension to OpenAI text-embedding-3-*. Going direct lets the caller + # use any model (Ollama bge-m3 at 1024, mxbai-embed-large, etc.); the + # EmbeddingFunc wrapper below still validates the configured dimension. + cfg = Config.get_instance().lightrag.embedding + from openai import AsyncOpenAI + client = AsyncOpenAI(api_key=cfg.api_key, base_url=cfg.api_base) + response = await client.embeddings.create(model=cfg.model_name, input=texts) + return np.array([item.embedding for item in response.data]) # Wrap the embedding function with the correct attributes from config embedding_func = EmbeddingFunc( diff --git a/knowledge_mcp/rag.py b/knowledge_mcp/rag.py index 553b9a6..0f46ded 100644 --- a/knowledge_mcp/rag.py +++ b/knowledge_mcp/rag.py @@ -135,8 +135,7 @@ async def create_rag_instance(self, kb_name: str) -> RAGAnything: working_dir=str(kb_path), llm_model_func=llm_func, llm_model_kwargs=llm_kwargs, - llm_model_name=llm_config.model_name, - llm_model_max_token_size=llm_model_max_tokens, + llm_model_name=llm_config.model_name, embedding_func=embed_func, embedding_cache_config={ "enabled": cache_config.enabled, @@ -248,7 +247,12 @@ async def query(self, kb_name: str, query_text: str, **kwargs: Any) -> Any: # Ensure 'description' is not passed as a query param final_query_params.pop("description", None) - + + # Filter out kwargs that older per-KB configs may still emit but + # which were removed from lightrag-hku's QueryParam in 1.4.x. + for stale_kwarg in ("ids", "model_func", "chunk_top_k", "enable_rerank"): + final_query_params.pop(stale_kwarg, None) + # Handle user_prompt: only include if not empty user_prompt_value = final_query_params.get('user_prompt', '') if not (user_prompt_value and user_prompt_value.strip()):