Update __version__

examples/lightrag_zhipu_demo.py

@@ -1,9 +1,6 @@
-import asyncio
 import os
-import inspect
 import logging
 
-from dotenv import load_dotenv
 
 from lightrag import LightRAG, QueryParam
 from lightrag.llm import zhipu_complete, zhipu_embedding
@@ -21,7 +18,6 @@ if api_key is None:
     raise Exception("Please set ZHIPU_API_KEY in your environment")
 
 
-
 rag = LightRAG(
     working_dir=WORKING_DIR,
     llm_model_func=zhipu_complete,
@@ -31,9 +27,7 @@ rag = LightRAG(
     embedding_func=EmbeddingFunc(
         embedding_dim=2048,  # Zhipu embedding-3 dimension
         max_token_size=8192,
-        func=lambda texts: zhipu_embedding(
+        func=lambda texts: zhipu_embedding(texts),
-            texts
-        ),
     ),
 )
 

lightrag/__init__.py

@@ -1,5 +1,5 @@
 from .lightrag import LightRAG as LightRAG, QueryParam as QueryParam
 
-__version__ = "1.0.5"
+__version__ = "1.0.6"
 __author__ = "Zirui Guo"
 __url__ = "https://github.com/HKUDS/LightRAG"

lightrag/kg/milvus_impl.py

@@ -63,7 +63,9 @@ class MilvusVectorDBStorge(BaseVectorStorage):
             return result
 
         embedding_tasks = [wrapped_task(batch) for batch in batches]
-        pbar = tqdm_async(total=len(embedding_tasks), desc="Generating embeddings", unit="batch")
+        pbar = tqdm_async(
+            total=len(embedding_tasks), desc="Generating embeddings", unit="batch"
+        )
         embeddings_list = await asyncio.gather(*embedding_tasks)
 
         embeddings = np.concatenate(embeddings_list)

lightrag/llm.py

@@ -608,7 +608,7 @@ async def zhipu_complete_if_cache(
     api_key: Optional[str] = None,
     system_prompt: Optional[str] = None,
     history_messages: List[Dict[str, str]] = [],
-    **kwargs
+    **kwargs,
 ) -> str:
     # dynamically load ZhipuAI
     try:
@@ -640,13 +640,11 @@ async def zhipu_complete_if_cache(
     logger.debug(f"System prompt: {system_prompt}")
 
     # Remove unsupported kwargs
-    kwargs = {k: v for k, v in kwargs.items() if k not in ['hashing_kv', 'keyword_extraction']}
+    kwargs = {
+        k: v for k, v in kwargs.items() if k not in ["hashing_kv", "keyword_extraction"]
+    }
 
-    response = client.chat.completions.create(
+    response = client.chat.completions.create(model=model, messages=messages, **kwargs)
-            model=model,
-            messages=messages,
-            **kwargs
-        )
 
     return response.choices[0].message.content
 
@@ -683,7 +681,7 @@ async def zhipu_complete(
                 prompt=prompt,
                 system_prompt=system_prompt,
                 history_messages=history_messages,
-                **kwargs
+                **kwargs,
             )
 
             # Try to parse as JSON
@@ -691,7 +689,7 @@ async def zhipu_complete(
                 data = json.loads(response)
                 return GPTKeywordExtractionFormat(
                     high_level_keywords=data.get("high_level_keywords", []),
-                    low_level_keywords=data.get("low_level_keywords", [])
+                    low_level_keywords=data.get("low_level_keywords", []),
                 )
             except json.JSONDecodeError:
                 # If direct JSON parsing fails, try to extract JSON from text
@@ -701,13 +699,15 @@ async def zhipu_complete(
                         data = json.loads(match.group())
                         return GPTKeywordExtractionFormat(
                             high_level_keywords=data.get("high_level_keywords", []),
-                            low_level_keywords=data.get("low_level_keywords", [])
+                            low_level_keywords=data.get("low_level_keywords", []),
                         )
                     except json.JSONDecodeError:
                         pass
 
                 # If all parsing fails, log warning and return empty format
-                logger.warning(f"Failed to parse keyword extraction response: {response}")
+                logger.warning(
+                    f"Failed to parse keyword extraction response: {response}"
+                )
                 return GPTKeywordExtractionFormat(
                     high_level_keywords=[], low_level_keywords=[]
                 )
@@ -722,7 +722,7 @@ async def zhipu_complete(
             prompt=prompt,
             system_prompt=system_prompt,
             history_messages=history_messages,
-            **kwargs
+            **kwargs,
         )
 
 
@@ -733,13 +733,9 @@ async def zhipu_complete(
     retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
 )
 async def zhipu_embedding(
-    texts: list[str],
+    texts: list[str], model: str = "embedding-3", api_key: str = None, **kwargs
-    model: str = "embedding-3",
-    api_key: str = None,
-    **kwargs
 ) -> np.ndarray:
-
+    # dynamically load ZhipuAI
-# dynamically load ZhipuAI
     try:
         from zhipuai import ZhipuAI
     except ImportError:
@@ -758,11 +754,7 @@ async def zhipu_embedding(
     embeddings = []
     for text in texts:
         try:
-            response = client.embeddings.create(
+            response = client.embeddings.create(model=model, input=[text], **kwargs)
-                model=model,
-                input=[text],
-                **kwargs
-            )
             embeddings.append(response.data[0].embedding)
         except Exception as e:
             raise Exception(f"Error calling ChatGLM Embedding API: {str(e)}")

lightrag/storage.py

@@ -103,7 +103,9 @@ class NanoVectorDBStorage(BaseVectorStorage):
             return result
 
         embedding_tasks = [wrapped_task(batch) for batch in batches]
-        pbar = tqdm_async(total=len(embedding_tasks), desc="Generating embeddings", unit="batch")
+        pbar = tqdm_async(
+            total=len(embedding_tasks), desc="Generating embeddings", unit="batch"
+        )
         embeddings_list = await asyncio.gather(*embedding_tasks)
 
         embeddings = np.concatenate(embeddings_list)