Revert "Attempt fix for ChromaDB ValueError"

This reverts commit 9477d3a838873d357a7dba604e1675933fb649c8.

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,11 +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
     if compute_mode == "remote":
@@ -222,7 +218,7 @@ def to_workflow(request: gr.Request, *args):
             yield value
 
 
-@spaces.GPU(duration=90)
+@spaces.GPU(duration=60)
 def run_workflow_local(*args):
     for value in run_workflow(*args):
         yield value
@@ -268,11 +264,13 @@ with gr.Blocks(
             "local",
             "remote",
         ],
-        # Default to remote because it provides a better first impression for most people
-        # value=("local" if torch.cuda.is_available() else "remote"),
-        value="remote",
+        value=("local" if torch.cuda.is_available() else "remote"),
         label="Compute Mode",
-        info="NOTE: remote mode **does not** use ZeroGPU",
+        info=(
+            "NOTE: remote mode **does not** use ZeroGPU"
+            if torch.cuda.is_available()
+            else "NOTE: local mode requires GPU"
+        ),
         render=False,
     )
 

graph.py

@@ -12,6 +12,9 @@ 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,10 +177,6 @@ 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}]")

requirements.txt

@@ -29,5 +29,3 @@ spaces==0.37.1
 #boto3==1.39.14
 # Others
 python-dotenv
-# For snapshot_download
-huggingface-hub==0.33.2

retriever.py

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