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, 'args') else 50,
337
  }
338
 
339
- if hasattr(self.rag, 'args') and self.rag.args.history_turns is not None:
 
 
 
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(cleaned_query, param=query_param)
 
 
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