Fix linting
This commit is contained in:
@@ -41,12 +41,14 @@ from .ollama_api import ollama_server_infos
|
|||||||
# Load environment variables
|
# Load environment variables
|
||||||
load_dotenv()
|
load_dotenv()
|
||||||
|
|
||||||
|
|
||||||
class RAGStorageConfig:
|
class RAGStorageConfig:
|
||||||
KV_STORAGE = "JsonKVStorage"
|
KV_STORAGE = "JsonKVStorage"
|
||||||
DOC_STATUS_STORAGE = "JsonDocStatusStorage"
|
DOC_STATUS_STORAGE = "JsonDocStatusStorage"
|
||||||
GRAPH_STORAGE = "NetworkXStorage"
|
GRAPH_STORAGE = "NetworkXStorage"
|
||||||
VECTOR_STORAGE = "NanoVectorDBStorage"
|
VECTOR_STORAGE = "NanoVectorDBStorage"
|
||||||
|
|
||||||
|
|
||||||
# Initialize rag storage config
|
# Initialize rag storage config
|
||||||
rag_storage_config = RAGStorageConfig()
|
rag_storage_config = RAGStorageConfig()
|
||||||
|
|
||||||
@@ -592,6 +594,7 @@ class SearchMode(str, Enum):
|
|||||||
hybrid = "hybrid"
|
hybrid = "hybrid"
|
||||||
mix = "mix"
|
mix = "mix"
|
||||||
|
|
||||||
|
|
||||||
class QueryRequest(BaseModel):
|
class QueryRequest(BaseModel):
|
||||||
query: str
|
query: str
|
||||||
mode: SearchMode = SearchMode.hybrid
|
mode: SearchMode = SearchMode.hybrid
|
||||||
|
@@ -12,6 +12,7 @@ import asyncio
|
|||||||
from ascii_colors import trace_exception
|
from ascii_colors import trace_exception
|
||||||
from lightrag import LightRAG, QueryParam
|
from lightrag import LightRAG, QueryParam
|
||||||
|
|
||||||
|
|
||||||
class OllamaServerInfos:
|
class OllamaServerInfos:
|
||||||
# Constants for emulated Ollama model information
|
# Constants for emulated Ollama model information
|
||||||
LIGHTRAG_NAME = "lightrag"
|
LIGHTRAG_NAME = "lightrag"
|
||||||
@@ -21,8 +22,10 @@ class OllamaServerInfos:
|
|||||||
LIGHTRAG_CREATED_AT = "2024-01-15T00:00:00Z"
|
LIGHTRAG_CREATED_AT = "2024-01-15T00:00:00Z"
|
||||||
LIGHTRAG_DIGEST = "sha256:lightrag"
|
LIGHTRAG_DIGEST = "sha256:lightrag"
|
||||||
|
|
||||||
|
|
||||||
ollama_server_infos = OllamaServerInfos()
|
ollama_server_infos = OllamaServerInfos()
|
||||||
|
|
||||||
|
|
||||||
# query mode according to query prefix (bypass is not LightRAG quer mode)
|
# query mode according to query prefix (bypass is not LightRAG quer mode)
|
||||||
class SearchMode(str, Enum):
|
class SearchMode(str, Enum):
|
||||||
naive = "naive"
|
naive = "naive"
|
||||||
@@ -32,11 +35,13 @@ class SearchMode(str, Enum):
|
|||||||
mix = "mix"
|
mix = "mix"
|
||||||
bypass = "bypass"
|
bypass = "bypass"
|
||||||
|
|
||||||
|
|
||||||
class OllamaMessage(BaseModel):
|
class OllamaMessage(BaseModel):
|
||||||
role: str
|
role: str
|
||||||
content: str
|
content: str
|
||||||
images: Optional[List[str]] = None
|
images: Optional[List[str]] = None
|
||||||
|
|
||||||
|
|
||||||
class OllamaChatRequest(BaseModel):
|
class OllamaChatRequest(BaseModel):
|
||||||
model: str
|
model: str
|
||||||
messages: List[OllamaMessage]
|
messages: List[OllamaMessage]
|
||||||
@@ -44,12 +49,14 @@ class OllamaChatRequest(BaseModel):
|
|||||||
options: Optional[Dict[str, Any]] = None
|
options: Optional[Dict[str, Any]] = None
|
||||||
system: Optional[str] = None
|
system: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
class OllamaChatResponse(BaseModel):
|
class OllamaChatResponse(BaseModel):
|
||||||
model: str
|
model: str
|
||||||
created_at: str
|
created_at: str
|
||||||
message: OllamaMessage
|
message: OllamaMessage
|
||||||
done: bool
|
done: bool
|
||||||
|
|
||||||
|
|
||||||
class OllamaGenerateRequest(BaseModel):
|
class OllamaGenerateRequest(BaseModel):
|
||||||
model: str
|
model: str
|
||||||
prompt: str
|
prompt: str
|
||||||
@@ -57,6 +64,7 @@ class OllamaGenerateRequest(BaseModel):
|
|||||||
stream: bool = False
|
stream: bool = False
|
||||||
options: Optional[Dict[str, Any]] = None
|
options: Optional[Dict[str, Any]] = None
|
||||||
|
|
||||||
|
|
||||||
class OllamaGenerateResponse(BaseModel):
|
class OllamaGenerateResponse(BaseModel):
|
||||||
model: str
|
model: str
|
||||||
created_at: str
|
created_at: str
|
||||||
@@ -70,9 +78,11 @@ class OllamaGenerateResponse(BaseModel):
|
|||||||
eval_count: Optional[int]
|
eval_count: Optional[int]
|
||||||
eval_duration: Optional[int]
|
eval_duration: Optional[int]
|
||||||
|
|
||||||
|
|
||||||
class OllamaVersionResponse(BaseModel):
|
class OllamaVersionResponse(BaseModel):
|
||||||
version: str
|
version: str
|
||||||
|
|
||||||
|
|
||||||
class OllamaModelDetails(BaseModel):
|
class OllamaModelDetails(BaseModel):
|
||||||
parent_model: str
|
parent_model: str
|
||||||
format: str
|
format: str
|
||||||
@@ -81,6 +91,7 @@ class OllamaModelDetails(BaseModel):
|
|||||||
parameter_size: str
|
parameter_size: str
|
||||||
quantization_level: str
|
quantization_level: str
|
||||||
|
|
||||||
|
|
||||||
class OllamaModel(BaseModel):
|
class OllamaModel(BaseModel):
|
||||||
name: str
|
name: str
|
||||||
model: str
|
model: str
|
||||||
@@ -89,9 +100,11 @@ class OllamaModel(BaseModel):
|
|||||||
modified_at: str
|
modified_at: str
|
||||||
details: OllamaModelDetails
|
details: OllamaModelDetails
|
||||||
|
|
||||||
|
|
||||||
class OllamaTagResponse(BaseModel):
|
class OllamaTagResponse(BaseModel):
|
||||||
models: List[OllamaModel]
|
models: List[OllamaModel]
|
||||||
|
|
||||||
|
|
||||||
def estimate_tokens(text: str) -> int:
|
def estimate_tokens(text: str) -> int:
|
||||||
"""Estimate the number of tokens in text
|
"""Estimate the number of tokens in text
|
||||||
Chinese characters: approximately 1.5 tokens per character
|
Chinese characters: approximately 1.5 tokens per character
|
||||||
@@ -106,6 +119,7 @@ def estimate_tokens(text: str) -> int:
|
|||||||
|
|
||||||
return int(tokens)
|
return int(tokens)
|
||||||
|
|
||||||
|
|
||||||
def parse_query_mode(query: str) -> tuple[str, SearchMode]:
|
def parse_query_mode(query: str) -> tuple[str, SearchMode]:
|
||||||
"""Parse query prefix to determine search mode
|
"""Parse query prefix to determine search mode
|
||||||
Returns tuple of (cleaned_query, search_mode)
|
Returns tuple of (cleaned_query, search_mode)
|
||||||
@@ -127,6 +141,7 @@ def parse_query_mode(query: str) -> tuple[str, SearchMode]:
|
|||||||
|
|
||||||
return query, SearchMode.hybrid
|
return query, SearchMode.hybrid
|
||||||
|
|
||||||
|
|
||||||
class OllamaAPI:
|
class OllamaAPI:
|
||||||
def __init__(self, rag: LightRAG):
|
def __init__(self, rag: LightRAG):
|
||||||
self.rag = rag
|
self.rag = rag
|
||||||
@@ -333,10 +348,13 @@ class OllamaAPI:
|
|||||||
"stream": request.stream,
|
"stream": request.stream,
|
||||||
"only_need_context": False,
|
"only_need_context": False,
|
||||||
"conversation_history": conversation_history,
|
"conversation_history": conversation_history,
|
||||||
"top_k": self.rag.args.top_k if hasattr(self.rag, 'args') else 50,
|
"top_k": self.rag.args.top_k if hasattr(self.rag, "args") else 50,
|
||||||
}
|
}
|
||||||
|
|
||||||
if hasattr(self.rag, 'args') and self.rag.args.history_turns is not None:
|
if (
|
||||||
|
hasattr(self.rag, "args")
|
||||||
|
and self.rag.args.history_turns is not None
|
||||||
|
):
|
||||||
param_dict["history_turns"] = self.rag.args.history_turns
|
param_dict["history_turns"] = self.rag.args.history_turns
|
||||||
|
|
||||||
query_param = QueryParam(**param_dict)
|
query_param = QueryParam(**param_dict)
|
||||||
@@ -521,7 +539,9 @@ class OllamaAPI:
|
|||||||
**self.rag.llm_model_kwargs,
|
**self.rag.llm_model_kwargs,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
response_text = await self.rag.aquery(cleaned_query, param=query_param)
|
response_text = await self.rag.aquery(
|
||||||
|
cleaned_query, param=query_param
|
||||||
|
)
|
||||||
|
|
||||||
last_chunk_time = time.time_ns()
|
last_chunk_time = time.time_ns()
|
||||||
|
|
||||||
|
Reference in New Issue
Block a user