diff --git a/ai-hub/app/core/services/utils/code_change.py b/ai-hub/app/core/services/utils/code_change.py
index 17f3863..c2eaad0 100644
--- a/ai-hub/app/core/services/utils/code_change.py
+++ b/ai-hub/app/core/services/utils/code_change.py
@@ -17,7 +17,7 @@
     A helper class to process and manage a sequence of code change instructions.
     """
 
-    def __init__(self, db: Session, provider_name: str, input_data: str):
+    def __init__(self, db: Session, provider_name: str, input_data: str, request_id: uuid.UUID):
         """
         Initializes the CodeChangeHelper, parsing the input and setting up dependencies.
         
@@ -37,7 +37,7 @@
         self.last_step_index: int = -1
         try:
             self._parse_input_data()
-            self._preload_original_files()
+            self._preload_original_files(request_id=request_id)
         except (json.JSONDecodeError, ValueError) as e:
             logger.error(f"Initialization failed due to invalid input: {e}")
             raise
@@ -59,7 +59,7 @@
         
         self.parsed_data = parsed_data
 
-    def _preload_original_files(self) -> None:
+    def _preload_original_files(self, request_id: uuid.UUID) -> None:
         """
         Fetches and caches the content of all required original files.
         """
@@ -70,11 +70,11 @@
                 unique_file_paths.add(path)
         
         for file_path in unique_file_paths:
-            content = self._fetch_file_content(file_path)
+            content = self._fetch_file_content(file_path, request_id=request_id)
             if content is not None:
                 self.original_files[file_path] = content
 
-    def _fetch_file_content(self, file_path: str) -> Optional[str]:
+    def _fetch_file_content(self, file_path: str, request_id: uuid.UUID) -> Optional[str]:
         """
         Fetches the content of a file from the database.
         
@@ -82,7 +82,7 @@
         """
         try:
             # Assuming a single request_id for simplicity; adjust if needed
-            retrieved_file = self.db.query(file_retriever_models.RetrievedFile).filter_by(file_path=file_path).first()
+            retrieved_file = self.db.query(file_retriever_models.RetrievedFile).filter_by(file_path=file_path, request_id=request_id).first()
             if retrieved_file:
                 return retrieved_file.content
             else:
@@ -155,11 +155,32 @@
         }
         await websocket.send_text(json.dumps(client_log))
     
-    async def _post_process(self) ->Dict[str, Dict[str, str]]:
-        result= {}
+    async def _post_process(self) -> Dict[str, Dict[str, str]]:
+        result = {}
+        
+        # Regex to find and extract content from a Markdown code block
+        # The `re.DOTALL` flag allows `.` to match newlines
+        code_block_pattern = re.compile(r'```[a-zA-Z]*\n(.*)```', re.DOTALL)
+        
         for file_path, detail in self.updated_files.items():
             original_content = self.original_files.get(file_path, "")
-            result[file_path] = {"old": original_content, "new": detail.get("content", ""), "reasoning": detail.get("reasoning", "")}
+            updated_content = detail.get("content", "")
+            
+            # Check if the content is wrapped in a code block
+            match = code_block_pattern.search(updated_content)
+            if match:
+                # If a match is found, use the captured group as the new content,
+                # and strip any leading/trailing whitespace
+                cleaned_content = match.group(1).strip()
+            else:
+                # If no code block is found, use the content as-is
+                cleaned_content = updated_content
+                
+            result[file_path] = {
+                "old": original_content,
+                "new": cleaned_content,
+                "reasoning": detail.get("reasoning", "")
+            }
         return result
 
     async def process(self, websocket: WebSocket) -> Dict[str, Dict[str, str]]:
diff --git a/ai-hub/app/core/services/workspace.py b/ai-hub/app/core/services/workspace.py
index 4dadb28..e309094 100644
--- a/ai-hub/app/core/services/workspace.py
+++ b/ai-hub/app/core/services/workspace.py
@@ -670,7 +670,7 @@
             
             try:
                 # The input_data is a JSON string of code change instructions
