diff --git a/packages/components/src/components/code-block/code-block.tsx b/packages/components/src/components/code-block/code-block.tsx
index 61ac677..af101a6 100644
--- a/packages/components/src/components/code-block/code-block.tsx
+++ b/packages/components/src/components/code-block/code-block.tsx
@@ -79,6 +79,8 @@ type CodeBlockProps = {
   feedbackButton?: ReactNode;
   children?: ReactNode;
   copyButtonProps?: CopyToClipboardButtonProps;
+  isParentCodeGroup?: boolean;
+  shouldHighlight?: boolean;
 };
 
 const CodeBlock = function CodeBlock(params: CodeBlockProps) {
@@ -87,12 +89,10 @@ const CodeBlock = function CodeBlock(params: CodeBlockProps) {
     children,
     className,
     icon,
-    isSmallText,
     hideAskAiButton,
     feedbackModalOpen,
     anchorRef,
     codeBlockTheme = "system",
-    codeBlockThemeObject,
     askAiButton,
     feedbackButton,
     copyButtonProps,
@@ -130,7 +130,7 @@ const CodeBlock = function CodeBlock(params: CodeBlockProps) {
             textToCopy={codeString}
             {...copyButtonProps}
           />
-          {askAiButton && askAiButton}
+          {askAiButton && !hideAskAiButton ? askAiButton : null}
         </CodeHeader>
       ) : (
         <div
@@ -143,16 +143,10 @@ const CodeBlock = function CodeBlock(params: CodeBlockProps) {
             textToCopy={codeString}
             {...copyButtonProps}
           />
-          {askAiButton && askAiButton}
+          {askAiButton && !hideAskAiButton ? askAiButton : null}
         </div>
       )}
-      <BaseCodeBlock
-        codeBlockTheme={codeBlockTheme}
-        codeBlockThemeObject={codeBlockThemeObject}
-        hideAskAiButton={hideAskAiButton}
-        isSmallText={isSmallText}
-        {...params}
-      />
+      <BaseCodeBlock {...params} />
     </div>
   );
 };