fixed linting
This commit is contained in:
@@ -23,10 +23,6 @@ Requirements:
|
|||||||
- Lightrag
|
- Lightrag
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from pathlib import Path
|
|
||||||
from typing import Optional, List, Dict, Union, Any
|
|
||||||
from datetime import datetime
|
|
||||||
|
|
||||||
# Tool version
|
# Tool version
|
||||||
__version__ = "1.0.0"
|
__version__ = "1.0.0"
|
||||||
__author__ = "ParisNeo"
|
__author__ = "ParisNeo"
|
||||||
|
@@ -298,14 +298,12 @@ def parse_args() -> argparse.Namespace:
|
|||||||
help=f"llm server host URL (default: from env or {default_llm_host})",
|
help=f"llm server host URL (default: from env or {default_llm_host})",
|
||||||
)
|
)
|
||||||
|
|
||||||
default_llm_api_key = get_env_value(
|
default_llm_api_key = get_env_value("LLM_BINDING_API_KEY", None)
|
||||||
"LLM_BINDING_API_KEY", None
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--llm-binding-api-key",
|
"--llm-binding-api-key",
|
||||||
default=default_llm_api_key,
|
default=default_llm_api_key,
|
||||||
help=f"llm server API key (default: from env or empty string)",
|
help="llm server API key (default: from env or empty string)",
|
||||||
)
|
)
|
||||||
|
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
@@ -324,13 +322,11 @@ def parse_args() -> argparse.Namespace:
|
|||||||
help=f"embedding server host URL (default: from env or {default_embedding_host})",
|
help=f"embedding server host URL (default: from env or {default_embedding_host})",
|
||||||
)
|
)
|
||||||
|
|
||||||
default_embedding_api_key = get_env_value(
|
default_embedding_api_key = get_env_value("EMBEDDING_BINDING_API_KEY", "")
|
||||||
"EMBEDDING_BINDING_API_KEY", ""
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--embedding-binding-api-key",
|
"--embedding-binding-api-key",
|
||||||
default=default_embedding_api_key,
|
default=default_embedding_api_key,
|
||||||
help=f"embedding server API key (default: from env or empty string)",
|
help="embedding server API key (default: from env or empty string)",
|
||||||
)
|
)
|
||||||
|
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
@@ -649,26 +645,26 @@ def create_app(args):
|
|||||||
texts,
|
texts,
|
||||||
embed_model=args.embedding_model,
|
embed_model=args.embedding_model,
|
||||||
host=args.embedding_binding_host,
|
host=args.embedding_binding_host,
|
||||||
api_key = args.embedding_binding_api_key
|
api_key=args.embedding_binding_api_key,
|
||||||
)
|
)
|
||||||
if args.embedding_binding == "lollms"
|
if args.embedding_binding == "lollms"
|
||||||
else ollama_embed(
|
else ollama_embed(
|
||||||
texts,
|
texts,
|
||||||
embed_model=args.embedding_model,
|
embed_model=args.embedding_model,
|
||||||
host=args.embedding_binding_host,
|
host=args.embedding_binding_host,
|
||||||
api_key = args.embedding_binding_api_key
|
api_key=args.embedding_binding_api_key,
|
||||||
)
|
)
|
||||||
if args.embedding_binding == "ollama"
|
if args.embedding_binding == "ollama"
|
||||||
else azure_openai_embedding(
|
else azure_openai_embedding(
|
||||||
texts,
|
texts,
|
||||||
model=args.embedding_model, # no host is used for openai,
|
model=args.embedding_model, # no host is used for openai,
|
||||||
api_key = args.embedding_binding_api_key
|
api_key=args.embedding_binding_api_key,
|
||||||
)
|
)
|
||||||
if args.embedding_binding == "azure_openai"
|
if args.embedding_binding == "azure_openai"
|
||||||
else openai_embedding(
|
else openai_embedding(
|
||||||
texts,
|
texts,
|
||||||
model=args.embedding_model, # no host is used for openai,
|
model=args.embedding_model, # no host is used for openai,
|
||||||
api_key = args.embedding_binding_api_key
|
api_key=args.embedding_binding_api_key,
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -686,7 +682,7 @@ def create_app(args):
|
|||||||
"host": args.llm_binding_host,
|
"host": args.llm_binding_host,
|
||||||
"timeout": args.timeout,
|
"timeout": args.timeout,
|
||||||
"options": {"num_ctx": args.max_tokens},
|
"options": {"num_ctx": args.max_tokens},
|
||||||
"api_key": args.llm_binding_api_key
|
"api_key": args.llm_binding_api_key,
|
||||||
},
|
},
|
||||||
embedding_func=embedding_func,
|
embedding_func=embedding_func,
|
||||||
)
|
)
|
||||||
|
@@ -350,7 +350,7 @@ async def ollama_model_if_cache(
|
|||||||
timeout = kwargs.pop("timeout", None)
|
timeout = kwargs.pop("timeout", None)
|
||||||
kwargs.pop("hashing_kv", None)
|
kwargs.pop("hashing_kv", None)
|
||||||
api_key = kwargs.pop("api_key", None)
|
api_key = kwargs.pop("api_key", None)
|
||||||
headers={'Authorization': f'Bearer {api_key}'} if api_key else None
|
headers = {"Authorization": f"Bearer {api_key}"} if api_key else None
|
||||||
ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers)
|
ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers)
|
||||||
messages = []
|
messages = []
|
||||||
if system_prompt:
|
if system_prompt:
|
||||||
@@ -383,7 +383,7 @@ async def lollms_model_if_cache(
|
|||||||
|
|
||||||
stream = True if kwargs.get("stream") else False
|
stream = True if kwargs.get("stream") else False
|
||||||
api_key = kwargs.pop("api_key", None)
|
api_key = kwargs.pop("api_key", None)
|
||||||
headers={'Authorization': f'Bearer {api_key}'} if api_key else None
|
headers = {"Authorization": f"Bearer {api_key}"} if api_key else None
|
||||||
|
|
||||||
# Extract lollms specific parameters
|
# Extract lollms specific parameters
|
||||||
request_data = {
|
request_data = {
|
||||||
@@ -626,7 +626,12 @@ async def nvidia_openai_complete(
|
|||||||
|
|
||||||
|
|
||||||
async def azure_openai_complete(
|
async def azure_openai_complete(
|
||||||
model: str = "gpt-4o-mini", prompt="", system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
|
model: str = "gpt-4o-mini",
|
||||||
|
prompt="",
|
||||||
|
system_prompt=None,
|
||||||
|
history_messages=[],
|
||||||
|
keyword_extraction=False,
|
||||||
|
**kwargs,
|
||||||
) -> str:
|
) -> str:
|
||||||
keyword_extraction = kwargs.pop("keyword_extraction", None)
|
keyword_extraction = kwargs.pop("keyword_extraction", None)
|
||||||
result = await azure_openai_complete_if_cache(
|
result = await azure_openai_complete_if_cache(
|
||||||
@@ -1153,7 +1158,11 @@ async def ollama_embedding(texts: list[str], embed_model, **kwargs) -> np.ndarra
|
|||||||
|
|
||||||
async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray:
|
async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray:
|
||||||
api_key = kwargs.pop("api_key", None)
|
api_key = kwargs.pop("api_key", None)
|
||||||
headers = {"Authorization": api_key, "Content-Type": "application/json"} if api_key else None
|
headers = (
|
||||||
|
{"Authorization": api_key, "Content-Type": "application/json"}
|
||||||
|
if api_key
|
||||||
|
else None
|
||||||
|
)
|
||||||
kwargs["headers"] = headers
|
kwargs["headers"] = headers
|
||||||
ollama_client = ollama.Client(**kwargs)
|
ollama_client = ollama.Client(**kwargs)
|
||||||
data = ollama_client.embed(model=embed_model, input=texts)
|
data = ollama_client.embed(model=embed_model, input=texts)
|
||||||
@@ -1176,14 +1185,19 @@ async def lollms_embed(
|
|||||||
np.ndarray: Array of embeddings
|
np.ndarray: Array of embeddings
|
||||||
"""
|
"""
|
||||||
api_key = kwargs.pop("api_key", None)
|
api_key = kwargs.pop("api_key", None)
|
||||||
headers = {"Authorization": api_key, "Content-Type": "application/json"} if api_key else None
|
headers = (
|
||||||
|
{"Authorization": api_key, "Content-Type": "application/json"}
|
||||||
|
if api_key
|
||||||
|
else None
|
||||||
|
)
|
||||||
async with aiohttp.ClientSession(headers=headers) as session:
|
async with aiohttp.ClientSession(headers=headers) as session:
|
||||||
embeddings = []
|
embeddings = []
|
||||||
for text in texts:
|
for text in texts:
|
||||||
request_data = {"text": text}
|
request_data = {"text": text}
|
||||||
|
|
||||||
async with session.post(
|
async with session.post(
|
||||||
f"{base_url}/lollms_embed", json=request_data,
|
f"{base_url}/lollms_embed",
|
||||||
|
json=request_data,
|
||||||
) as response:
|
) as response:
|
||||||
result = await response.json()
|
result = await response.json()
|
||||||
embeddings.append(result["vector"])
|
embeddings.append(result["vector"])
|
||||||
|
Reference in New Issue
Block a user