-                cch = CodeChangeHelper(db=self.db, provider_name="gemini", input_data=raw_answer_text)
+                cch = CodeChangeHelper(db=self.db, provider_name="gemini", input_data=raw_answer_text,request_id= uuid.UUID(request_id))
                 
                 # Use the CodeChangeHelper to process all code changes
                 final_changes = await cch.process(websocket=websocket)
diff --git a/ui/client-app/src/components/ChatWindow.js b/ui/client-app/src/components/ChatWindow.js
index 36cf1d2..4e37863 100644
--- a/ui/client-app/src/components/ChatWindow.js
+++ b/ui/client-app/src/components/ChatWindow.js
@@ -47,7 +47,7 @@
       {message.code_changes && (
         <FileListComponent code_changes={message.code_changes} onFileClick={handleFileClick} />
       )}
-      {selectedFile && <DiffViewer oldContent={selectedFile.old} newContent={selectedFile.new} filePath="" onClose={handleCloseDiff} />}
+      {selectedFile && <DiffViewer oldContent={selectedFile.old} newContent={selectedFile.new} filePath={selectedFile.filepath} onClose={handleCloseDiff} />}
     </div>
   );
 };
diff --git a/ui/client-app/src/components/DiffViewer.css b/ui/client-app/src/components/DiffViewer.css
index 5c24a89..eac3042 100644
--- a/ui/client-app/src/components/DiffViewer.css
+++ b/ui/client-app/src/components/DiffViewer.css
@@ -5,4 +5,8 @@
   position: static !important;
   background: inherit !important;
   z-index: auto !important;
+}
+
+ol li::before {
+      content: "";
 }
\ No newline at end of file
diff --git a/ui/client-app/src/components/DiffViewer.js b/ui/client-app/src/components/DiffViewer.js
index ff6d979..ba5d04c 100644
--- a/ui/client-app/src/components/DiffViewer.js
+++ b/ui/client-app/src/components/DiffViewer.js
@@ -80,13 +80,14 @@
           <h2 className="text-xl font-semibold text-gray-900 dark:text-gray-100">Changes in {filePath}</h2>
           <div className="flex items-center space-x-2">
             <div className="relative">
-              <button
+             <button
                 onClick={() => setIsDropdownOpen(!isDropdownOpen)}
-                className="text-gray-500 hover:text-gray-700 dark:text-gray-400 dark:hover:text-gray-200 transition-colors duration-200"
+                className="flex items-center space-x-1 px-4 py-2 rounded-md text-gray-700 bg-gray-100 hover:bg-gray-200 dark:text-gray-200 dark:bg-gray-700 dark:hover:bg-gray-600 transition-colors duration-200"
               >
-                <svg xmlns="http://www.w3.org/2000/svg" className="h-6 w-6" viewBox="0 0 20 20" fill="currentColor">
-                  <path fillRule="evenodd" d="M5.293 7.293a1 1 0 011.414 0L10 10.586l3.293-3.293a1 1 0 111.414 1.414l-4 4a1 1 0 01-1.414 0l-4-4a1 1 0 010-1.414z" clipRule="evenodd" />
+               <svg xmlns="http://www.w3.org/2000/svg" className="h-6 w-6" viewBox="0 0 20 20" fill="currentColor">
+                <path fillRule="evenodd" d="M5.293 7.293a1 1 0 011.414 0L10 10.586l3.293-3.293a1 1 0 111.414 1.414l-4 4a1 1 0 01-1.414 0l-4-4a1 1 0 010-1.414z" clipRule="evenodd" />
                 </svg>
+                <span>Download</span>
               </button>
               {isDropdownOpen && (
                 <div className="absolute right-0 mt-2 w-48 rounded-md shadow-lg bg-white dark:bg-gray-700 ring-1 ring-black ring-opacity-5 focus:outline-none z-10">