fix(bedrock): Add Gov Cloud regions (#6105)

backend/Dockerfile

@@ -90,6 +90,10 @@ nltk.download('stopwords', quiet=True); \
 nltk.download('punkt_tab', quiet=True);"
 # nltk.download('wordnet', quiet=True); introduce this back if lemmatization is needed
 
+# Pre-downloading tiktoken for setups with limited egress
+RUN python -c "import tiktoken; \
+tiktoken.get_encoding('cl100k_base')"
+
 # Set up application files
 WORKDIR /app
 

backend/onyx/llm/llm_provider_options.py

@@ -103,6 +103,8 @@ def _fallback_bedrock_regions() -> list[str]:
     return [
         "us-east-1",
         "us-east-2",
+        "us-gov-east-1",
+        "us-gov-west-1",
         "us-west-2",
         "ap-northeast-1",
         "ap-south-1",
@@ -120,8 +122,22 @@ def _build_bedrock_region_options() -> list[CustomConfigOption]:
         import boto3
 
         session = boto3.session.Session()
-        regions = set(session.get_available_regions("bedrock"))
-        regions.update(session.get_available_regions("bedrock-runtime"))
+        regions: set[str] = set()
+        # Include both commercial and GovCloud partitions so GovCloud users can select their region.
+        for partition_name in ("aws", "aws-us-gov"):
+            try:
+                regions.update(
+                    session.get_available_regions(
+                        "bedrock", partition_name=partition_name
+                    )
+                )
+                regions.update(
+                    session.get_available_regions(
+                        "bedrock-runtime", partition_name=partition_name
+                    )
+                )
+            except Exception:
+                continue
         if not regions:
             raise ValueError("No Bedrock regions returned from boto3")
         sorted_regions = sorted(regions)

web/src/app/chat/message/messageComponents/AIMessage.tsx

@@ -169,7 +169,6 @@ export default function AIMessage({
   // Track indices for graceful SECTION_END injection
   const seenIndicesRef = useRef<Set<number>>(new Set());
   const indicesWithSectionEndRef = useRef<Set<number>>(new Set());
-  const toolIndicesRef = useRef<Set<number>>(new Set());
 
   // Reset incremental state when switching messages or when stream resets
   const resetState = () => {
@@ -184,7 +183,6 @@ export default function AIMessage({
     stopPacketSeenRef.current = isStreamingComplete(rawPackets);
     seenIndicesRef.current = new Set();
     indicesWithSectionEndRef.current = new Set();
-    toolIndicesRef.current = new Set();
   };
   useEffect(() => {
     resetState();
@@ -240,10 +238,7 @@ export default function AIMessage({
       // If we see a new index, inject SECTION_END for previous tool indices
       if (isNewIndex && seenIndicesRef.current.size > 0) {
         Array.from(seenIndicesRef.current).forEach((prevInd) => {
-          if (
-            toolIndicesRef.current.has(prevInd) &&
-            !indicesWithSectionEndRef.current.has(prevInd)
-          ) {
+          if (!indicesWithSectionEndRef.current.has(prevInd)) {
             injectSectionEnd(prevInd);
           }
         });
@@ -252,11 +247,6 @@ export default function AIMessage({
       // Track this index
       seenIndicesRef.current.add(currentInd);
 
-      // Track if this is a tool packet
-      if (isToolPacket(packet, false)) {
-        toolIndicesRef.current.add(currentInd);
-      }
-
       // Track SECTION_END packets
       if (packet.obj.type === PacketType.SECTION_END) {
         indicesWithSectionEndRef.current.add(currentInd);
@@ -310,10 +300,10 @@ export default function AIMessage({
 
       if (packet.obj.type === PacketType.STOP && !stopPacketSeenRef.current) {
         setStopPacketSeen(true);
-        // Inject SECTION_END for all tool indices that don't have one
-        Array.from(toolIndicesRef.current).forEach((toolInd) => {
-          if (!indicesWithSectionEndRef.current.has(toolInd)) {
-            injectSectionEnd(toolInd);
+        // Inject SECTION_END for all indices that don't have one
+        Array.from(seenIndicesRef.current).forEach((ind) => {
+          if (!indicesWithSectionEndRef.current.has(ind)) {
+            injectSectionEnd(ind);
           }
         });
       }
@@ -355,19 +345,6 @@ export default function AIMessage({
   const updateCurrentSelectedNodeForDocDisplay = useChatSessionStore(
     (state) => state.updateCurrentSelectedNodeForDocDisplay
   );
-  // Calculate unique source count
-  const _uniqueSourceCount = useMemo(() => {
-    const uniqueDocIds = new Set<string>();
-    for (const citation of citations) {
-      if (citation.document_id) {
-        uniqueDocIds.add(citation.document_id);
-      }
-    }
-    documentMap.forEach((_, docId) => {
-      uniqueDocIds.add(docId);
-    });
-    return uniqueDocIds.size;
-  }, [citations.length, documentMap.size]);
 
   // Message switching logic
   const {