gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Added the option to only return context without doing a query to the AI (useful for automated systems like lollms)

Browse Source
This commit is contained in:
Saifeddine ALOUI
2024-12-26 23:32:02 +01:00
parent 5fcfb05e62
commit 4d0184189f
4 changed files with 12 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
lightrag/api/azure_openai_lightrag_server.py
View File

@@ -136,6 +136,7 @@ class SearchMode(str, Enum):
class QueryRequest(BaseModel): class QueryRequest(BaseModel):
query: str query: str
mode: SearchMode = SearchMode.hybrid mode: SearchMode = SearchMode.hybrid
only_need_context: bool = False
# stream: bool = False # stream: bool = False
@@ -308,7 +309,7 @@ def create_app(args):
try: try:
response = await rag.aquery( response = await rag.aquery(
request.query, request.query,
param=QueryParam(mode=request.mode, stream=False), param=QueryParam(mode=request.mode, stream=False, only_need_context=request.only_need_context),
) )
return QueryResponse(response=response) return QueryResponse(response=response)
except Exception as e: except Exception as e:
@@ -319,7 +320,7 @@ def create_app(args):
try: try:
response = await rag.aquery( response = await rag.aquery(
request.query, request.query,
param=QueryParam(mode=request.mode, stream=True), param=QueryParam(mode=request.mode, stream=True, only_need_context=request.only_need_context),
) )
if inspect.isasyncgen(response): if inspect.isasyncgen(response):

5
lightrag/api/lollms_lightrag_server.py
View File

@@ -130,6 +130,7 @@ class QueryRequest(BaseModel):
query: str query: str
mode: SearchMode = SearchMode.hybrid mode: SearchMode = SearchMode.hybrid
stream: bool = False stream: bool = False
only_need_context: bool = False
class QueryResponse(BaseModel): class QueryResponse(BaseModel):
@@ -266,7 +267,7 @@ def create_app(args):
try: try:
response = await rag.aquery( response = await rag.aquery(
request.query, request.query,
param=QueryParam(mode=request.mode, stream=request.stream), param=QueryParam(mode=request.mode, stream=request.stream, only_need_context=request.only_need_context),
) )
if request.stream: if request.stream:
@@ -283,7 +284,7 @@ def create_app(args):
async def query_text_stream(request: QueryRequest): async def query_text_stream(request: QueryRequest):
try: try:
response = rag.query( response = rag.query(
request.query, param=QueryParam(mode=request.mode, stream=True) request.query, param=QueryParam(mode=request.mode, stream=True, only_need_context=request.only_need_context)
) )
async def stream_generator(): async def stream_generator():

5
lightrag/api/ollama_lightrag_server.py
View File

@@ -130,6 +130,7 @@ class QueryRequest(BaseModel):
query: str query: str
mode: SearchMode = SearchMode.hybrid mode: SearchMode = SearchMode.hybrid
stream: bool = False stream: bool = False
only_need_context: bool = False
class QueryResponse(BaseModel): class QueryResponse(BaseModel):
@@ -266,7 +267,7 @@ def create_app(args):
try: try:
response = await rag.aquery( response = await rag.aquery(
request.query, request.query,
param=QueryParam(mode=request.mode, stream=request.stream), param=QueryParam(mode=request.mode, stream=request.stream, only_need_context=request.only_need_context),
) )
if request.stream: if request.stream:
@@ -283,7 +284,7 @@ def create_app(args):
async def query_text_stream(request: QueryRequest): async def query_text_stream(request: QueryRequest):
try: try:
response = rag.query( response = rag.query(
request.query, param=QueryParam(mode=request.mode, stream=True) request.query, param=QueryParam(mode=request.mode, stream=True, only_need_context=request.only_need_context)
) )
async def stream_generator(): async def stream_generator():

5
lightrag/api/openai_lightrag_server.py
View File

@@ -119,6 +119,7 @@ class QueryRequest(BaseModel):
query: str query: str
mode: SearchMode = SearchMode.hybrid mode: SearchMode = SearchMode.hybrid
stream: bool = False stream: bool = False
only_need_context: bool = False
class QueryResponse(BaseModel): class QueryResponse(BaseModel):
@@ -270,7 +271,7 @@ def create_app(args):
try: try:
response = await rag.aquery( response = await rag.aquery(
request.query, request.query,
param=QueryParam(mode=request.mode, stream=request.stream), param=QueryParam(mode=request.mode, stream=request.stream, only_need_context=request.only_need_context),
) )
if request.stream: if request.stream:
@@ -287,7 +288,7 @@ def create_app(args):
async def query_text_stream(request: QueryRequest): async def query_text_stream(request: QueryRequest):
try: try:
response = rag.query( response = rag.query(
request.query, param=QueryParam(mode=request.mode, stream=True) request.query, param=QueryParam(mode=request.mode, stream=True, only_need_context=request.only_need_context)
) )
async def stream_generator(): async def stream_generator():