Add RAG configuration options and enhance parameter configurability

- Add top-k and cosine-threshold parms for api server
- Update .env and cli parms handling with new parameters
- Improve splash screen display
- Update bash and storage classes to read new parameters from .env file.
This commit is contained in:
yangdx
2025-01-29 21:34:34 +08:00
parent d0052456d4
commit 7aedc08caf
4 changed files with 42 additions and 12 deletions

View File

@@ -207,8 +207,12 @@ def display_splash_screen(args: argparse.Namespace) -> None:
ASCIIColors.yellow(f"{args.chunk_size}")
ASCIIColors.white(" ├─ Chunk Overlap Size: ", end="")
ASCIIColors.yellow(f"{args.chunk_overlap_size}")
ASCIIColors.white(" ─ History Turns: ", end="")
ASCIIColors.white(" ─ History Turns: ", end="")
ASCIIColors.yellow(f"{args.history_turns}")
ASCIIColors.white(" ├─ Cosine Threshold: ", end="")
ASCIIColors.yellow(f"{args.cosine_threshold}")
ASCIIColors.white(" └─ Top-K: ", end="")
ASCIIColors.yellow(f"{args.top_k}")
# System Configuration
ASCIIColors.magenta("\n🛠️ System Configuration:")
@@ -484,6 +488,20 @@ def parse_args() -> argparse.Namespace:
help="Number of conversation history turns to include (default: from env or 3)",
)
# Search parameters
parser.add_argument(
"--top-k",
type=int,
default=get_env_value("TOP_K", 50, int),
help="Number of most similar results to return (default: from env or 50)",
)
parser.add_argument(
"--cosine-threshold",
type=float,
default=get_env_value("COSINE_THRESHOLD", 0.4, float),
help="Cosine similarity threshold (default: from env or 0.4)",
)
args = parser.parse_args()
return args
@@ -846,6 +864,9 @@ def create_app(args):
graph_storage=GRAPH_STORAGE,
vector_storage=VECTOR_STORAGE,
doc_status_storage=DOC_STATUS_STORAGE,
vector_db_storage_cls_kwargs={
"cosine_better_than_threshold": args.cosine_threshold
},
)
else:
rag = LightRAG(
@@ -863,6 +884,9 @@ def create_app(args):
graph_storage=GRAPH_STORAGE,
vector_storage=VECTOR_STORAGE,
doc_status_storage=DOC_STATUS_STORAGE,
vector_db_storage_cls_kwargs={
"cosine_better_than_threshold": args.cosine_threshold
},
)
async def index_file(file_path: Union[str, Path]) -> None:
@@ -1052,6 +1076,7 @@ def create_app(args):
mode=request.mode,
stream=request.stream,
only_need_context=request.only_need_context,
top_k=args.top_k,
),
)
@@ -1093,6 +1118,7 @@ def create_app(args):
mode=request.mode,
stream=True,
only_need_context=request.only_need_context,
top_k=args.top_k,
),
)
@@ -1632,6 +1658,7 @@ def create_app(args):
"stream": request.stream,
"only_need_context": False,
"conversation_history": conversation_history,
"top_k": args.top_k,
}
if args.history_turns is not None: