yangdx
commited on
Commit
·
3f47dc3
1
Parent(s):
efd74df
Fix linting
Browse files
lightrag/api/lightrag_server.py
CHANGED
@@ -41,12 +41,14 @@ from .ollama_api import ollama_server_infos
|
|
41 |
# Load environment variables
|
42 |
load_dotenv()
|
43 |
|
|
|
44 |
class RAGStorageConfig:
|
45 |
KV_STORAGE = "JsonKVStorage"
|
46 |
DOC_STATUS_STORAGE = "JsonDocStatusStorage"
|
47 |
GRAPH_STORAGE = "NetworkXStorage"
|
48 |
VECTOR_STORAGE = "NanoVectorDBStorage"
|
49 |
|
|
|
50 |
# Initialize rag storage config
|
51 |
rag_storage_config = RAGStorageConfig()
|
52 |
|
@@ -592,6 +594,7 @@ class SearchMode(str, Enum):
|
|
592 |
hybrid = "hybrid"
|
593 |
mix = "mix"
|
594 |
|
|
|
595 |
class QueryRequest(BaseModel):
|
596 |
query: str
|
597 |
mode: SearchMode = SearchMode.hybrid
|
|
|
41 |
# Load environment variables
|
42 |
load_dotenv()
|
43 |
|
44 |
+
|
45 |
class RAGStorageConfig:
|
46 |
KV_STORAGE = "JsonKVStorage"
|
47 |
DOC_STATUS_STORAGE = "JsonDocStatusStorage"
|
48 |
GRAPH_STORAGE = "NetworkXStorage"
|
49 |
VECTOR_STORAGE = "NanoVectorDBStorage"
|
50 |
|
51 |
+
|
52 |
# Initialize rag storage config
|
53 |
rag_storage_config = RAGStorageConfig()
|
54 |
|
|
|
594 |
hybrid = "hybrid"
|
595 |
mix = "mix"
|
596 |
|
597 |
+
|
598 |
class QueryRequest(BaseModel):
|
599 |
query: str
|
600 |
mode: SearchMode = SearchMode.hybrid
|
lightrag/api/ollama_api.py
CHANGED
@@ -12,6 +12,7 @@ import asyncio
|
|
12 |
from ascii_colors import trace_exception
|
13 |
from lightrag import LightRAG, QueryParam
|
14 |
|
|
|
15 |
class OllamaServerInfos:
|
16 |
# Constants for emulated Ollama model information
|
17 |
LIGHTRAG_NAME = "lightrag"
|
@@ -21,8 +22,10 @@ class OllamaServerInfos:
|
|
21 |
LIGHTRAG_CREATED_AT = "2024-01-15T00:00:00Z"
|
22 |
LIGHTRAG_DIGEST = "sha256:lightrag"
|
23 |
|
|
|
24 |
ollama_server_infos = OllamaServerInfos()
|
25 |
|
|
|
26 |
# query mode according to query prefix (bypass is not LightRAG quer mode)
|
27 |
class SearchMode(str, Enum):
|
28 |
naive = "naive"
|
@@ -32,11 +35,13 @@ class SearchMode(str, Enum):
|
|
32 |
mix = "mix"
|
33 |
bypass = "bypass"
|
34 |
|
|
|
35 |
class OllamaMessage(BaseModel):
|
36 |
role: str
|
37 |
content: str
|
38 |
images: Optional[List[str]] = None
|
39 |
|
|
|
40 |
class OllamaChatRequest(BaseModel):
|
41 |
model: str
|
42 |
messages: List[OllamaMessage]
|
@@ -44,12 +49,14 @@ class OllamaChatRequest(BaseModel):
|
|
44 |
options: Optional[Dict[str, Any]] = None
|
45 |
system: Optional[str] = None
|
46 |
|
|
|
47 |
class OllamaChatResponse(BaseModel):
|
48 |
model: str
|
49 |
created_at: str
|
50 |
message: OllamaMessage
|
51 |
done: bool
|
52 |
|
|
|
53 |
class OllamaGenerateRequest(BaseModel):
|
54 |
model: str
|
55 |
prompt: str
|
@@ -57,6 +64,7 @@ class OllamaGenerateRequest(BaseModel):
|
|
57 |
stream: bool = False
|
58 |
options: Optional[Dict[str, Any]] = None
|
59 |
|
|
|
60 |
class OllamaGenerateResponse(BaseModel):
|
61 |
model: str
|
62 |
created_at: str
|
@@ -70,9 +78,11 @@ class OllamaGenerateResponse(BaseModel):
|
|
70 |
eval_count: Optional[int]
|
71 |
eval_duration: Optional[int]
|
72 |
|
|
|
73 |
class OllamaVersionResponse(BaseModel):
|
74 |
version: str
|
75 |
|
|
|
76 |
class OllamaModelDetails(BaseModel):
|
77 |
parent_model: str
|
78 |
format: str
|
@@ -81,6 +91,7 @@ class OllamaModelDetails(BaseModel):
|
|
81 |
parameter_size: str
|
82 |
quantization_level: str
|
83 |
|
|
|
84 |
class OllamaModel(BaseModel):
|
85 |
name: str
|
86 |
model: str
|
@@ -89,9 +100,11 @@ class OllamaModel(BaseModel):
|
|
89 |
modified_at: str
|
90 |
details: OllamaModelDetails
|
91 |
|
|
|
92 |
class OllamaTagResponse(BaseModel):
|
93 |
models: List[OllamaModel]
|
94 |
|
|
|
95 |
def estimate_tokens(text: str) -> int:
|
96 |
"""Estimate the number of tokens in text
|
97 |
Chinese characters: approximately 1.5 tokens per character
|
@@ -106,6 +119,7 @@ def estimate_tokens(text: str) -> int:
|
|
106 |
|
107 |
return int(tokens)
|
108 |
|
|
|
109 |
def parse_query_mode(query: str) -> tuple[str, SearchMode]:
|
110 |
"""Parse query prefix to determine search mode
|
111 |
Returns tuple of (cleaned_query, search_mode)
|
@@ -127,6 +141,7 @@ def parse_query_mode(query: str) -> tuple[str, SearchMode]:
|
|
127 |
|
128 |
return query, SearchMode.hybrid
|
129 |
|
|
|
130 |
class OllamaAPI:
|
131 |
def __init__(self, rag: LightRAG):
|
132 |
self.rag = rag
|
@@ -333,10 +348,13 @@ class OllamaAPI:
|
|
333 |
"stream": request.stream,
|
334 |
"only_need_context": False,
|
335 |
"conversation_history": conversation_history,
|
336 |
-
"top_k": self.rag.args.top_k if hasattr(self.rag,
|
337 |
}
|
338 |
|
339 |
-
if
|
|
|
|
|
|
|
340 |
param_dict["history_turns"] = self.rag.args.history_turns
|
341 |
|
342 |
query_param = QueryParam(**param_dict)
|
@@ -521,7 +539,9 @@ class OllamaAPI:
|
|
521 |
**self.rag.llm_model_kwargs,
|
522 |
)
|
523 |
else:
|
524 |
-
response_text = await self.rag.aquery(
|
|
|
|
|
525 |
|
526 |
last_chunk_time = time.time_ns()
|
527 |
|
|
|
12 |
from ascii_colors import trace_exception
|
13 |
from lightrag import LightRAG, QueryParam
|
14 |
|
15 |
+
|
16 |
class OllamaServerInfos:
|
17 |
# Constants for emulated Ollama model information
|
18 |
LIGHTRAG_NAME = "lightrag"
|
|
|
22 |
LIGHTRAG_CREATED_AT = "2024-01-15T00:00:00Z"
|
23 |
LIGHTRAG_DIGEST = "sha256:lightrag"
|
24 |
|
25 |
+
|
26 |
ollama_server_infos = OllamaServerInfos()
|
27 |
|
28 |
+
|
29 |
# query mode according to query prefix (bypass is not LightRAG quer mode)
|
30 |
class SearchMode(str, Enum):
|
31 |
naive = "naive"
|
|
|
35 |
mix = "mix"
|
36 |
bypass = "bypass"
|
37 |
|
38 |
+
|
39 |
class OllamaMessage(BaseModel):
|
40 |
role: str
|
41 |
content: str
|
42 |
images: Optional[List[str]] = None
|
43 |
|
44 |
+
|
45 |
class OllamaChatRequest(BaseModel):
|
46 |
model: str
|
47 |
messages: List[OllamaMessage]
|
|
|
49 |
options: Optional[Dict[str, Any]] = None
|
50 |
system: Optional[str] = None
|
51 |
|
52 |
+
|
53 |
class OllamaChatResponse(BaseModel):
|
54 |
model: str
|
55 |
created_at: str
|
56 |
message: OllamaMessage
|
57 |
done: bool
|
58 |
|
59 |
+
|
60 |
class OllamaGenerateRequest(BaseModel):
|
61 |
model: str
|
62 |
prompt: str
|
|
|
64 |
stream: bool = False
|
65 |
options: Optional[Dict[str, Any]] = None
|
66 |
|
67 |
+
|
68 |
class OllamaGenerateResponse(BaseModel):
|
69 |
model: str
|
70 |
created_at: str
|
|
|
78 |
eval_count: Optional[int]
|
79 |
eval_duration: Optional[int]
|
80 |
|
81 |
+
|
82 |
class OllamaVersionResponse(BaseModel):
|
83 |
version: str
|
84 |
|
85 |
+
|
86 |
class OllamaModelDetails(BaseModel):
|
87 |
parent_model: str
|
88 |
format: str
|
|
|
91 |
parameter_size: str
|
92 |
quantization_level: str
|
93 |
|
94 |
+
|
95 |
class OllamaModel(BaseModel):
|
96 |
name: str
|
97 |
model: str
|
|
|
100 |
modified_at: str
|
101 |
details: OllamaModelDetails
|
102 |
|
103 |
+
|
104 |
class OllamaTagResponse(BaseModel):
|
105 |
models: List[OllamaModel]
|
106 |
|
107 |
+
|
108 |
def estimate_tokens(text: str) -> int:
|
109 |
"""Estimate the number of tokens in text
|
110 |
Chinese characters: approximately 1.5 tokens per character
|
|
|
119 |
|
120 |
return int(tokens)
|
121 |
|
122 |
+
|
123 |
def parse_query_mode(query: str) -> tuple[str, SearchMode]:
|
124 |
"""Parse query prefix to determine search mode
|
125 |
Returns tuple of (cleaned_query, search_mode)
|
|
|
141 |
|
142 |
return query, SearchMode.hybrid
|
143 |
|
144 |
+
|
145 |
class OllamaAPI:
|
146 |
def __init__(self, rag: LightRAG):
|
147 |
self.rag = rag
|
|
|
348 |
"stream": request.stream,
|
349 |
"only_need_context": False,
|
350 |
"conversation_history": conversation_history,
|
351 |
+
"top_k": self.rag.args.top_k if hasattr(self.rag, "args") else 50,
|
352 |
}
|
353 |
|
354 |
+
if (
|
355 |
+
hasattr(self.rag, "args")
|
356 |
+
and self.rag.args.history_turns is not None
|
357 |
+
):
|
358 |
param_dict["history_turns"] = self.rag.args.history_turns
|
359 |
|
360 |
query_param = QueryParam(**param_dict)
|
|
|
539 |
**self.rag.llm_model_kwargs,
|
540 |
)
|
541 |
else:
|
542 |
+
response_text = await self.rag.aquery(
|
543 |
+
cleaned_query, param=query_param
|
544 |
+
)
|
545 |
|
546 |
last_chunk_time = time.time_ns()
|
547 |
|