fix(environment): Split WRAPPER_MODEL_NAME from COLLECTION_NAME

Still, the former will default to the later when no defined, so
this change is BC. But enables us to use better model names, not
subject to the collection name restrictions (alphanum + low hyphen).

Author: stronk7

env_template

@@ -29,6 +29,8 @@ WRAPPER_API_BASE="0.0.0.0:8080"
 WRAPPER_CHAT_MAX_TURNS=10
 # To limit the number of OpenAI history tokens allowed in a chat (0 = no limit).
 WRAPPER_CHAT_MAX_TOKENS=1536
+# Public name that the wrapper will use to identify itself as a model. Will default to COLLECTION_NAME is not set.
+WRAPPER_MODEL_NAME="Your Model Name"
 
 # Validate requests auth against these bearer tokens.
 AUTH_TOKENS="11111111,22222222,33333333"

wiki_rag/search/util.py

@@ -20,7 +20,8 @@
 )
 from langchain_core.runnables import RunnableConfig
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
-from langgraph.graph.state import START, CompiledStateGraph, StateGraph
+from langgraph.constants import START
+from langgraph.graph.state import CompiledStateGraph, StateGraph
 from pymilvus import AnnSearchRequest, MilvusClient, WeightedRanker
 
 import wiki_rag.index as index
@@ -51,6 +52,7 @@ class ConfigSchema(TypedDict):
     stream: bool
     wrapper_chat_max_turns: int
     wrapper_chat_max_tokens: int
+    wrapper_model_name: str
 
 
 class RagState(TypedDict):

wiki_rag/server/main.py

@@ -67,6 +67,7 @@ def main():
     if not collection_name:
         logger.error("Collection name not found in environment. Exiting.")
         sys.exit(1)
+        # TODO: Validate that only numbers, letters and underscores are used.
 
     index.milvus_url = os.getenv("MILVUS_URL")
     if not index.milvus_url:
@@ -113,6 +114,10 @@ def main():
     # These are optional, default to 0 (unlimited).
     wrapper_chat_max_turns = int(os.getenv("WRAPPER_CHAT_MAX_TURNS", 0))
     wrapper_chat_max_tokens = int(os.getenv("WRAPPER_CHAT_MAX_TOKENS", 0))
+    wrapper_model_name = os.getenv("WRAPPER_MODEL_NAME") or os.getenv("COLLECTION_NAME")
+    if not wrapper_model_name:
+        logger.error("Public wrapper name not found in environment. Exiting.")  # This is unreachable.
+        sys.exit(1)
 
     logger.info("Building the graph")
     server.graph = build_graph()
@@ -136,6 +141,7 @@ def main():
         stream=False,
         wrapper_chat_max_turns=wrapper_chat_max_turns,
         wrapper_chat_max_tokens=wrapper_chat_max_tokens,
+        wrapper_model_name=wrapper_model_name,
     ).items()
 
     # Prepare the configuration.

wiki_rag/server/server.py

@@ -82,7 +82,7 @@ async def models_list() -> ModelsListResponse:
         object="list",
         data=[
             ModelResponse(
-                id=server.config["configurable"]["collection_name"],
+                id=server.config["configurable"]["wrapper_model_name"],
                 object="model",
                 created=int(time.time()),
                 owned_by=server.config["configurable"]["kb_name"],
@@ -113,7 +113,7 @@ async def chat_completions(request: ChatCompletionRequest) -> ChatCompletionResp
     if not request.model:
         raise HTTPException(status_code=400, detail="No model provided.")
 
-    if request.model != server.config["configurable"]["collection_name"]:
+    if request.model != server.config["configurable"]["wrapper_model_name"]:
         raise HTTPException(status_code=400, detail="Model not supported.")
 
     logger.debug(f"Request: {request}")

wiki_rag/server/util.py

@@ -42,7 +42,7 @@ class ChatCompletionRequest(BaseModel):
     max_completion_tokens: int | None = 768  # Max tokens to generate (not all models support this).
     temperature: float | None = 0.05  # Temperature for sampling (0.0, deterministic to 2.0, creative).
     top_p: float | None = 0.85  # Which probability (0.0 - 1.0) is used to consider the next token (0.85 default).
-    model: str = server.config["configurable"]["collection_name"]
+    model: str = server.config["configurable"]["wrapper_model_name"]
     messages: list[Message] = [Message(role="user", content="Hello!")]
     stream: bool | None = False
 
@@ -60,14 +60,14 @@ class ChatCompletionResponse(BaseModel):
     id: UUID4 = uuid.uuid4()
     object: str = "chat.completion"
     created: int = int(time.time())
-    model: str = server.config["configurable"]["collection_name"]
+    model: str = server.config["configurable"]["wrapper_model_name"]
     choices: list[ChoiceResponse] = [ChoiceResponse()]
 
 
 class ModelResponse(BaseModel):
     """Information about a LLM model."""
 
-    id: str = server.config["configurable"]["collection_name"]
+    id: str = server.config["configurable"]["wrapper_model_name"]
     object: str = "model"
     created: int = int(time.time())
     owned_by: str = server.config["configurable"]["kb_name"]