wren
104 lines
3.7 KiB
Python
104 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
from fastapi_admin.config._settings import embedding, llm
|
|
|
|
|
|
def _get_str(name: str, default: str = "") -> str:
|
|
import os
|
|
return os.getenv(name, default)
|
|
|
|
|
|
def _get_bool(name: str, default: bool = False) -> bool:
|
|
import os
|
|
return os.getenv(name, str(default).lower()).lower() == "true"
|
|
|
|
|
|
def _get_int(name: str, default: int) -> int:
|
|
import os
|
|
try:
|
|
return int(os.getenv(name, str(default)))
|
|
except ValueError:
|
|
return default
|
|
|
|
|
|
def _get_float(name: str, default: float) -> float:
|
|
import os
|
|
try:
|
|
return float(os.getenv(name, str(default)))
|
|
except ValueError:
|
|
return default
|
|
|
|
|
|
RAG_CONFIG = {
|
|
"USE_SELF_HOSTED": True,
|
|
"CHROMA_PERSIST_DIR": _get_str("RAG_CHROMA_PERSIST_DIR", ".chromadb_rag"),
|
|
"CHROMA_HOST": _get_str("RAG_CHROMA_HOST", ""),
|
|
"CHROMA_PORT": _get_int("RAG_CHROMA_PORT", 8010),
|
|
"CHROMA_TOKEN": _get_str("RAG_CHROMA_TOKEN", ""),
|
|
"CHROMA_AUTH_HEADER": _get_str("RAG_CHROMA_AUTH_HEADER", "X-Chroma-Token"),
|
|
"EMBED_URL": _get_str(
|
|
"RAG_EMBED_URL",
|
|
_get_str("GRAPH_RAG_EMBED_URL", _get_str("EMBEDDING_BASE_URL", embedding.EMBEDDING_BASE_URL)),
|
|
),
|
|
"EMBED_KEY": _get_str(
|
|
"RAG_EMBED_KEY",
|
|
_get_str("GRAPH_RAG_EMBED_KEY", _get_str("EMBEDDING_API_KEY", embedding.EMBEDDING_API_KEY)),
|
|
),
|
|
"EMBED_MODEL": _get_str(
|
|
"RAG_EMBED_MODEL",
|
|
_get_str("GRAPH_RAG_EMBED_MODEL", _get_str("EMBEDDING_MODEL", embedding.EMBEDDING_MODEL)),
|
|
),
|
|
"EMBED_DIM": _get_int("RAG_EMBED_DIM", _get_int("EMBEDDING_DIM", embedding.EMBEDDING_DIM)),
|
|
"EMBED_BATCH_SIZE": _get_int(
|
|
"RAG_EMBED_BATCH_SIZE",
|
|
_get_int("EMBEDDING_BATCH_SIZE", embedding.EMBEDDING_BATCH_SIZE),
|
|
),
|
|
"RERANKER_URL": _get_str("RAG_RERANKER_URL", _get_str("GRAPH_RAG_RERANKER_URL", "")),
|
|
"RERANKER_KEY": _get_str("RAG_RERANKER_KEY", _get_str("GRAPH_RAG_RERANKER_KEY", "")),
|
|
"RERANKER_MODEL": _get_str("RAG_RERANKER_MODEL", _get_str("GRAPH_RAG_RERANKER_MODEL", "")),
|
|
"LLM_BASE_URL": _get_str("LLM_BASE_URL", llm.LLM_BASE_URL),
|
|
"LLM_MODEL": _get_str("LLM_MODEL", llm.LLM_MODEL),
|
|
"LLM_API_KEY": _get_str("LLM_API_KEY", llm.LLM_API_KEY),
|
|
"VECTOR_TOP_K": _get_int("RAG_VECTOR_TOP_K", 15),
|
|
"RERANK_TOP_K": _get_int("RAG_RERANK_TOP_K", 5),
|
|
"BM25_TOP_K": _get_int("RAG_BM25_TOP_K", 15),
|
|
"RRF_K": _get_int("RAG_RRF_K", 60),
|
|
"LLM_TEMPERATURE": _get_float("RAG_LLM_TEMPERATURE", 0.3),
|
|
"LLM_MAX_TOKENS": _get_int("RAG_LLM_MAX_TOKENS", 2048),
|
|
"LLM_TIMEOUT": _get_int("RAG_LLM_TIMEOUT", 120),
|
|
"QUERY_REWRITING": _get_bool("RAG_QUERY_REWRITING", False),
|
|
"HYBRID_SEARCH": _get_bool("RAG_HYBRID_SEARCH", True),
|
|
"RERANKING": _get_bool("RAG_RERANKING", True),
|
|
}
|
|
|
|
|
|
def build_openai_chat_completions_url(base_url: str) -> str:
|
|
normalized = (base_url or "").strip().rstrip("/")
|
|
if not normalized:
|
|
return "/chat/completions"
|
|
if normalized.endswith("/chat/completions"):
|
|
return normalized
|
|
return f"{normalized}/chat/completions"
|
|
|
|
|
|
def build_openai_embeddings_url(base_url: str) -> str:
|
|
normalized = (base_url or "").strip().rstrip("/")
|
|
if not normalized:
|
|
return "/embeddings"
|
|
if normalized.endswith("/chat/completions"):
|
|
normalized = normalized[:-len("/chat/completions")]
|
|
if normalized.endswith("/embeddings"):
|
|
return normalized
|
|
return f"{normalized}/embeddings"
|
|
|
|
|
|
def normalize_openai_base_url(base_url: str) -> str:
|
|
normalized = (base_url or "").strip().rstrip("/")
|
|
if not normalized:
|
|
return ""
|
|
if normalized.endswith("/chat/completions"):
|
|
return normalized[:-len("/chat/completions")]
|
|
if normalized.endswith("/embeddings"):
|
|
return normalized[:-len("/embeddings")]
|
|
return normalized
|