yujunjun/LightRAG
1
0
Fork 0
mirror of https://github.com/HKUDS/LightRAG.git synced 2025-06-26 22:00:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix linting

Browse Source
This commit is contained in:
yangdx 2025-05-14 01:14:45 +08:00
parent 5c9fd9c4d2
commit 0e26cbebd0
2 changed files with 19 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
examples/lightrag_ollama_demo.py
View File

@ -26,9 +26,7 @@ def configure_logging():
# Get log directory path from environment variable or use current directory # Get log directory path from environment variable or use current directory
log_dir = os.getenv("LOG_DIR", os.getcwd()) log_dir = os.getenv("LOG_DIR", os.getcwd())
log_file_path = os.path.abspath( log_file_path = os.path.abspath(os.path.join(log_dir, "lightrag_ollama_demo.log"))
os.path.join(log_dir, "lightrag_ollama_demo.log")
)
print(f"\nLightRAG compatible demo log file: {log_file_path}\n") print(f"\nLightRAG compatible demo log file: {log_file_path}\n")
os.makedirs(os.path.dirname(log_file_path), exist_ok=True) os.makedirs(os.path.dirname(log_file_path), exist_ok=True)
@ -211,6 +209,7 @@ async def main():
await rag.llm_response_cache.index_done_callback() await rag.llm_response_cache.index_done_callback()
await rag.finalize_storages() await rag.finalize_storages()
if __name__ == "__main__": if __name__ == "__main__":
# Configure logging before running the main function # Configure logging before running the main function
configure_logging() configure_logging()

26
lightrag/llm/ollama.py
View File

@ -62,9 +62,9 @@ async def _ollama_model_if_cache(
} }
if api_key: if api_key:
headers["Authorization"] = f"Bearer {api_key}" headers["Authorization"] = f"Bearer {api_key}"
ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers) ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers)
try: try:
messages = [] messages = []
if system_prompt: if system_prompt:
@ -106,15 +106,21 @@ async def _ollama_model_if_cache(
await ollama_client._client.aclose() await ollama_client._client.aclose()
logger.debug("Successfully closed Ollama client after exception") logger.debug("Successfully closed Ollama client after exception")
except Exception as close_error: except Exception as close_error:
logger.warning(f"Failed to close Ollama client after exception: {close_error}") logger.warning(
f"Failed to close Ollama client after exception: {close_error}"
)
raise e raise e
finally: finally:
if not stream: if not stream:
try: try:
await ollama_client._client.aclose() await ollama_client._client.aclose()
logger.debug("Successfully closed Ollama client for non-streaming response") logger.debug(
"Successfully closed Ollama client for non-streaming response"
)
except Exception as close_error: except Exception as close_error:
logger.warning(f"Failed to close Ollama client in finally block: {close_error}") logger.warning(
f"Failed to close Ollama client in finally block: {close_error}"
)
async def ollama_model_complete( async def ollama_model_complete(
@ -141,12 +147,12 @@ async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray:
} }
if api_key: if api_key:
headers["Authorization"] = f"Bearer {api_key}" headers["Authorization"] = f"Bearer {api_key}"
host = kwargs.pop("host", None) host = kwargs.pop("host", None)
timeout = kwargs.pop("timeout", None) or 90 # Default time out 90s timeout = kwargs.pop("timeout", None) or 90 # Default time out 90s
ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers) ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers)
try: try:
data = await ollama_client.embed(model=embed_model, input=texts) data = await ollama_client.embed(model=embed_model, input=texts)
return np.array(data["embeddings"]) return np.array(data["embeddings"])
@ -156,7 +162,9 @@ async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray:
await ollama_client._client.aclose() await ollama_client._client.aclose()
logger.debug("Successfully closed Ollama client after exception in embed") logger.debug("Successfully closed Ollama client after exception in embed")
except Exception as close_error: except Exception as close_error:
logger.warning(f"Failed to close Ollama client after exception in embed: {close_error}") logger.warning(
f"Failed to close Ollama client after exception in embed: {close_error}"
)
raise e raise e
finally: finally:
try: try: