zrguo
commited on
Commit
·
b6cfd8e
1
Parent(s):
46fd949
fix lint
Browse files- lightrag/llm/ollama.py +3 -3
lightrag/llm/ollama.py
CHANGED
@@ -50,7 +50,7 @@ async def _ollama_model_if_cache(
|
|
50 |
kwargs.pop("max_tokens", None)
|
51 |
# kwargs.pop("response_format", None) # allow json
|
52 |
host = kwargs.pop("host", None)
|
53 |
-
timeout = kwargs.pop("timeout", None) or 600 # Default timeout 600s
|
54 |
kwargs.pop("hashing_kv", None)
|
55 |
api_key = kwargs.pop("api_key", None)
|
56 |
headers = {
|
@@ -146,7 +146,7 @@ async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray:
|
|
146 |
headers["Authorization"] = f"Bearer {api_key}"
|
147 |
|
148 |
host = kwargs.pop("host", None)
|
149 |
-
timeout = kwargs.pop("timeout", None) or 300 # Default time out 300s
|
150 |
|
151 |
ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers)
|
152 |
|
@@ -168,4 +168,4 @@ async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray:
|
|
168 |
await ollama_client._client.aclose()
|
169 |
logger.debug("Successfully closed Ollama client after embed")
|
170 |
except Exception as close_error:
|
171 |
-
logger.warning(f"Failed to close Ollama client after embed: {close_error}")
|
|
|
50 |
kwargs.pop("max_tokens", None)
|
51 |
# kwargs.pop("response_format", None) # allow json
|
52 |
host = kwargs.pop("host", None)
|
53 |
+
timeout = kwargs.pop("timeout", None) or 600 # Default timeout 600s
|
54 |
kwargs.pop("hashing_kv", None)
|
55 |
api_key = kwargs.pop("api_key", None)
|
56 |
headers = {
|
|
|
146 |
headers["Authorization"] = f"Bearer {api_key}"
|
147 |
|
148 |
host = kwargs.pop("host", None)
|
149 |
+
timeout = kwargs.pop("timeout", None) or 300 # Default time out 300s
|
150 |
|
151 |
ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers)
|
152 |
|
|
|
168 |
await ollama_client._client.aclose()
|
169 |
logger.debug("Successfully closed Ollama client after embed")
|
170 |
except Exception as close_error:
|
171 |
+
logger.warning(f"Failed to close Ollama client after embed: {close_error}")
|