Revert model download before running workflow

app.py

@@ -4,7 +4,7 @@ from graph import BuildGraph
 from retriever import db_dir
 from langgraph.checkpoint.memory import MemorySaver
 from dotenv import load_dotenv
-from main import openai_model, model_id, DownloadChatModel
+from main import openai_model, model_id
 from util import get_sources, get_start_end_months
 from mods.tool_calling_llm import extract_think
 import requests
@@ -82,6 +82,7 @@ def run_workflow(input, history, compute_mode, thread_id, session_hash):
         if compute_mode == "local":
             gr.Info(
                 f"Please wait for the local model to load",
+                duration=15,
                 title=f"Model loading...",
             )
         # Get the chat model and build the graph
@@ -210,10 +211,6 @@ def to_workflow(request: gr.Request, *args):
     # Add session_hash to arguments
     new_args = args + (request.session_hash,)
     if compute_mode == "local":
-        # If graph hasn't been instantiated, download model before running workflow
-        graph = graph_instances[compute_mode].get(request.session_hash)
-        if graph is None:
-            DownloadChatModel()
         # Call the workflow function with the @spaces.GPU decorator
         for value in run_workflow_local(*new_args):
             yield value

graph.py

@@ -12,9 +12,6 @@ from retriever import BuildRetriever
 from prompts import query_prompt, generate_prompt, generic_tools_template
 from mods.tool_calling_llm import ToolCallingLLM
 
-# Local modules
-from retriever import BuildRetriever
-
 # For tracing (disabled)
 # os.environ["LANGSMITH_TRACING"] = "true"
 # os.environ["LANGSMITH_PROJECT"] = "R-help-chat"

main.py

@@ -5,7 +5,6 @@ from langchain_core.output_parsers import StrOutputParser
 from langgraph.checkpoint.memory import MemorySaver
 from langchain_core.messages import ToolMessage
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from huggingface_hub import snapshot_download
 from datetime import datetime
 from dotenv import load_dotenv
 import os
@@ -129,13 +128,6 @@ def ProcessDirectory(path, compute_mode):
             print(f"Chroma: no change for {file_path}")
 
 
-def DownloadChatModel():
-    """
-    Downloads a chat model to the local Hugging Face cache.
-    """
-    snapshot_download(model_id)
-
-
 def GetChatModel(compute_mode):
     """
     Get a chat model.

mods/tool_calling_llm.py

@@ -177,6 +177,10 @@ class ToolCallingLLM(BaseChatModel, ABC):
         # Extract <think>...</think> content and text after </think> for further processing 20250726 jmd
         think_text, post_think = extract_think(response_message.content)
 
+        ## For debugging
+        # print("post_think")
+        # print(post_think)
+
         # Parse output for JSON (support multiple objects separated by commas)
         try:
             parsed_json_results = json.loads(f"[{post_think}]")

retriever.py

@@ -174,9 +174,9 @@ def BuildRetrieverDense(compute_mode: str, top_k=6):
         # Get top k documents
         search_kwargs={"k": top_k},
     )
-    ## Release GPU memory
-    ## https://github.com/langchain-ai/langchain/discussions/10668
-    # torch.cuda.empty_cache()
+    # Fix for ValueError('Could not connect to tenant default_tenant. Are you sure it exists?')
+    # https://github.com/langchain-ai/langchain/issues/26884
+    chromadb.api.client.SharedSystemClient.clear_system_cache()
     return retriever