Added system prompt support in all modes

This commit is contained in:
MdNazishArmanShorthillsAI
2025-02-17 16:45:00 +05:30
parent da9560d706
commit c2ff17d343
3 changed files with 38 additions and 18 deletions

View File

@@ -984,7 +984,10 @@ class LightRAG:
await self._insert_done()
def query(
self, query: str, param: QueryParam = QueryParam(), prompt: str | None = None
self,
query: str,
param: QueryParam = QueryParam(),
system_prompt: str | None = None,
) -> str | Iterator[str]:
"""
Perform a sync query.
@@ -999,13 +1002,13 @@ class LightRAG:
"""
loop = always_get_an_event_loop()
return loop.run_until_complete(self.aquery(query, param, prompt)) # type: ignore
return loop.run_until_complete(self.aquery(query, param, system_prompt)) # type: ignore
async def aquery(
self,
query: str,
param: QueryParam = QueryParam(),
prompt: str | None = None,
system_prompt: str | None = None,
) -> str | AsyncIterator[str]:
"""
Perform a async query.
@@ -1037,7 +1040,7 @@ class LightRAG:
global_config=asdict(self),
embedding_func=self.embedding_func,
),
prompt=prompt,
system_prompt=system_prompt,
)
elif param.mode == "naive":
response = await naive_query(
@@ -1056,6 +1059,7 @@ class LightRAG:
global_config=asdict(self),
embedding_func=self.embedding_func,
),
system_prompt=system_prompt,
)
elif param.mode == "mix":
response = await mix_kg_vector_query(
@@ -1077,6 +1081,7 @@ class LightRAG:
global_config=asdict(self),
embedding_func=self.embedding_func,
),
system_prompt=system_prompt,
)
else:
raise ValueError(f"Unknown mode {param.mode}")